Upload 244 files

.gitattributes

@@ -33,3 +33,19 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/13[[:space:]]00_00_00-00_00_30.gif filter=lfs diff=lfs merge=lfs -text
+assets/5[[:space:]]00_00_00-00_00_30.gif filter=lfs diff=lfs merge=lfs -text
+assets/6[[:space:]]00_00_00-00_00_30.gif filter=lfs diff=lfs merge=lfs -text
+assets/7[[:space:]]00_00_00-00_00_30.gif filter=lfs diff=lfs merge=lfs -text
+assets/dpvj8-y3ubn.gif filter=lfs diff=lfs merge=lfs -text
+assets/i1ude-11d4e.gif filter=lfs diff=lfs merge=lfs -text
+assets/kntw7-iuluy.gif filter=lfs diff=lfs merge=lfs -text
+assets/nr2a2-oe6qj.gif filter=lfs diff=lfs merge=lfs -text
+assets/ns4et-xj8ax.gif filter=lfs diff=lfs merge=lfs -text
+assets/open-sora-plan.png filter=lfs diff=lfs merge=lfs -text
+assets/ozg76-g1aqh.gif filter=lfs diff=lfs merge=lfs -text
+assets/pvvm5-5hm65.gif filter=lfs diff=lfs merge=lfs -text
+assets/rrdqk-puoud.gif filter=lfs diff=lfs merge=lfs -text
+assets/y70q9-y5tip.gif filter=lfs diff=lfs merge=lfs -text
+opensora/models/captioner/caption_refiner/dataset/test_videos/video1.gif filter=lfs diff=lfs merge=lfs -text
+opensora/models/captioner/caption_refiner/dataset/test_videos/video2.gif filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 PKU-YUAN's Group (袁粒课题组-北大信工)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

docker/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 SimonLee
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

docker/README.md

@@ -0,0 +1,87 @@
+# Docker4ML
+
+Useful docker scripts for ML developement.
+[https://github.com/SimonLeeGit/Docker4ML](https://github.com/SimonLeeGit/Docker4ML)
+
+## Build Docker Image
+
+```bash
+bash docker_build.sh
+```
+
+![build_docker](build_docker.png)
+
+## Run Docker Container as Development Envirnoment
+
+```bash
+bash docker_run.sh
+```
+
+![run_docker](run_docker.png)
+
+## Custom Docker Config
+
+### Config [setup_env.sh](./setup_env.sh)
+
+You can modify this file to custom your settings.
+
+```bash
+TAG=ml:dev
+BASE_TAG=nvcr.io/nvidia/pytorch:23.12-py3
+```
+
+#### TAG
+
+Your built docker image tag, you can set it as what you what.
+
+#### BASE_TAG
+
+The base docker image tag for your built docker image, here we use nvidia pytorch images.
+You can check it from [https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags)
+
+Also, you can use other docker image as base, such as: [ubuntu](https://hub.docker.com/_/ubuntu/tags)
+
+### USER_NAME
+
+Your user name used in docker container.
+
+### USER_PASSWD
+
+Your user password used in docker container.
+
+### Config [requriements.txt](./requirements.txt)
+
+You can add your default installed python libraries here.
+
+```txt
+transformers==4.27.1
+```
+
+By default, it has some libs installed, you can check it from [https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html](https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-01.html)
+
+### Config [packages.txt](./packages.txt)
+
+You can add your default apt-get installed packages here.
+
+```txt
+wget
+curl
+git
+```
+
+### Config [ports.txt](./ports.txt)
+
+You can add some ports enabled for docker container here.
+
+```txt
+-p 6006:6006
+-p 8080:8080
+```
+
+### Config [postinstallscript.sh](./postinstallscript.sh)
+
+You can add your custom script to run when build docker image.
+
+## Q&A
+
+If you have any use problems, please contact to <simonlee235@gmail.com>.

docker/docker_build.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+WORK_DIR=$(dirname "$(readlink -f "$0")")
+cd $WORK_DIR
+
+source setup_env.sh
+
+docker build -t $TAG --build-arg BASE_TAG=$BASE_TAG --build-arg USER_NAME=$USER_NAME --build-arg USER_PASSWD=$USER_PASSWD . -f dockerfile.base

docker/docker_run.sh

@@ -0,0 +1,45 @@
+#!/usr/bin/env bash
+
+WORK_DIR=$(dirname "$(readlink -f "$0")")
+source $WORK_DIR/setup_env.sh
+
+RUNNING_IDS="$(docker ps --filter ancestor=$TAG --format "{{.ID}}")"
+
+if [ -n "$RUNNING_IDS" ]; then
+    # Initialize an array to hold the container IDs
+    declare -a container_ids=($RUNNING_IDS)
+
+    # Get the first container ID using array indexing
+    ID=${container_ids[0]}
+
+    # Print the first container ID
+    echo ' '
+    echo "The running container ID is: $ID, enter it!"
+else
+    echo ' '
+    echo "Not found running containers, run it!"
+
+    # Run a new docker container instance
+    ID=$(docker run \
+        --rm \
+        --gpus all \
+        -itd \
+        --ipc=host \
+        --ulimit memlock=-1 \
+        --ulimit stack=67108864 \
+        -e DISPLAY=$DISPLAY \
+        -v /tmp/.X11-unix/:/tmp/.X11-unix/ \
+        -v $PWD:/home/$USER_NAME/workspace \
+        -w /home/$USER_NAME/workspace \
+        $(cat $WORK_DIR/ports.txt) \
+        $TAG)
+fi
+
+docker logs $ID
+
+echo ' '
+echo ' '
+echo '========================================='
+echo ' '
+
+docker exec -it $ID bash

docker/dockerfile.base

@@ -0,0 +1,24 @@
+ARG BASE_TAG
+FROM ${BASE_TAG}
+ARG USER_NAME=myuser
+ARG USER_PASSWD=111111
+ARG DEBIAN_FRONTEND=noninteractive
+
+# Pre-install packages, pip install requirements and run post install script.
+COPY packages.txt .
+COPY requirements.txt .
+COPY postinstallscript.sh .
+RUN apt-get update && apt-get install -y sudo $(cat packages.txt)
+RUN pip install --no-cache-dir -r requirements.txt
+RUN bash postinstallscript.sh
+
+# Create a new user and group using the username argument
+RUN groupadd -r ${USER_NAME} && useradd -r -m -g${USER_NAME} ${USER_NAME}
+RUN echo "${USER_NAME}:${USER_PASSWD}" | chpasswd
+RUN usermod -aG sudo ${USER_NAME}
+USER ${USER_NAME}
+ENV USER=${USER_NAME}
+WORKDIR /home/${USER_NAME}/workspace
+
+# Set the prompt to highlight the username
+RUN echo "export PS1='\[\033[01;32m\]\u\[\033[00m\]@\[\033[01;34m\]\h\[\033[00m\]:\[\033[01;36m\]\w\[\033[00m\]\$'" >> /home/${USER_NAME}/.bashrc

docker/packages.txt

@@ -0,0 +1,3 @@
+wget
+curl
+git

docker/ports.txt

@@ -0,0 +1 @@
+-p 6006:6006

docker/postinstallscript.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+# this script will run when build docker image.
+

docker/requirements.txt

@@ -0,0 +1,40 @@
+setuptools>=61.0
+torch==2.0.1
+torchvision==0.15.2
+transformers==4.32.0
+albumentations==1.4.0
+av==11.0.0
+decord==0.6.0
+einops==0.3.0
+fastapi==0.110.0
+accelerate==0.21.0
+gdown==5.1.0
+h5py==3.10.0
+idna==3.6
+imageio==2.34.0
+matplotlib==3.7.5
+numpy==1.24.4
+omegaconf==2.1.1
+opencv-python==4.9.0.80
+opencv-python-headless==4.9.0.80
+pandas==2.0.3
+pillow==10.2.0
+pydub==0.25.1
+pytorch-lightning==1.4.2
+pytorchvideo==0.1.5
+PyYAML==6.0.1
+regex==2023.12.25
+requests==2.31.0
+scikit-learn==1.3.2
+scipy==1.10.1
+six==1.16.0
+tensorboard==2.14.0
+test-tube==0.7.5
+timm==0.9.16
+torchdiffeq==0.2.3
+torchmetrics==0.5.0
+tqdm==4.66.2
+urllib3==2.2.1
+uvicorn==0.27.1
+diffusers==0.24.0
+scikit-video==1.1.11

docker/setup_env.sh

@@ -0,0 +1,11 @@
+# Docker tag for new build image
+TAG=open_sora_plan:dev
+
+# Base docker image tag used by docker build
+BASE_TAG=nvcr.io/nvidia/pytorch:23.05-py3
+
+# User name used in docker container
+USER_NAME=developer
+
+# User password used in docker container
+USER_PASSWD=666666

docs/Contribution_Guidelines.md

@@ -0,0 +1,87 @@
+# Contributing to the Open-Sora Plan Community
+
+The Open-Sora Plan open-source community is a collaborative initiative driven by the community, emphasizing a commitment to being free and void of exploitation. Organized spontaneously by community members, we invite you to contribute to the Open-Sora Plan open-source community and help elevate it to new heights!
+
+## Submitting a Pull Request (PR)
+
+As a contributor, before submitting your request, kindly follow these guidelines:
+
+1. Start by checking the [Open-Sora Plan GitHub](https://github.com/PKU-YuanGroup/Open-Sora-Plan/pulls) to see if there are any open or closed pull requests related to your intended submission. Avoid duplicating existing work.
+
+2. [Fork](https://github.com/PKU-YuanGroup/Open-Sora-Plan/fork) the [open-sora plan](https://github.com/PKU-YuanGroup/Open-Sora-Plan) repository and download your forked repository to your local machine.
+
+   ```bash
+   git clone [your-forked-repository-url]
+   ```
+
+3. Add the original Open-Sora Plan repository as a remote to sync with the latest updates:
+
+   ```bash
+   git remote add upstream https://github.com/PKU-YuanGroup/Open-Sora-Plan
+   ```
+
+4. Sync the code from the main repository to your local machine, and then push it back to your forked remote repository.
+
+   ```
+   # Pull the latest code from the upstream branch
+   git fetch upstream
+   
+   # Switch to the main branch
+   git checkout main
+   
+   # Merge the updates from the upstream branch into main, synchronizing the local main branch with the upstream
+   git merge upstream/main
+   
+   # Additionally, sync the local main branch to the remote branch of your forked repository
+   git push origin main
+   ```
+
+
+   > Note: Sync the code from the main repository before each submission.
+
+5. Create a branch in your forked repository for your changes, ensuring the branch name is meaningful.
+
+   ```bash
+   git checkout -b my-docs-branch main
+   ```
+
+6. While making modifications and committing changes, adhere to our [Commit Message Format](#Commit-Message-Format).
+
+   ```bash
+   git commit -m "[docs]: xxxx"
+   ```
+
+7. Push your changes to your GitHub repository.
+
+   ```bash
+   git push origin my-docs-branch
+   ```
+
+8. Submit a pull request to `Open-Sora-Plan:main` on the GitHub repository page.
+
+## Commit Message Format
+
+Commit messages must include both `<type>` and `<summary>` sections.
+
+```bash
+[<type>]: <summary>
+  │        │
+  │        └─⫸ Briefly describe your changes, without ending with a period.
+  │
+  └─⫸ Commit Type: |docs|feat|fix|refactor|
+```
+
+### Type 
+
+* **docs**: Modify or add documents.
+* **feat**: Introduce a new feature.
+* **fix**: Fix a bug.
+* **refactor**: Restructure code, excluding new features or bug fixes.
+
+### Summary
+
+Describe modifications in English, without ending with a period.
+
+> e.g., git commit -m "[docs]: add a contributing.md file"
+
+This guideline is borrowed by [minisora](https://github.com/mini-sora/minisora). We sincerely appreciate MiniSora authors for their awesome templates. 

docs/Data.md

@@ -0,0 +1,35 @@
+
+**We need more dataset**, please refer to the [open-sora-Dataset](https://github.com/shaodong233/open-sora-Dataset) for details.
+
+
+## Sky
+
+
+This is an un-condition datasets. [Link](https://drive.google.com/open?id=1xWLiU-MBGN7MrsFHQm4_yXmfHBsMbJQo)
+
+```
+sky_timelapse
+├── readme
+├── sky_test
+├── sky_train
+├── test_videofolder.py
+└── video_folder.py
+```
+
+## UCF101
+
+We test the code with UCF-101 dataset. In order to download UCF-101 dataset, you can download the necessary files in [here](https://www.crcv.ucf.edu/data/UCF101.php). The code assumes a `ucf101` directory with the following structure
+```
+UCF-101/
+    ApplyEyeMakeup/
+        v1.avi
+        ...
+    ...
+    YoYo/
+        v1.avi
+        ...
+```
+
+
+## Offline feature extraction
+Coming soon...

docs/EVAL.md

@@ -0,0 +1,110 @@
+# Evaluate the generated videos quality
+
+You can easily calculate the following video quality metrics, which supports the batch-wise process.
+- **CLIP-SCORE**: It uses the pretrained CLIP model to measure the cosine similarity between two modalities.
+- **FVD**: Frechét Video Distance
+- **SSIM**: structural similarity index measure
+- **LPIPS**: learned perceptual image patch similarity
+- **PSNR**: peak-signal-to-noise ratio
+
+# Requirement 
+## Environment
+- install Pytorch (torch>=1.7.1)
+- install CLIP
+    ```
+        pip install git+https://github.com/openai/CLIP.git
+    ```
+- install clip-cose from PyPi
+    ```
+        pip install clip-score
+    ```
+- Other package
+    ```
+        pip install lpips
+        pip install scipy (scipy==1.7.3/1.9.3, if you use 1.11.3, **you will calculate a WRONG FVD VALUE!!!**)
+        pip install numpy
+        pip install pillow
+        pip install torchvision>=0.8.2
+        pip install ftfy
+        pip install regex
+        pip install tqdm
+    ```
+## Pretrain model
+- FVD
+    Before you cacluate FVD, you should first download the FVD pre-trained model. You can manually download any of the following and put it into FVD folder. 
+    - `i3d_torchscript.pt` from [here](https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt) 
+    - `i3d_pretrained_400.pt` from [here](https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI)
+
+
+## Other Notices
+1. Make sure the pixel value of videos should be in [0, 1].
+2. We average SSIM when images have 3 channels, ssim is the only metric extremely sensitive to gray being compared to b/w.
+3. Because the i3d model downsamples in the time dimension, `frames_num` should > 10 when calculating FVD, so FVD calculation begins from 10-th frame, like upper example.
+4. For grayscale videos, we multiply to 3 channels
+5. data input specifications for clip_score
+> - Image Files:All images should be stored in a single directory. The image files can be in either .png or .jpg format.
+> 
+> - Text Files: All text data should be contained in plain text files in a separate directory. These text files should have the extension .txt.    
+>
+> Note: The number of files in the image directory should be exactly equal to the number of files in the text directory. Additionally, the files in the image directory and text directory should be paired by file name. For instance, if there is a cat.png in the image directory, there should be a corresponding cat.txt in the text directory.
+>
+> Directory Structure Example:
+> ```
+>   ├── path/to/image
+>   │   ├── cat.png
+>   │   ├── dog.png
+>   │   └── bird.jpg
+>   └── path/to/text
+>       ├── cat.txt
+>       ├── dog.txt
+>       └── bird.txt
+> ```
+
+6. data input specifications for fvd, psnr, ssim, lpips
+
+> Directory Structure Example:
+> ```
+>   ├── path/to/generated_image
+>   │   ├── cat.mp4
+>   │   ├── dog.mp4
+>   │   └── bird.mp4
+>   └── path/to/real_image
+>       ├── cat.mp4
+>       ├── dog.mp4
+>       └── bird.mp4
+> ```
+
+
+
+# Usage
+
+```
+# you change the file path and need to set the frame_num, resolution etc...
+
+# clip_score cross modality
+cd opensora/eval
+bash script/cal_clip_score.sh
+
+
+
+# fvd 
+cd opensora/eval
+bash script/cal_fvd.sh
+
+# psnr
+cd opensora/eval
+bash eval/script/cal_psnr.sh
+
+
+# ssim
+cd opensora/eval
+bash eval/script/cal_ssim.sh
+
+
+# lpips
+cd opensora/eval
+bash eval/script/cal_lpips.sh
+```
+
+# Acknowledgement
+The evaluation codebase refers to [clip-score](https://github.com/Taited/clip-score) and [common_metrics](https://github.com/JunyaoHu/common_metrics_on_video_quality).

docs/Report-v1.0.0.md

@@ -0,0 +1,131 @@
+# Report v1.0.0
+
+In March 2024, we launched a plan called Open-Sora-Plan, which aims to reproduce the OpenAI [Sora](https://openai.com/sora) through an open-source framework. As a foundational open-source framework, it enables training of video generation models, including Unconditioned Video Generation, Class Video Generation, and Text-to-Video Generation.
+
+**Today, we are thrilled to present Open-Sora-Plan v1.0.0, which significantly enhances video generation quality and text control capabilities.**
+
+Compared with previous video generation model, Open-Sora-Plan v1.0.0 has several improvements:
+
+1. **Efficient training and inference with CausalVideoVAE**. We apply a spatial-temporal compression to the videos by 4×8×8.
+2. **Joint image-video training for better quality**. Our CausalVideoVAE considers the first frame as an image, allowing for the simultaneous encoding of both images and videos in a natural manner. This allows the diffusion model to grasp more spatial-visual details to improve visual quality.
+
+### Open-Source Release
+We open-source the Open-Sora-Plan to facilitate future development of Video Generation in the community. Code, data, model will be made publicly available.
+- Demo: Hugging Face demo [here](https://huggingface.co/spaces/LanguageBind/Open-Sora-Plan-v1.0.0). 🤝 Enjoying the [![Replicate demo and cloud API](https://replicate.com/camenduru/open-sora-plan-512x512/badge)](https://replicate.com/camenduru/open-sora-plan-512x512) and [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/Open-Sora-Plan-jupyter/blob/main/Open_Sora_Plan_jupyter.ipynb), created by [@camenduru](https://github.com/camenduru), who generously supports our research!
+- Code: All training scripts and sample scripts.
+- Model: Both Diffusion Model and CausalVideoVAE [here](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0).
+- Data: Both raw videos and captions [here](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.0.0).
+
+## Gallery
+
+Open-Sora-Plan v1.0.0 supports joint training of images and videos. Here, we present the capabilities of Video/Image Reconstruction and Generation:
+
+### CausalVideoVAE Reconstruction
+
+**Video Reconstruction** with 720×1280. Since github can't upload large video, we put it here: [1](https://streamable.com/gqojal), [2](https://streamable.com/6nu3j8). 
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/c100bb02-2420-48a3-9d7b-4608a41f14aa
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/8aa8f587-d9f1-4e8b-8a82-d3bf9ba91d68
+
+**Image Reconstruction** in 1536×1024.
+
+<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/1684c3ec-245d-4a60-865c-b8946d788eb9" width="45%"/> <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/46ef714e-3e5b-492c-aec4-3793cb2260b5" width="45%"/>
+
+**Text-to-Video Generation** with 65×1024×1024
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/2641a8aa-66ac-4cda-8279-86b2e6a6e011
+
+**Text-to-Video Generation** with 65×512×512
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/37e3107e-56b3-4b09-8920-fa1d8d144b9e
+
+
+**Text-to-Image Generation** with 512×512
+
+![download](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/491d72bc-e762-48ff-bdcc-cc69350f56d6)
+
+## Detailed Technical Report
+
+### CausalVideoVAE
+
+#### Model Structure
+
+![image](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/e3c8b35d-a217-4d96-b2e9-5c248a2859c8)
+
+The CausalVideoVAE architecture inherits from the [Stable-Diffusion Image VAE](https://github.com/CompVis/stable-diffusion/tree/main). To ensure that the pretrained weights of the Image VAE can be seamlessly applied to the Video VAE, the model structure has been designed as follows:
+
+1. **CausalConv3D**: Converting Conv2D to CausalConv3D enables joint training of image and video data. CausalConv3D applies a special treatment to the first frame, as it does not have access to subsequent frames. For more specific details, please refer to https://github.com/PKU-YuanGroup/Open-Sora-Plan/pull/145
+
+2. **Initialization**: There are two common [methods](https://github.com/hassony2/inflated_convnets_pytorch/blob/master/src/inflate.py#L5) to expand Conv2D to Conv3D: average initialization and center initialization. But we employ a specific initialization method (tail initialization). This initialization method ensures that without any training, the model is capable of directly reconstructing images, and even videos.
+
+#### Training Details
+
+<img width="833" alt="image" src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/9ffb6dc4-23f6-4274-a066-bbebc7522a14">
+
+We present the loss curves for two distinct initialization methods under 17×256×256. The yellow curve represents the loss using tail init, while the blue curve corresponds to the loss from center initialization. As shown in the graph, tail initialization demonstrates better performance on the loss curve. Additionally, we found that center initialization leads to error accumulation, causing the collapse over extended durations.
+
+#### Inference Tricks
+Despite the VAE in Diffusion training being frozen, we still find it challenging to afford the cost of the CausalVideoVAE. In our case, with 80GB of GPU memory, we can only infer a video of either 256×512×512 or 32×1024×1024 resolution using half-precision, which limits our ability to scale up to longer and higher-resolution videos. Therefore, we adopt tile convolution, which allows us to infer videos of arbitrary duration or resolution with nearly constant memory usage.
+
+### Data Construction
+We define a high-quality video dataset based on two core principles: (1) No content-unrelated watermarks. (2) High-quality and dense captions.
+
+**For principles 1**, we crawled approximately 40,000 videos from open-source websites under the CC0 license. Specifically, we obtained 1,244 videos from [mixkit](https://mixkit.co/), 7,408 videos from [pexels](https://www.pexels.com/), and 31,617 videos from [pixabay](https://pixabay.com/). These videos adhere to the principle of having no content-unrelated watermarks. According to the scene transformation and clipping script provided by [Panda70M](https://github.com/snap-research/Panda-70M/blob/main/splitting/README.md), we have divided these videos into approximately 434,000 video clips. In fact, based on our clipping results, 99% of the videos obtained from these online sources are found to contain single scenes. Additionally, we have observed that over 60% of the crawled data comprises landscape videos.
+
+**For principles 2**, it is challenging to directly crawl a large quantity of high-quality dense captions from the internet. Therefore, we utilize a mature Image-captioner model to obtain high-quality dense captions. We conducted ablation experiments on two multimodal large models: [ShareGPT4V-Captioner-7B](https://github.com/InternLM/InternLM-XComposer/blob/main/projects/ShareGPT4V/README.md) and [LLaVA-1.6-34B](https://github.com/haotian-liu/LLaVA). The former is specifically designed for caption generation, while the latter is a general-purpose multimodal large model. After conducting our ablation experiments, we found that they are comparable in performance. However, there is a significant difference in their inference speed on the A800 GPU: 40s/it of batch size of 12 for ShareGPT4V-Captioner-7B, 15s/it of batch size of 1 for LLaVA-1.6-34B. We open-source all annotations [here](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.0.0)。
+
+### Training Diffusion Model
+Similar to previous work, we employ a multi-stage cascaded training approach, which consumes a total of 2,048 A800 GPU hours. We found that joint training with images significantly accelerates model convergence and enhances visual perception, aligning with the findings of [Latte](https://github.com/Vchitect/Latte). Below is our training card:
+
+| Name | Stage 1 | Stage 2 | Stage 3 | Stage 4 |
+|---|---|---|---|---|
+| Training Video Size | 17×256×256 |  65×256×256 | 65×512×512 |  65×1024×1024 | 
+| Compute (#A800 GPU x #Hours) | 32 × 40 |  32 × 18 |  32 × 6 |  Under training | 
+| Checkpoint | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/17x256x256) | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/65x256x256) |  [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/65x512x512) |  Under training | 
+| Log | [wandb](https://api.wandb.ai/links/linbin/p6n3evym) |  [wandb](https://api.wandb.ai/links/linbin/t2g53sew) |  [wandb](https://api.wandb.ai/links/linbin/uomr0xzb) | Under training | 
+| Training Data | ~40k videos |  ~40k videos |  ~40k videos |  ~40k videos | 
+
+## Next Release Preview
+### CausalVideoVAE
+Currently, the released version of CausalVideoVAE (v1.0.0) has two main drawbacks: **motion blurring** and **gridding effect**. We have made a series of improvements to CausalVideoVAE to reduce its inference cost and enhance its performance. We are currently referring to this enhanced version as the "preview version," which will be released in the next update. Preview reconstruction is as follows:
+
+**1 min Video Reconstruction with 720×1280**. Since github can't put too big video, we put it here: [origin video](https://streamable.com/u4onbb), [reconstruction video](https://streamable.com/qt8ncc). 
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/cdcfa9a3-4de0-42d4-94c0-0669710e407b
+
+We randomly selected 100 samples from the validation set of Kinetics-400 for evaluation, and the results are presented in the following table:
+
+|  | SSIM↑ | LPIPS↓ | PSNR↑ | FLOLPIPS↓ |
+|---|---|---|---|---|
+| v1.0.0 | 0.829 |  0.106 |  27.171 |  0.119 | 
+| Preview | 0.877 |  0.064 |  29.695 |  0.070 | 
+
+#### Motion Blurring
+
+| **v1.0.0** | **Preview** |
+| --- | --- |
+| ![6862cae0-b1b6-48d1-bd11-84348cf42b42](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/f815636f-fb38-4891-918b-50b1f9aa086d)  | ![9189da06-ef2c-42e6-ad34-bd702a6f538e](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/1e413f50-a785-485a-9851-a1449f952f1c)  |
+
+#### Gridding effect
+
+| **v1.0.0** | **Preview** |
+| --- | --- |
+| ![img](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/7fec5bed-3c83-4ee9-baef-4a3dacafc658)  | ![img](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/4f41b432-a3ef-484e-a492-8afd8a691bf7)  |
+
+### Data Construction
+
+**Data source**. As mentioned earlier, over 60% of our dataset consists of landscape videos. This implies that our ability to generate videos in other domains is limited. However, most of the current large-scale open-source datasets are primarily obtained through web scraping from platforms like YouTube. While these datasets provide a vast quantity of videos, we have concerns about the quality of the videos themselves. Therefore, we will continue to collect high-quality datasets and also welcome recommendations from the open-source community.
+
+**Caption Generation Pipeline**. As the video duration increases, we need to consider more efficient methods for video caption generation instead of relying solely on large multimodal image models. We are currently developing a new video caption generation pipeline that provides robust support for long videos. We are excited to share more details with you in the near future. Stay tuned!
+
+### Training Diffusion Model
+Although v1.0.0 has shown promising results, we acknowledge that we still have a ways to go to reach the level of Sora. In our upcoming work, we will primarily focus on three aspects:
+
+1. **Training support for dynamic resolution and duration**: We aim to develop techniques that enable training models with varying resolutions and durations, allowing for more flexible and adaptable training processes.
+
+2. **Support for longer video generation**: We will explore methods to extend the generation capabilities of our models, enabling them to produce longer videos beyond the current limitations.
+
+3. **Enhanced conditional control**: We seek to enhance the conditional control capabilities of our models, providing users with more options and control over the generated videos.
+
+Furthermore, through careful observation of the generated videos, we have noticed the presence of some non-physiological speckles or abnormal flow. This can be attributed to the limited performance of CausalVideoVAE, as mentioned earlier. In future experiments, we plan to retrain a diffusion model using a more powerful version of CausalVideoVAE to address these issues.

docs/VQVAE.md

@@ -0,0 +1,57 @@
+# VQVAE Documentation
+
+# Introduction
+
+Vector Quantized Variational AutoEncoders (VQ-VAE) is a type of autoencoder that uses a discrete latent representation. It is particularly useful for tasks that require discrete latent variables, such as text-to-speech and video generation.
+
+# Usage
+
+## Initialization
+
+To initialize a VQVAE model, you can use the `VideoGPTVQVAE` class. This class is a part of the `opensora.models.ae` module.
+
+```python
+from opensora.models.ae import VideoGPTVQVAE
+
+vqvae = VideoGPTVQVAE()
+```
+
+### Training
+
+To train the VQVAE model, you can use the `train_videogpt.sh` script. This script will train the model using the parameters specified in the script.
+
+```bash
+bash scripts/videogpt/train_videogpt.sh
+```
+
+### Loading Pretrained Models
+
+You can load a pretrained model using the `download_and_load_model` method. This method will download the checkpoint file and load the model.
+
+```python
+vqvae = VideoGPTVQVAE.download_and_load_model("bair_stride4x2x2")
+```
+
+Alternatively, you can load a model from a checkpoint using the `load_from_checkpoint` method.
+
+```python
+vqvae = VQVAEModel.load_from_checkpoint("results/VQVAE/checkpoint-1000")
+```
+
+### Encoding and Decoding
+
+You can encode a video using the `encode` method. This method will return the encodings and embeddings of the video.
+
+```python
+encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True)
+```
+
+You can reconstruct a video from its encodings using the decode method.
+
+```python
+video_recon = vqvae.decode(encodings)
+```
+
+## Testing
+
+You can test the VQVAE model by reconstructing a video. The `examples/rec_video.py` script provides an example of how to do this.

examples/get_latents_std.py

@@ -0,0 +1,38 @@
+import torch
+from torch.utils.data import DataLoader, Subset
+import sys
+sys.path.append(".")
+from opensora.models.ae.videobase import CausalVAEModel, CausalVAEDataset
+
+num_workers = 4
+batch_size = 12
+
+torch.manual_seed(0)
+torch.set_grad_enabled(False)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+pretrained_model_name_or_path = 'results/causalvae/checkpoint-26000'
+data_path = '/remote-home1/dataset/UCF-101'
+video_num_frames = 17
+resolution = 128
+sample_rate = 10
+
+vae = CausalVAEModel.load_from_checkpoint(pretrained_model_name_or_path)
+vae.to(device)
+
+dataset = CausalVAEDataset(data_path, sequence_length=video_num_frames, resolution=resolution, sample_rate=sample_rate)
+subset_indices = list(range(1000))
+subset_dataset = Subset(dataset, subset_indices)
+loader = DataLoader(subset_dataset, batch_size=8, pin_memory=True)
+
+all_latents = []
+for video_data in loader:
+    video_data = video_data['video'].to(device)
+    latents = vae.encode(video_data).sample()
+    all_latents.append(video_data.cpu())
+
+all_latents_tensor = torch.cat(all_latents)
+std = all_latents_tensor.std().item()
+normalizer = 1 / std
+print(f'{normalizer = }')

examples/prompt_list_0.txt

@@ -0,0 +1,16 @@
+A quiet beach at dawn, the waves gently lapping at the shore and the sky painted in pastel hues.
+A quiet beach at dawn, the waves softly lapping at the shore, pink and orange hues painting the sky, offering a moment of solitude and reflection.
+The majestic beauty of a waterfall cascading down a cliff into a serene lake.
+Sunset over the sea.
+a cat wearing sunglasses and working as a lifeguard at pool.
+Slow pan upward of blazing oak fire in an indoor fireplace.
+Yellow and black tropical fish dart through the sea.
+a serene winter scene in a forest. The forest is blanketed in a thick layer of snow, which has settled on the branches of the trees, creating a canopy of white. The trees, a mix of evergreens and deciduous, stand tall and silent, their forms partially obscured by the snow. The ground is a uniform white, with no visible tracks or signs of human activity. The sun is low in the sky, casting a warm glow that contrasts with the cool tones of the snow. The light filters through the trees, creating a soft, diffused illumination that highlights the texture of the snow and the contours of the trees. The overall style of the scene is naturalistic, with a focus on the tranquility and beauty of the winter landscape.
+a dynamic interaction between the ocean and a large rock. The rock, with its rough texture and jagged edges, is partially submerged in the water, suggesting it is a natural feature of the coastline. The water around the rock is in motion, with white foam and waves crashing against the rock, indicating the force of the ocean's movement. The background is a vast expanse of the ocean, with small ripples and waves, suggesting a moderate sea state. The overall style of the scene is a realistic depiction of a natural landscape, with a focus on the interplay between the rock and the water.
+A serene waterfall cascading down moss-covered rocks, its soothing sound creating a harmonious symphony with nature.
+A soaring drone footage captures the majestic beauty of a coastal cliff, its red and yellow stratified rock faces rich in color and against the vibrant turquoise of the sea. Seabirds can be seen taking flight around the cliff's precipices. As the drone slowly moves from different angles, the changing sunlight casts shifting shadows that highlight the rugged textures of the cliff and the surrounding calm sea. The water gently laps at the rock base and the greenery that clings to the top of the cliff, and the scene gives a sense of peaceful isolation at the fringes of the ocean. The video captures the essence of pristine natural beauty untouched by human structures.
+The video captures the majestic beauty of a waterfall cascading down a cliff into a serene lake. The waterfall, with its powerful flow, is the central focus of the video. The surrounding landscape is lush and green, with trees and foliage adding to the natural beauty of the scene. The camera angle provides a bird's eye view of the waterfall, allowing viewers to appreciate the full height and grandeur of the waterfall. The video is a stunning representation of nature's power and beauty.
+A vibrant scene of a snowy mountain landscape. The sky is filled with a multitude of colorful hot air balloons, each floating at different heights, creating a dynamic and lively atmosphere. The balloons are scattered across the sky, some closer to the viewer, others further away, adding depth to the scene.  Below, the mountainous terrain is blanketed in a thick layer of snow, with a few patches of bare earth visible here and there. The snow-covered mountains provide a stark contrast to the colorful balloons, enhancing the visual appeal of the scene.
+A serene underwater scene featuring a sea turtle swimming through a coral reef. The turtle, with its greenish-brown shell, is the main focus of the video, swimming gracefully towards the right side of the frame. The coral reef, teeming with life, is visible in the background, providing a vibrant and colorful backdrop to the turtle's journey. Several small fish, darting around the turtle, add a sense of movement and dynamism to the scene.
+A snowy forest landscape with a dirt road running through it. The road is flanked by trees covered in snow, and the ground is also covered in snow. The sun is shining, creating a bright and serene atmosphere. The road appears to be empty, and there are no people or animals visible in the video. The style of the video is a natural landscape shot, with a focus on the beauty of the snowy forest and the peacefulness of the road.
+The dynamic movement of tall, wispy grasses swaying in the wind. The sky above is filled with clouds, creating a dramatic backdrop. The sunlight pierces through the clouds, casting a warm glow on the scene. The grasses are a mix of green and brown, indicating a change in seasons. The overall style of the video is naturalistic, capturing the beauty of the landscape in a realistic manner. The focus is on the grasses and their movement, with the sky serving as a secondary element. The video does not contain any human or animal elements.

examples/rec_image.py

@@ -0,0 +1,57 @@
+import sys
+sys.path.append(".")
+from PIL import Image
+import torch
+from torchvision.transforms import ToTensor, Compose, Resize, Normalize
+from torch.nn import functional as F
+from opensora.models.ae.videobase import CausalVAEModel
+import argparse
+import numpy as np
+
+def preprocess(video_data: torch.Tensor, short_size: int = 128) -> torch.Tensor:
+    transform = Compose(
+        [
+            ToTensor(),
+            Normalize((0.5), (0.5)),
+            Resize(size=short_size),
+        ]
+    )
+    outputs = transform(video_data)
+    outputs = outputs.unsqueeze(0).unsqueeze(2)
+    return outputs
+
+def main(args: argparse.Namespace):
+    image_path = args.image_path
+    resolution = args.resolution
+    device = args.device
+    
+    vqvae = CausalVAEModel.load_from_checkpoint(args.ckpt)
+    vqvae.eval()
+    vqvae = vqvae.to(device)
+    
+    with torch.no_grad():
+        x_vae = preprocess(Image.open(image_path), resolution)
+        x_vae = x_vae.to(device)
+        latents = vqvae.encode(x_vae)
+        recon = vqvae.decode(latents.sample())
+    x = recon[0, :, 0, :, :]
+    x = x.squeeze()
+    x = x.detach().cpu().numpy()
+    x = np.clip(x, -1, 1)
+    x = (x + 1) / 2
+    x = (255*x).astype(np.uint8)
+    x = x.transpose(1,2,0)
+    image = Image.fromarray(x)
+    image.save(args.rec_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image-path', type=str, default='')
+    parser.add_argument('--rec-path', type=str, default='')
+    parser.add_argument('--ckpt', type=str, default='')
+    parser.add_argument('--resolution', type=int, default=336)
+    parser.add_argument('--device', type=str, default='cuda')
+    
+    args = parser.parse_args()
+    main(args)

examples/rec_imvi_vae.py

@@ -0,0 +1,159 @@
+import math
+import random
+import argparse
+from typing import Optional
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+import torch
+from PIL import Image
+from decord import VideoReader, cpu
+from torch.nn import functional as F
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.transforms import Lambda, Compose
+
+import sys
+sys.path.append(".")
+from opensora.dataset.transform import CenterCropVideo, resize
+from opensora.models.ae.videobase import CausalVAEModel
+
+
+def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None:
+    height, width, channels = image_array[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
+
+    for image in image_array:
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        video_writer.write(image_rgb)
+
+    video_writer.release()
+
+def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None:
+    x = x.detach().cpu()
+    x = torch.clamp(x, -1, 1)
+    x = (x + 1) / 2
+    x = x.permute(1, 2, 3, 0).numpy()
+    x = (255*x).astype(np.uint8)
+    array_to_video(x, fps=fps, output_file=output_file)
+    return
+
+def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor:
+    decord_vr = VideoReader(video_path, ctx=cpu(0))
+    total_frames = len(decord_vr)
+    sample_frames_len = sample_rate * num_frames
+
+    if total_frames > sample_frames_len:
+        s = random.randint(0, total_frames - sample_frames_len - 1)
+        s = 0
+        e = s + sample_frames_len
+        num_frames = num_frames
+    else:
+        s = 0
+        e = total_frames
+        num_frames = int(total_frames / sample_frames_len * num_frames)
+        print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path,
+              total_frames)
+
+
+    frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
+    video_data = decord_vr.get_batch(frame_id_list).asnumpy()
+    video_data = torch.from_numpy(video_data)
+    video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
+    return video_data
+
+
+class ResizeVideo:
+    def __init__(
+            self,
+            size,
+            interpolation_mode="bilinear",
+    ):
+        self.size = size
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        _, _, h, w = clip.shape
+        if w < h:
+            new_h = int(math.floor((float(h) / w) * self.size))
+            new_w = self.size
+        else:
+            new_h = self.size
+            new_w = int(math.floor((float(w) / h) * self.size))
+        return torch.nn.functional.interpolate(
+            clip, size=(new_h, new_w), mode=self.interpolation_mode, align_corners=False, antialias=True
+        )
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor:
+
+    transform = Compose(
+        [
+            Lambda(lambda x: ((x / 255.0) * 2 - 1)),
+            ResizeVideo(size=short_size),
+            CenterCropVideo(crop_size) if crop_size is not None else Lambda(lambda x: x),
+        ]
+    )
+
+    video_outputs = transform(video_data)
+    video_outputs = torch.unsqueeze(video_outputs, 0)
+
+    return video_outputs
+
+
+def main(args: argparse.Namespace):
+    video_path = args.video_path
+    num_frames = args.num_frames
+    resolution = args.resolution
+    crop_size = args.crop_size
+    sample_fps = args.sample_fps
+    sample_rate = args.sample_rate
+    device = args.device
+    vqvae = CausalVAEModel.from_pretrained(args.ckpt)
+    if args.enable_tiling:
+        vqvae.enable_tiling()
+        vqvae.tile_overlap_factor = args.tile_overlap_factor
+    vqvae.eval()
+    vqvae = vqvae.to(device)
+    vqvae = vqvae.to(torch.float16)
+
+    with torch.no_grad():
+        x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size)
+        x_vae = x_vae.to(device)  # b c t h w
+        x_vae = x_vae.to(torch.float16)
+        latents = vqvae.encode(x_vae).sample().to(torch.float16)
+        video_recon = vqvae.decode(latents)
+        
+    if video_recon.shape[2] == 1:
+        x = video_recon[0, :, 0, :, :]
+        x = x.squeeze()
+        x = x.detach().cpu().numpy()
+        x = np.clip(x, -1, 1)
+        x = (x + 1) / 2
+        x = (255 * x).astype(np.uint8)
+        x = x.transpose(1, 2, 0)
+        image = Image.fromarray(x)
+        image.save(args.rec_path.replace('mp4', 'jpg'))
+    else:
+        custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--video-path', type=str, default='')
+    parser.add_argument('--rec-path', type=str, default='')
+    parser.add_argument('--ckpt', type=str, default='results/pretrained')
+    parser.add_argument('--sample-fps', type=int, default=30)
+    parser.add_argument('--resolution', type=int, default=336)
+    parser.add_argument('--crop-size', type=int, default=None)
+    parser.add_argument('--num-frames', type=int, default=100)
+    parser.add_argument('--sample-rate', type=int, default=1)
+    parser.add_argument('--device', type=str, default="cuda")
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument('--enable_tiling', action='store_true')
+    
+    args = parser.parse_args()
+    main(args)

examples/rec_video.py

@@ -0,0 +1,120 @@
+import random
+import argparse
+from typing import Optional
+
+import cv2
+import imageio
+import numpy as np
+import numpy.typing as npt
+import torch
+from decord import VideoReader, cpu
+from torch.nn import functional as F
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms._transforms_video import RandomCropVideo
+
+import sys
+sys.path.append(".")
+from opensora.models.ae import VQVAEModel
+
+
+def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None:
+    height, width, channels = image_array[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')    # type: ignore
+    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
+
+    for image in image_array:
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        video_writer.write(image_rgb)
+
+    video_writer.release()
+
+def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None:
+    x = x.detach().cpu()
+    x = torch.clamp(x, -0.5, 0.5)
+    x = (x + 0.5)
+    x = x.permute(1, 2, 3, 0).numpy()  # (C, T, H, W) -> (T, H, W, C)
+    x = (255*x).astype(np.uint8)
+    # array_to_video(x, fps=fps, output_file=output_file)
+    imageio.mimwrite(output_file, x, fps=fps, quality=9)
+    return
+
+def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor:
+    decord_vr = VideoReader(video_path, ctx=cpu(0))
+    total_frames = len(decord_vr)
+    sample_frames_len = sample_rate * num_frames
+
+    if total_frames > sample_frames_len:
+        s = random.randint(0, total_frames - sample_frames_len - 1)
+        e = s + sample_frames_len
+        num_frames = num_frames
+    else:
+        s = 0
+        e = total_frames
+        num_frames = int(total_frames / sample_frames_len * num_frames)
+        print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path,
+              total_frames)
+
+
+    frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
+    video_data = decord_vr.get_batch(frame_id_list).asnumpy()
+    video_data = torch.from_numpy(video_data)
+    video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
+    return video_data
+
+def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor:
+
+    transform = Compose(
+        [
+            # UniformTemporalSubsample(num_frames),
+            Lambda(lambda x: ((x / 255.0) - 0.5)),
+            # NormalizeVideo(mean=OPENAI_DATASET_MEAN, std=OPENAI_DATASET_STD),
+            ShortSideScale(size=short_size),
+            RandomCropVideo(size=crop_size) if crop_size is not None else Lambda(lambda x: x),
+            # RandomHorizontalFlipVideo(p=0.5),
+        ]
+    )
+
+    video_outputs = transform(video_data)
+    video_outputs = torch.unsqueeze(video_outputs, 0)
+
+    return video_outputs
+
+
+def main(args: argparse.Namespace):
+    video_path = args.video_path
+    num_frames = args.num_frames
+    resolution = args.resolution
+    crop_size = args.crop_size
+    sample_fps = args.sample_fps
+    sample_rate = args.sample_rate
+    device = torch.device('cuda')
+    if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']:
+        vqvae = VQVAEModel.download_and_load_model(args.ckpt)
+    else:
+        vqvae = VQVAEModel.load_from_checkpoint(args.ckpt)
+    vqvae.eval()
+    vqvae = vqvae.to(device)
+
+    with torch.no_grad():
+        x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size)
+        x_vae = x_vae.to(device)
+        encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True)
+        video_recon = vqvae.decode(encodings)
+
+    # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4')
+    custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--video-path', type=str, default='')
+    parser.add_argument('--rec-path', type=str, default='')
+    parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4')
+    parser.add_argument('--sample-fps', type=int, default=30)
+    parser.add_argument('--resolution', type=int, default=336)
+    parser.add_argument('--crop-size', type=int, default=None)
+    parser.add_argument('--num-frames', type=int, default=100)
+    parser.add_argument('--sample-rate', type=int, default=1)
+    args = parser.parse_args()
+    main(args)

examples/rec_video_ae.py

@@ -0,0 +1,120 @@
+import random
+import argparse
+from typing import Optional
+
+import cv2
+import imageio
+import numpy as np
+import numpy.typing as npt
+import torch
+from decord import VideoReader, cpu
+from torch.nn import functional as F
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms._transforms_video import RandomCropVideo
+
+import sys
+sys.path.append(".")
+from opensora.models.ae import VQVAEModel
+
+
+def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None:
+    height, width, channels = image_array[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')    # type: ignore
+    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
+
+    for image in image_array:
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        video_writer.write(image_rgb)
+
+    video_writer.release()
+
+def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None:
+    x = x.detach().cpu()
+    x = torch.clamp(x, -0.5, 0.5)
+    x = (x + 0.5)
+    x = x.permute(1, 2, 3, 0).numpy()  # (C, T, H, W) -> (T, H, W, C)
+    x = (255*x).astype(np.uint8)
+    # array_to_video(x, fps=fps, output_file=output_file)
+    imageio.mimwrite(output_file, x, fps=fps, quality=9)
+    return
+
+def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor:
+    decord_vr = VideoReader(video_path, ctx=cpu(0))
+    total_frames = len(decord_vr)
+    sample_frames_len = sample_rate * num_frames
+
+    if total_frames > sample_frames_len:
+        s = random.randint(0, total_frames - sample_frames_len - 1)
+        e = s + sample_frames_len
+        num_frames = num_frames
+    else:
+        s = 0
+        e = total_frames
+        num_frames = int(total_frames / sample_frames_len * num_frames)
+        print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path,
+              total_frames)
+
+
+    frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
+    video_data = decord_vr.get_batch(frame_id_list).asnumpy()
+    video_data = torch.from_numpy(video_data)
+    video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
+    return video_data
+
+def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor:
+
+    transform = Compose(
+        [
+            # UniformTemporalSubsample(num_frames),
+            Lambda(lambda x: ((x / 255.0) - 0.5)),
+            # NormalizeVideo(mean=OPENAI_DATASET_MEAN, std=OPENAI_DATASET_STD),
+            ShortSideScale(size=short_size),
+            RandomCropVideo(size=crop_size) if crop_size is not None else Lambda(lambda x: x),
+            # RandomHorizontalFlipVideo(p=0.5),
+        ]
+    )
+
+    video_outputs = transform(video_data)
+    video_outputs = torch.unsqueeze(video_outputs, 0)
+
+    return video_outputs
+
+
+def main(args: argparse.Namespace):
+    video_path = args.video_path
+    num_frames = args.num_frames
+    resolution = args.resolution
+    crop_size = args.crop_size
+    sample_fps = args.sample_fps
+    sample_rate = args.sample_rate
+    device = torch.device('cuda')
+    if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']:
+        vqvae = VQVAEModel.download_and_load_model(args.ckpt)
+    else:
+        vqvae = VQVAEModel.load_from_checkpoint(args.ckpt)
+    vqvae.eval()
+    vqvae = vqvae.to(device)
+
+    with torch.no_grad():
+        x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size)
+        x_vae = x_vae.to(device)
+        encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True)
+        video_recon = vqvae.decode(encodings)
+
+    # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4')
+    custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--video-path', type=str, default='')
+    parser.add_argument('--rec-path', type=str, default='')
+    parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4')
+    parser.add_argument('--sample-fps', type=int, default=30)
+    parser.add_argument('--resolution', type=int, default=336)
+    parser.add_argument('--crop-size', type=int, default=None)
+    parser.add_argument('--num-frames', type=int, default=100)
+    parser.add_argument('--sample-rate', type=int, default=1)
+    args = parser.parse_args()
+    main(args)

examples/rec_video_vae.py

@@ -0,0 +1,275 @@
+import random
+import argparse
+import cv2
+from tqdm import tqdm
+import numpy as np
+import numpy.typing as npt
+import torch
+from decord import VideoReader, cpu
+from torch.nn import functional as F
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms._transforms_video import CenterCropVideo
+import sys
+from torch.utils.data import Dataset, DataLoader, Subset
+import os
+
+sys.path.append(".")
+from opensora.models.ae.videobase import CausalVAEModel
+import torch.nn as nn
+
+
+def array_to_video(
+    image_array: npt.NDArray, fps: float = 30.0, output_file: str = "output_video.mp4"
+) -> None:
+    height, width, channels = image_array[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
+
+    for image in image_array:
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        video_writer.write(image_rgb)
+
+    video_writer.release()
+
+
+def custom_to_video(
+    x: torch.Tensor, fps: float = 2.0, output_file: str = "output_video.mp4"
+) -> None:
+    x = x.detach().cpu()
+    x = torch.clamp(x, -1, 1)
+    x = (x + 1) / 2
+    x = x.permute(1, 2, 3, 0).float().numpy()
+    x = (255 * x).astype(np.uint8)
+    array_to_video(x, fps=fps, output_file=output_file)
+    return
+
+
+def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor:
+    decord_vr = VideoReader(video_path, ctx=cpu(0), num_threads=8)
+    total_frames = len(decord_vr)
+    sample_frames_len = sample_rate * num_frames
+
+    if total_frames > sample_frames_len:
+        s = 0
+        e = s + sample_frames_len
+        num_frames = num_frames
+    else:
+        s = 0
+        e = total_frames
+        num_frames = int(total_frames / sample_frames_len * num_frames)
+        print(
+            f"sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}",
+            video_path,
+            total_frames,
+        )
+
+    frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
+    video_data = decord_vr.get_batch(frame_id_list).asnumpy()
+    video_data = torch.from_numpy(video_data)
+    video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
+    return video_data
+
+
+class RealVideoDataset(Dataset):
+    def __init__(
+        self,
+        real_video_dir,
+        num_frames,
+        sample_rate=1,
+        crop_size=None,
+        resolution=128,
+    ) -> None:
+        super().__init__()
+        self.real_video_files = self._combine_without_prefix(real_video_dir)
+        self.num_frames = num_frames
+        self.sample_rate = sample_rate
+        self.crop_size = crop_size
+        self.short_size = resolution
+
+    def __len__(self):
+        return len(self.real_video_files)
+
+    def __getitem__(self, index):
+        if index >= len(self):
+            raise IndexError
+        real_video_file = self.real_video_files[index]
+        real_video_tensor = self._load_video(real_video_file)
+        video_name = os.path.basename(real_video_file)
+        return {'video': real_video_tensor, 'file_name': video_name }
+
+    def _load_video(self, video_path):
+        num_frames = self.num_frames
+        sample_rate = self.sample_rate
+        decord_vr = VideoReader(video_path, ctx=cpu(0))
+        total_frames = len(decord_vr)
+        sample_frames_len = sample_rate * num_frames
+
+        if total_frames > sample_frames_len:
+            s = 0
+            e = s + sample_frames_len
+            num_frames = num_frames
+        else:
+            s = 0
+            e = total_frames
+            num_frames = int(total_frames / sample_frames_len * num_frames)
+            print(
+                f"sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}",
+                video_path,
+                total_frames,
+            )
+
+        frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
+        video_data = decord_vr.get_batch(frame_id_list).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(3, 0, 1, 2)
+        return _preprocess(
+            video_data, short_size=self.short_size, crop_size=self.crop_size
+        )
+
+    def _combine_without_prefix(self, folder_path, prefix="."):
+        folder = []
+        for name in os.listdir(folder_path):
+            if name[0] == prefix:
+                continue
+            folder.append(os.path.join(folder_path, name))
+        folder.sort()
+        return folder
+
+def resize(x, resolution):
+    height, width = x.shape[-2:]
+    aspect_ratio = width / height
+    if width <= height:
+        new_width = resolution
+        new_height = int(resolution / aspect_ratio)
+    else:
+        new_height = resolution
+        new_width = int(resolution * aspect_ratio)
+    resized_x = F.interpolate(x, size=(new_height, new_width), mode='bilinear', align_corners=True, antialias=True)
+    return resized_x
+
+def _preprocess(video_data, short_size=128, crop_size=None):
+    transform = Compose(
+        [
+            Lambda(lambda x: ((x / 255.0) * 2 - 1)),
+            Lambda(lambda x: resize(x, short_size)),
+            (
+                CenterCropVideo(crop_size=crop_size)
+                if crop_size is not None
+                else Lambda(lambda x: x)
+            ),
+        ]
+    )
+    video_outputs = transform(video_data)
+    video_outputs = _format_video_shape(video_outputs)
+    return video_outputs
+
+
+def _format_video_shape(video, time_compress=4, spatial_compress=8):
+    time = video.shape[1]
+    height = video.shape[2]
+    width = video.shape[3]
+    new_time = (
+        (time - (time - 1) % time_compress)
+        if (time - 1) % time_compress != 0
+        else time
+    )
+    new_height = (
+        (height - (height) % spatial_compress)
+        if height % spatial_compress != 0
+        else height
+    )
+    new_width = (
+        (width - (width) % spatial_compress) if width % spatial_compress != 0 else width
+    )
+    return video[:, :new_time, :new_height, :new_width]
+
+
+@torch.no_grad()
+def main(args: argparse.Namespace):
+    real_video_dir = args.real_video_dir
+    generated_video_dir = args.generated_video_dir
+    ckpt = args.ckpt
+    sample_rate = args.sample_rate
+    resolution = args.resolution
+    crop_size = args.crop_size
+    num_frames = args.num_frames
+    sample_rate = args.sample_rate
+    device = args.device
+    sample_fps = args.sample_fps
+    batch_size = args.batch_size
+    num_workers = args.num_workers
+    subset_size = args.subset_size
+    
+    if not os.path.exists(args.generated_video_dir):
+        os.makedirs(args.generated_video_dir, exist_ok=True)
+    
+    data_type = torch.bfloat16
+    
+    # ---- Load Model ----
+    device = args.device
+    vqvae = CausalVAEModel.from_pretrained(args.ckpt)
+    vqvae = vqvae.to(device).to(data_type)
+    if args.enable_tiling:
+        vqvae.enable_tiling()
+        vqvae.tile_overlap_factor = args.tile_overlap_factor
+    # ---- Load Model ----
+
+    # ---- Prepare Dataset ----
+    dataset = RealVideoDataset(
+        real_video_dir=real_video_dir,
+        num_frames=num_frames,
+        sample_rate=sample_rate,
+        crop_size=crop_size,
+        resolution=resolution,
+    )
+    
+    if subset_size:
+        indices = range(subset_size)
+        dataset = Subset(dataset, indices=indices)
+        
+    dataloader = DataLoader(
+        dataset, batch_size=batch_size, pin_memory=True, num_workers=num_workers
+    )
+    # ---- Prepare Dataset
+
+    # ---- Inference ----
+    for batch in tqdm(dataloader):
+        x, file_names = batch['video'], batch['file_name']
+        x = x.to(device=device, dtype=data_type)  # b c t h w
+        latents = vqvae.encode(x).sample().to(data_type)
+        video_recon = vqvae.decode(latents)
+        for idx, video in enumerate(video_recon):
+            output_path = os.path.join(generated_video_dir, file_names[idx])
+            if args.output_origin:
+                os.makedirs(os.path.join(generated_video_dir, "origin/"), exist_ok=True)
+                origin_output_path = os.path.join(generated_video_dir, "origin/", file_names[idx])
+                custom_to_video(
+                    x[idx], fps=sample_fps / sample_rate, output_file=origin_output_path
+                )
+            custom_to_video(
+                video, fps=sample_fps / sample_rate, output_file=output_path
+            )
+    # ---- Inference ----
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--real_video_dir", type=str, default="")
+    parser.add_argument("--generated_video_dir", type=str, default="")
+    parser.add_argument("--ckpt", type=str, default="")
+    parser.add_argument("--sample_fps", type=int, default=30)
+    parser.add_argument("--resolution", type=int, default=336)
+    parser.add_argument("--crop_size", type=int, default=None)
+    parser.add_argument("--num_frames", type=int, default=17)
+    parser.add_argument("--sample_rate", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--subset_size", type=int, default=None)
+    parser.add_argument("--tile_overlap_factor", type=float, default=0.25)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--output_origin', action='store_true')
+    parser.add_argument("--device", type=str, default="cuda")
+
+    args = parser.parse_args()
+    main(args)
+    

opensora/__init__.py

@@ -0,0 +1 @@
+# 

opensora/dataset/__init__.py

@@ -0,0 +1,99 @@
+from torchvision.transforms import Compose
+from transformers import AutoTokenizer
+
+from .feature_datasets import T2V_Feature_dataset, T2V_T5_Feature_dataset
+from torchvision import transforms
+from torchvision.transforms import Lambda
+
+from .landscope import Landscope
+from .t2v_datasets import T2V_dataset
+from .transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVideo, CenterCropResizeVideo
+from .ucf101 import UCF101
+from .sky_datasets import Sky
+
+ae_norm = {
+    'CausalVAEModel_4x8x8': Lambda(lambda x: 2. * x - 1.),
+    'CausalVQVAEModel_4x4x4': Lambda(lambda x: x - 0.5),
+    'CausalVQVAEModel_4x8x8': Lambda(lambda x: x - 0.5),
+    'VQVAEModel_4x4x4': Lambda(lambda x: x - 0.5),
+    'VQVAEModel_4x8x8': Lambda(lambda x: x - 0.5),
+    "bair_stride4x2x2": Lambda(lambda x: x - 0.5),
+    "ucf101_stride4x4x4": Lambda(lambda x: x - 0.5),
+    "kinetics_stride4x4x4": Lambda(lambda x: x - 0.5),
+    "kinetics_stride2x4x4": Lambda(lambda x: x - 0.5),
+    'stabilityai/sd-vae-ft-mse': transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+    'stabilityai/sd-vae-ft-ema': transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+    'vqgan_imagenet_f16_1024': Lambda(lambda x: 2. * x - 1.),
+    'vqgan_imagenet_f16_16384': Lambda(lambda x: 2. * x - 1.),
+    'vqgan_gumbel_f8': Lambda(lambda x: 2. * x - 1.),
+
+}
+ae_denorm = {
+    'CausalVAEModel_4x8x8': lambda x: (x + 1.) / 2.,
+    'CausalVQVAEModel_4x4x4': lambda x: x + 0.5,
+    'CausalVQVAEModel_4x8x8': lambda x: x + 0.5,
+    'VQVAEModel_4x4x4': lambda x: x + 0.5,
+    'VQVAEModel_4x8x8': lambda x: x + 0.5,
+    "bair_stride4x2x2": lambda x: x + 0.5,
+    "ucf101_stride4x4x4": lambda x: x + 0.5,
+    "kinetics_stride4x4x4": lambda x: x + 0.5,
+    "kinetics_stride2x4x4": lambda x: x + 0.5,
+    'stabilityai/sd-vae-ft-mse': lambda x: 0.5 * x + 0.5,
+    'stabilityai/sd-vae-ft-ema': lambda x: 0.5 * x + 0.5,
+    'vqgan_imagenet_f16_1024': lambda x: (x + 1.) / 2.,
+    'vqgan_imagenet_f16_16384': lambda x: (x + 1.) / 2.,
+    'vqgan_gumbel_f8': lambda x: (x + 1.) / 2.,
+}
+
+def getdataset(args):
+    temporal_sample = TemporalRandomCrop(args.num_frames * args.sample_rate)  # 16 x
+    norm_fun = ae_norm[args.ae]
+    if args.dataset == 'ucf101':
+        transform = Compose(
+            [
+                ToTensorVideo(),  # TCHW
+                CenterCropResizeVideo(size=args.max_image_size),
+                RandomHorizontalFlipVideo(p=0.5),
+                norm_fun,
+            ]
+        )
+        return UCF101(args, transform=transform, temporal_sample=temporal_sample)
+    if args.dataset == 'landscope':
+        transform = Compose(
+            [
+                ToTensorVideo(),  # TCHW
+                CenterCropResizeVideo(size=args.max_image_size),
+                RandomHorizontalFlipVideo(p=0.5),
+                norm_fun,
+            ]
+        )
+        return Landscope(args, transform=transform, temporal_sample=temporal_sample)
+    elif args.dataset == 'sky':
+        transform = transforms.Compose([
+            ToTensorVideo(),
+            CenterCropResizeVideo(args.max_image_size),
+            RandomHorizontalFlipVideo(p=0.5),
+            norm_fun
+        ])
+        return Sky(args, transform=transform, temporal_sample=temporal_sample)
+    elif args.dataset == 't2v':
+        transform = transforms.Compose([
+            ToTensorVideo(),
+            CenterCropResizeVideo(args.max_image_size),
+            RandomHorizontalFlipVideo(p=0.5),
+            norm_fun
+        ])
+        tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir='./cache_dir')
+        return T2V_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer)
+    elif args.dataset == 't2v_feature':
+        return T2V_Feature_dataset(args, temporal_sample)
+    elif args.dataset == 't2v_t5_feature':
+        transform = transforms.Compose([
+            ToTensorVideo(),
+            CenterCropResizeVideo(args.max_image_size),
+            RandomHorizontalFlipVideo(p=0.5),
+            norm_fun
+        ])
+        return T2V_T5_Feature_dataset(args, transform, temporal_sample)
+    else:
+        raise NotImplementedError(args.dataset)

opensora/dataset/extract_feature_dataset.py

@@ -0,0 +1,64 @@
+import os
+from glob import glob
+
+import numpy as np
+import torch
+import torchvision
+from PIL import Image
+from torch.utils.data import Dataset
+
+from opensora.utils.dataset_utils import DecordInit, is_image_file
+
+
+class ExtractVideo2Feature(Dataset):
+    def __init__(self, args, transform):
+        self.data_path = args.data_path
+        self.transform = transform
+        self.v_decoder = DecordInit()
+        self.samples = list(glob(f'{self.data_path}'))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        video_path = self.samples[idx]
+        video = self.decord_read(video_path)
+        video = self.transform(video)  # T C H W -> T C H W
+        return video, video_path
+
+    def tv_read(self, path):
+        vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+        total_frames = len(vframes)
+        frame_indice = list(range(total_frames))
+        video = vframes[frame_indice]
+        return video
+
+    def decord_read(self, path):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        frame_indice = list(range(total_frames))
+        video_data = decord_vr.get_batch(frame_indice).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+
+
+class ExtractImage2Feature(Dataset):
+    def __init__(self, args, transform):
+        self.data_path = args.data_path
+        self.transform = transform
+        self.data_all = list(glob(f'{self.data_path}'))
+
+    def __len__(self):
+        return len(self.data_all)
+
+    def __getitem__(self, index):
+        path = self.data_all[index]
+        video_frame = torch.as_tensor(np.array(Image.open(path), dtype=np.uint8, copy=True)).unsqueeze(0)
+        video_frame = video_frame.permute(0, 3, 1, 2)
+        video_frame = self.transform(video_frame)  # T C H W
+        # video_frame = video_frame.transpose(0, 1)  # T C H W -> C T H W
+
+        return video_frame, path
+

opensora/dataset/feature_datasets.py

@@ -0,0 +1,213 @@
+import json
+import os
+import torch
+import random
+import torch.utils.data as data
+
+import numpy as np
+from glob import glob
+from PIL import Image
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+from opensora.dataset.transform import center_crop, RandomCropVideo
+from opensora.utils.dataset_utils import DecordInit
+
+
+class T2V_Feature_dataset(Dataset):
+    def __init__(self, args, temporal_sample):
+
+        self.video_folder = args.video_folder
+        self.num_frames = args.video_length
+        self.temporal_sample = temporal_sample
+
+        print('Building dataset...')
+        if os.path.exists('samples_430k.json'):
+            with open('samples_430k.json', 'r') as f:
+                self.samples = json.load(f)
+        else:
+            self.samples = self._make_dataset()
+            with open('samples_430k.json', 'w') as f:
+                json.dump(self.samples, f, indent=2)
+
+        self.use_image_num = args.use_image_num
+        self.use_img_from_vid = args.use_img_from_vid
+        if self.use_image_num != 0 and not self.use_img_from_vid:
+            self.img_cap_list = self.get_img_cap_list()
+
+    def _make_dataset(self):
+        all_mp4 = list(glob(os.path.join(self.video_folder, '**', '*.mp4'), recursive=True))
+        # all_mp4 = all_mp4[:1000]
+        samples = []
+        for i in tqdm(all_mp4):
+            video_id = os.path.basename(i).split('.')[0]
+            ae = os.path.split(i)[0].replace('data_split_tt', 'lb_causalvideovae444_feature')
+            ae = os.path.join(ae, f'{video_id}_causalvideovae444.npy')
+            if not os.path.exists(ae):
+                continue
+
+            t5 = os.path.split(i)[0].replace('data_split_tt', 'lb_t5_feature')
+            cond_list = []
+            cond_llava = os.path.join(t5, f'{video_id}_t5_llava_fea.npy')
+            mask_llava = os.path.join(t5, f'{video_id}_t5_llava_mask.npy')
+            if os.path.exists(cond_llava) and os.path.exists(mask_llava):
+                llava = dict(cond=cond_llava, mask=mask_llava)
+                cond_list.append(llava)
+            cond_sharegpt4v = os.path.join(t5, f'{video_id}_t5_sharegpt4v_fea.npy')
+            mask_sharegpt4v = os.path.join(t5, f'{video_id}_t5_sharegpt4v_mask.npy')
+            if os.path.exists(cond_sharegpt4v) and os.path.exists(mask_sharegpt4v):
+                sharegpt4v = dict(cond=cond_sharegpt4v, mask=mask_sharegpt4v)
+                cond_list.append(sharegpt4v)
+            if len(cond_list) > 0:
+                sample = dict(ae=ae, t5=cond_list)
+                samples.append(sample)
+        return samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        # try:
+        sample = self.samples[idx]
+        ae, t5 = sample['ae'], sample['t5']
+        t5 = random.choice(t5)
+        video_origin = np.load(ae)[0]  # C T H W
+        _, total_frames, _, _ = video_origin.shape
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        assert end_frame_ind - start_frame_ind >= self.num_frames
+        select_video_idx = np.linspace(start_frame_ind, end_frame_ind - 1, num=self.num_frames, dtype=int)  # start, stop, num=50
+        # print('select_video_idx', total_frames, select_video_idx)
+        video = video_origin[:, select_video_idx]  # C num_frames H W
+        video = torch.from_numpy(video)
+
+        cond = torch.from_numpy(np.load(t5['cond']))[0]  # L
+        cond_mask = torch.from_numpy(np.load(t5['mask']))[0]  # L D
+
+        if self.use_image_num != 0 and self.use_img_from_vid:
+            select_image_idx = np.random.randint(0, total_frames, self.use_image_num)
+            # print('select_image_idx', total_frames, self.use_image_num, select_image_idx)
+            images = video_origin[:, select_image_idx]  # c, num_img, h, w
+            images = torch.from_numpy(images)
+            video = torch.cat([video, images], dim=1)  # c, num_frame+num_img, h, w
+            cond = torch.stack([cond] * (1+self.use_image_num))  # 1+self.use_image_num, l
+            cond_mask = torch.stack([cond_mask] * (1+self.use_image_num))  # 1+self.use_image_num, l
+        elif self.use_image_num != 0 and not self.use_img_from_vid:
+            images, captions = self.img_cap_list[idx]
+            raise NotImplementedError
+        else:
+            pass
+
+        return video, cond, cond_mask
+        # except Exception as e:
+        #     print(f'Error with {e}, {sample}')
+        #     return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+    def get_img_cap_list(self):
+        raise NotImplementedError
+
+
+
+
+class T2V_T5_Feature_dataset(Dataset):
+    def __init__(self, args, transform, temporal_sample):
+
+        self.video_folder = args.video_folder
+        self.num_frames = args.num_frames
+        self.transform = transform
+        self.temporal_sample = temporal_sample
+        self.v_decoder = DecordInit()
+
+        print('Building dataset...')
+        if os.path.exists('samples_430k.json'):
+            with open('samples_430k.json', 'r') as f:
+                self.samples = json.load(f)
+                self.samples = [dict(ae=i['ae'].replace('lb_causalvideovae444_feature', 'data_split_1024').replace('_causalvideovae444.npy', '.mp4'), t5=i['t5']) for i in self.samples]
+        else:
+            self.samples = self._make_dataset()
+            with open('samples_430k.json', 'w') as f:
+                json.dump(self.samples, f, indent=2)
+
+        self.use_image_num = args.use_image_num
+        self.use_img_from_vid = args.use_img_from_vid
+        if self.use_image_num != 0 and not self.use_img_from_vid:
+            self.img_cap_list = self.get_img_cap_list()
+
+    def _make_dataset(self):
+        all_mp4 = list(glob(os.path.join(self.video_folder, '**', '*.mp4'), recursive=True))
+        # all_mp4 = all_mp4[:1000]
+        samples = []
+        for i in tqdm(all_mp4):
+            video_id = os.path.basename(i).split('.')[0]
+            # ae = os.path.split(i)[0].replace('data_split', 'lb_causalvideovae444_feature')
+            # ae = os.path.join(ae, f'{video_id}_causalvideovae444.npy')
+            ae = i
+            if not os.path.exists(ae):
+                continue
+
+            t5 = os.path.split(i)[0].replace('data_split_1024', 'lb_t5_feature')
+            cond_list = []
+            cond_llava = os.path.join(t5, f'{video_id}_t5_llava_fea.npy')
+            mask_llava = os.path.join(t5, f'{video_id}_t5_llava_mask.npy')
+            if os.path.exists(cond_llava) and os.path.exists(mask_llava):
+                llava = dict(cond=cond_llava, mask=mask_llava)
+                cond_list.append(llava)
+            cond_sharegpt4v = os.path.join(t5, f'{video_id}_t5_sharegpt4v_fea.npy')
+            mask_sharegpt4v = os.path.join(t5, f'{video_id}_t5_sharegpt4v_mask.npy')
+            if os.path.exists(cond_sharegpt4v) and os.path.exists(mask_sharegpt4v):
+                sharegpt4v = dict(cond=cond_sharegpt4v, mask=mask_sharegpt4v)
+                cond_list.append(sharegpt4v)
+            if len(cond_list) > 0:
+                sample = dict(ae=ae, t5=cond_list)
+                samples.append(sample)
+        return samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        try:
+            sample = self.samples[idx]
+            ae, t5 = sample['ae'], sample['t5']
+            t5 = random.choice(t5)
+
+            video = self.decord_read(ae)
+            video = self.transform(video)  # T C H W -> T C H W
+            video = video.transpose(0, 1)  # T C H W -> C T H W
+            total_frames = video.shape[1]
+            cond = torch.from_numpy(np.load(t5['cond']))[0]  # L
+            cond_mask = torch.from_numpy(np.load(t5['mask']))[0]  # L D
+
+            if self.use_image_num != 0 and self.use_img_from_vid:
+                select_image_idx = np.random.randint(0, total_frames, self.use_image_num)
+                # print('select_image_idx', total_frames, self.use_image_num, select_image_idx)
+                images = video.numpy()[:, select_image_idx]  # c, num_img, h, w
+                images = torch.from_numpy(images)
+                video = torch.cat([video, images], dim=1)  # c, num_frame+num_img, h, w
+                cond = torch.stack([cond] * (1+self.use_image_num))  # 1+self.use_image_num, l
+                cond_mask = torch.stack([cond_mask] * (1+self.use_image_num))  # 1+self.use_image_num, l
+            elif self.use_image_num != 0 and not self.use_img_from_vid:
+                images, captions = self.img_cap_list[idx]
+                raise NotImplementedError
+            else:
+                pass
+
+            return video, cond, cond_mask
+        except Exception as e:
+            print(f'Error with {e}, {sample}')
+            return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+    def decord_read(self, path):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+        video_data = decord_vr.get_batch(frame_indice).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+    def get_img_cap_list(self):
+        raise NotImplementedError

opensora/dataset/landscope.py

@@ -0,0 +1,90 @@
+import math
+import os
+from glob import glob
+
+import decord
+import numpy as np
+import torch
+import torchvision
+from decord import VideoReader, cpu
+from torch.utils.data import Dataset
+from torchvision.transforms import Compose, Lambda, ToTensor
+from torchvision.transforms._transforms_video import NormalizeVideo, RandomCropVideo, RandomHorizontalFlipVideo
+from pytorchvideo.transforms import ApplyTransformToKey, ShortSideScale, UniformTemporalSubsample
+from torch.nn import functional as F
+import random
+
+from opensora.utils.dataset_utils import DecordInit
+
+
+class Landscope(Dataset):
+    def __init__(self, args, transform, temporal_sample):
+        self.data_path = args.data_path
+        self.num_frames = args.num_frames
+        self.transform = transform
+        self.temporal_sample = temporal_sample
+        self.v_decoder = DecordInit()
+
+        self.samples = self._make_dataset()
+        self.use_image_num = args.use_image_num
+        self.use_img_from_vid = args.use_img_from_vid
+        if self.use_image_num != 0 and not self.use_img_from_vid:
+            self.img_cap_list = self.get_img_cap_list()
+
+
+    def _make_dataset(self):
+        paths = list(glob(os.path.join(self.data_path, '**', '*.mp4'), recursive=True))
+
+        return paths
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        video_path = self.samples[idx]
+        try:
+            video = self.tv_read(video_path)
+            video = self.transform(video)  # T C H W -> T C H W
+            video = video.transpose(0, 1)  # T C H W -> C T H W
+            if self.use_image_num != 0 and self.use_img_from_vid:
+                select_image_idx = np.linspace(0, self.num_frames - 1, self.use_image_num, dtype=int)
+                assert self.num_frames >= self.use_image_num
+                images = video[:, select_image_idx]  # c, num_img, h, w
+                video = torch.cat([video, images], dim=1)  # c, num_frame+num_img, h, w
+            elif self.use_image_num != 0 and not self.use_img_from_vid:
+                images, captions = self.img_cap_list[idx]
+                raise NotImplementedError
+            else:
+                pass
+            return video, 1
+        except Exception as e:
+            print(f'Error with {e}, {video_path}')
+            return self.__getitem__(random.randint(0, self.__len__()-1))
+
+    def tv_read(self, path):
+        vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+        total_frames = len(vframes)
+
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+        video = vframes[frame_indice]  # (T, C, H, W)
+
+        return video
+
+    def decord_read(self, path):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+
+        video_data = decord_vr.get_batch(frame_indice).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+    def get_img_cap_list(self):
+        raise NotImplementedError

opensora/dataset/sky_datasets.py

@@ -0,0 +1,128 @@
+import os
+import torch
+import random
+import torch.utils.data as data
+
+import numpy as np
+
+from PIL import Image
+
+from opensora.utils.dataset_utils import is_image_file
+
+
+class Sky(data.Dataset):
+    def __init__(self, args, transform, temporal_sample=None, train=True):
+
+        self.args = args
+        self.data_path = args.data_path
+        self.transform = transform
+        self.temporal_sample = temporal_sample
+        self.num_frames = self.args.num_frames
+        self.sample_rate = self.args.sample_rate
+        self.data_all = self.load_video_frames(self.data_path)
+        self.use_image_num = args.use_image_num
+        self.use_img_from_vid = args.use_img_from_vid
+        if self.use_image_num != 0 and not self.use_img_from_vid:
+            self.img_cap_list = self.get_img_cap_list()
+
+    def __getitem__(self, index):
+
+        vframes = self.data_all[index]
+        total_frames = len(vframes)
+
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind-1, num=self.num_frames, dtype=int) # start, stop, num=50
+
+        select_video_frames = vframes[frame_indice[0]: frame_indice[-1]+1: self.sample_rate]
+
+        video_frames = []
+        for path in select_video_frames:
+            video_frame = torch.as_tensor(np.array(Image.open(path), dtype=np.uint8, copy=True)).unsqueeze(0)
+            video_frames.append(video_frame)
+        video_clip = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2)
+        video_clip = self.transform(video_clip)
+        video_clip = video_clip.transpose(0, 1)  # T C H W -> C T H W
+
+        if self.use_image_num != 0 and self.use_img_from_vid:
+            select_image_idx = np.linspace(0, self.num_frames - 1, self.use_image_num, dtype=int)
+            assert self.num_frames >= self.use_image_num
+            images = video_clip[:, select_image_idx]  # c, num_img, h, w
+            video_clip = torch.cat([video_clip, images], dim=1)  # c, num_frame+num_img, h, w
+        elif self.use_image_num != 0 and not self.use_img_from_vid:
+            images, captions = self.img_cap_list[index]
+            raise NotImplementedError
+        else:
+            pass
+
+        return video_clip, 1
+
+    def __len__(self):
+        return self.video_num
+    
+    def load_video_frames(self, dataroot):
+        data_all = []
+        frame_list = os.walk(dataroot)
+        for _, meta in enumerate(frame_list):
+            root = meta[0]
+            try:
+                frames = [i for i in meta[2] if is_image_file(i)]
+                frames = sorted(frames, key=lambda item: int(item.split('.')[0].split('_')[-1]))
+            except:
+                pass
+                # print(meta[0]) # root
+                # print(meta[2]) # files
+            frames = [os.path.join(root, item) for item in frames if is_image_file(item)]
+            if len(frames) > max(0, self.num_frames * self.sample_rate): # need all > (16 * frame-interval) videos
+            # if len(frames) >= max(0, self.target_video_len): # need all > 16 frames videos
+                data_all.append(frames)
+        self.video_num = len(data_all)
+        return data_all
+
+    def get_img_cap_list(self):
+        raise NotImplementedError
+
+if __name__ == '__main__':
+
+    import argparse
+    import torchvision
+    import video_transforms
+    import torch.utils.data as data
+
+    from torchvision import transforms
+    from torchvision.utils import save_image
+
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_frames", type=int, default=16)
+    parser.add_argument("--frame_interval", type=int, default=4)
+    parser.add_argument("--data-path", type=str, default="/path/to/datasets/sky_timelapse/sky_train/")
+    config = parser.parse_args()
+
+
+    target_video_len = config.num_frames
+
+    temporal_sample = video_transforms.TemporalRandomCrop(target_video_len * config.frame_interval)
+    trans = transforms.Compose([
+        video_transforms.ToTensorVideo(),
+        # video_transforms.CenterCropVideo(256),
+        video_transforms.CenterCropResizeVideo(256),
+        # video_transforms.RandomHorizontalFlipVideo(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+    ])
+
+    taichi_dataset = Sky(config, transform=trans, temporal_sample=temporal_sample)
+    print(len(taichi_dataset))
+    taichi_dataloader = data.DataLoader(dataset=taichi_dataset, batch_size=1, shuffle=False, num_workers=1)
+
+    for i, video_data in enumerate(taichi_dataloader):
+        print(video_data['video'].shape)
+        
+        # print(video_data.dtype)
+        # for i in range(target_video_len):
+        #     save_image(video_data[0][i], os.path.join('./test_data', '%04d.png' % i), normalize=True, value_range=(-1, 1))
+
+        # video_ = ((video_data[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1)
+        # torchvision.io.write_video('./test_data' + 'test.mp4', video_, fps=8)
+        # exit()

opensora/dataset/t2v_datasets.py

@@ -0,0 +1,111 @@
+import json
+import os, io, csv, math, random
+import numpy as np
+import torchvision
+from einops import rearrange
+from decord import VideoReader
+
+import torch
+import torchvision.transforms as transforms
+from torch.utils.data.dataset import Dataset
+from tqdm import tqdm
+
+from opensora.utils.dataset_utils import DecordInit
+from opensora.utils.utils import text_preprocessing
+
+
+
+class T2V_dataset(Dataset):
+    def __init__(self, args, transform, temporal_sample, tokenizer):
+
+        # with open(args.data_path, 'r') as csvfile:
+        #     self.samples = list(csv.DictReader(csvfile))
+        self.video_folder = args.video_folder
+        self.num_frames = args.num_frames
+        self.transform = transform
+        self.temporal_sample = temporal_sample
+        self.tokenizer = tokenizer
+        self.model_max_length = args.model_max_length
+        self.v_decoder = DecordInit()
+
+        with open(args.data_path, 'r') as f:
+            self.samples = json.load(f)
+        self.use_image_num = args.use_image_num
+        self.use_img_from_vid = args.use_img_from_vid
+        if self.use_image_num != 0 and not self.use_img_from_vid:
+            self.img_cap_list = self.get_img_cap_list()
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        try:
+            # video = torch.randn(3, 16, 128, 128)
+            # input_ids = torch.ones(1, 120).to(torch.long).squeeze(0)
+            # cond_mask = torch.cat([torch.ones(1, 60).to(torch.long), torch.ones(1, 60).to(torch.long)], dim=1).squeeze(0)
+            # return video, input_ids, cond_mask
+            video_path = self.samples[idx]['path']
+            video = self.decord_read(video_path)
+            video = self.transform(video)  # T C H W -> T C H W
+            video = video.transpose(0, 1)  # T C H W -> C T H W
+            text = self.samples[idx]['cap'][0]
+
+            text = text_preprocessing(text)
+            text_tokens_and_mask = self.tokenizer(
+                text,
+                max_length=self.model_max_length,
+                padding='max_length',
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors='pt'
+            )
+            input_ids = text_tokens_and_mask['input_ids'].squeeze(0)
+            cond_mask = text_tokens_and_mask['attention_mask'].squeeze(0)
+
+            if self.use_image_num != 0 and self.use_img_from_vid:
+                select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int)
+                assert self.num_frames >= self.use_image_num
+                images = video[:, select_image_idx]  # c, num_img, h, w
+                video = torch.cat([video, images], dim=1)  # c, num_frame+num_img, h, w
+                input_ids = torch.stack([input_ids] * (1+self.use_image_num))  # 1+self.use_image_num, l
+                cond_mask = torch.stack([cond_mask] * (1+self.use_image_num))  # 1+self.use_image_num, l
+            elif self.use_image_num != 0 and not self.use_img_from_vid:
+                images, captions = self.img_cap_list[idx]
+                raise NotImplementedError
+            else:
+                pass
+
+            return video, input_ids, cond_mask
+        except Exception as e:
+            print(f'Error with {e}, {self.samples[idx]}')
+            return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+    def tv_read(self, path):
+        vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+        total_frames = len(vframes)
+
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+
+        video = vframes[frame_indice]  # (T, C, H, W)
+
+        return video
+
+    def decord_read(self, path):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+
+        video_data = decord_vr.get_batch(frame_indice).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+    def get_img_cap_list(self):
+        raise NotImplementedError

opensora/dataset/transform.py

@@ -0,0 +1,489 @@
+import torch
+import random
+import numbers
+from torchvision.transforms import RandomCrop, RandomResizedCrop
+
+
+def _is_tensor_video_clip(clip):
+    if not torch.is_tensor(clip):
+        raise TypeError("clip should be Tensor. Got %s" % type(clip))
+
+    if not clip.ndimension() == 4:
+        raise ValueError("clip should be 4D. Got %dD" % clip.dim())
+
+    return True
+
+
+def center_crop_arr(pil_image, image_size):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+
+def crop(clip, i, j, h, w):
+    """
+    Args:
+        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+    """
+    if len(clip.size()) != 4:
+        raise ValueError("clip should be a 4D tensor")
+    return clip[..., i: i + h, j: j + w]
+
+
+def resize(clip, target_size, interpolation_mode):
+    if len(target_size) != 2:
+        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+    return torch.nn.functional.interpolate(clip, size=target_size, mode=interpolation_mode, align_corners=True, antialias=True)
+
+
+def resize_scale(clip, target_size, interpolation_mode):
+    if len(target_size) != 2:
+        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+    H, W = clip.size(-2), clip.size(-1)
+    scale_ = target_size[0] / min(H, W)
+    return torch.nn.functional.interpolate(clip, scale_factor=scale_, mode=interpolation_mode, align_corners=True, antialias=True)
+
+
+def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"):
+    """
+    Do spatial cropping and resizing to the video clip
+    Args:
+        clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        i (int): i in (i,j) i.e coordinates of the upper left corner.
+        j (int): j in (i,j) i.e coordinates of the upper left corner.
+        h (int): Height of the cropped region.
+        w (int): Width of the cropped region.
+        size (tuple(int, int)): height and width of resized clip
+    Returns:
+        clip (torch.tensor): Resized and cropped clip. Size is (T, C, H, W)
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    clip = crop(clip, i, j, h, w)
+    clip = resize(clip, size, interpolation_mode)
+    return clip
+
+
+def center_crop(clip, crop_size):
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    h, w = clip.size(-2), clip.size(-1)
+    th, tw = crop_size
+    if h < th or w < tw:
+        raise ValueError("height and width must be no smaller than crop_size")
+
+    i = int(round((h - th) / 2.0))
+    j = int(round((w - tw) / 2.0))
+    return crop(clip, i, j, th, tw)
+
+
+def center_crop_using_short_edge(clip):
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    h, w = clip.size(-2), clip.size(-1)
+    if h < w:
+        th, tw = h, h
+        i = 0
+        j = int(round((w - tw) / 2.0))
+    else:
+        th, tw = w, w
+        i = int(round((h - th) / 2.0))
+        j = 0
+    return crop(clip, i, j, th, tw)
+
+
+def random_shift_crop(clip):
+    '''
+    Slide along the long edge, with the short edge as crop size
+    '''
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    h, w = clip.size(-2), clip.size(-1)
+
+    if h <= w:
+        long_edge = w
+        short_edge = h
+    else:
+        long_edge = h
+        short_edge = w
+
+    th, tw = short_edge, short_edge
+
+    i = torch.randint(0, h - th + 1, size=(1,)).item()
+    j = torch.randint(0, w - tw + 1, size=(1,)).item()
+    return crop(clip, i, j, th, tw)
+
+
+def to_tensor(clip):
+    """
+    Convert tensor data type from uint8 to float, divide value by 255.0 and
+    permute the dimensions of clip tensor
+    Args:
+        clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
+    Return:
+        clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W)
+    """
+    _is_tensor_video_clip(clip)
+    if not clip.dtype == torch.uint8:
+        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
+    # return clip.float().permute(3, 0, 1, 2) / 255.0
+    return clip.float() / 255.0
+
+
+def normalize(clip, mean, std, inplace=False):
+    """
+    Args:
+        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
+        mean (tuple): pixel RGB mean. Size is (3)
+        std (tuple): pixel standard deviation. Size is (3)
+    Returns:
+        normalized clip (torch.tensor): Size is (T, C, H, W)
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    if not inplace:
+        clip = clip.clone()
+    mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device)
+    # print(mean)
+    std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device)
+    clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None])
+    return clip
+
+
+def hflip(clip):
+    """
+    Args:
+        clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W)
+    Returns:
+        flipped clip (torch.tensor): Size is (T, C, H, W)
+    """
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    return clip.flip(-1)
+
+
+class RandomCropVideo:
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: randomly cropped video clip.
+                size is (T, C, OH, OW)
+        """
+        i, j, h, w = self.get_params(clip)
+        return crop(clip, i, j, h, w)
+
+    def get_params(self, clip):
+        h, w = clip.shape[-2:]
+        th, tw = self.size
+
+        if h < th or w < tw:
+            raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
+
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = torch.randint(0, h - th + 1, size=(1,)).item()
+        j = torch.randint(0, w - tw + 1, size=(1,)).item()
+
+        return i, j, th, tw
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+
+class CenterCropResizeVideo:
+    '''
+    First use the short side for cropping length,
+    center crop video, then resize to the specified size
+    '''
+
+    def __init__(
+            self,
+            size,
+            interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: scale resized / center cropped video clip.
+                size is (T, C, crop_size, crop_size)
+        """
+        clip_center_crop = center_crop_using_short_edge(clip)
+        clip_center_crop_resize = resize(clip_center_crop, target_size=self.size,
+                                         interpolation_mode=self.interpolation_mode)
+        return clip_center_crop_resize
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class UCFCenterCropVideo:
+    '''
+    First scale to the specified size in equal proportion to the short edge,
+    then center cropping
+    '''
+
+    def __init__(
+            self,
+            size,
+            interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: scale resized / center cropped video clip.
+                size is (T, C, crop_size, crop_size)
+        """
+        clip_resize = resize_scale(clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode)
+        clip_center_crop = center_crop(clip_resize, self.size)
+        return clip_center_crop
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class KineticsRandomCropResizeVideo:
+    '''
+    Slide along the long edge, with the short edge as crop size. And resie to the desired size.
+    '''
+
+    def __init__(
+            self,
+            size,
+            interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        clip_random_crop = random_shift_crop(clip)
+        clip_resize = resize(clip_random_crop, self.size, self.interpolation_mode)
+        return clip_resize
+
+
+class CenterCropVideo:
+    def __init__(
+            self,
+            size,
+            interpolation_mode="bilinear",
+    ):
+        if isinstance(size, tuple):
+            if len(size) != 2:
+                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            self.size = size
+        else:
+            self.size = (size, size)
+
+        self.interpolation_mode = interpolation_mode
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: center cropped video clip.
+                size is (T, C, crop_size, crop_size)
+        """
+        clip_center_crop = center_crop(clip, self.size)
+        return clip_center_crop
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+
+class NormalizeVideo:
+    """
+    Normalize the video clip by mean subtraction and division by standard deviation
+    Args:
+        mean (3-tuple): pixel RGB mean
+        std (3-tuple): pixel RGB standard deviation
+        inplace (boolean): whether do in-place normalization
+    """
+
+    def __init__(self, mean, std, inplace=False):
+        self.mean = mean
+        self.std = std
+        self.inplace = inplace
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): video clip must be normalized. Size is (C, T, H, W)
+        """
+        return normalize(clip, self.mean, self.std, self.inplace)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})"
+
+
+class ToTensorVideo:
+    """
+    Convert tensor data type from uint8 to float, divide value by 255.0 and
+    permute the dimensions of clip tensor
+    """
+
+    def __init__(self):
+        pass
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W)
+        Return:
+            clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W)
+        """
+        return to_tensor(clip)
+
+    def __repr__(self) -> str:
+        return self.__class__.__name__
+
+
+class RandomHorizontalFlipVideo:
+    """
+    Flip the video clip along the horizontal direction with a given probability
+    Args:
+        p (float): probability of the clip being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Size is (T, C, H, W)
+        Return:
+            clip (torch.tensor): Size is (T, C, H, W)
+        """
+        if random.random() < self.p:
+            clip = hflip(clip)
+        return clip
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(p={self.p})"
+
+
+#  ------------------------------------------------------------
+#  ---------------------  Sampling  ---------------------------
+#  ------------------------------------------------------------
+class TemporalRandomCrop(object):
+    """Temporally crop the given frame indices at a random location.
+
+    Args:
+        size (int): Desired length of frames will be seen in the model.
+    """
+
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, total_frames):
+        rand_end = max(0, total_frames - self.size - 1)
+        begin_index = random.randint(0, rand_end)
+        end_index = min(begin_index + self.size, total_frames)
+        return begin_index, end_index
+
+
+if __name__ == '__main__':
+    from torchvision import transforms
+    import torchvision.io as io
+    import numpy as np
+    from torchvision.utils import save_image
+    import os
+
+    vframes, aframes, info = io.read_video(
+        filename='./v_Archery_g01_c03.avi',
+        pts_unit='sec',
+        output_format='TCHW'
+    )
+
+    trans = transforms.Compose([
+        ToTensorVideo(),
+        RandomHorizontalFlipVideo(),
+        UCFCenterCropVideo(512),
+        # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+    ])
+
+    target_video_len = 32
+    frame_interval = 1
+    total_frames = len(vframes)
+    print(total_frames)
+
+    temporal_sample = TemporalRandomCrop(target_video_len * frame_interval)
+
+    # Sampling video frames
+    start_frame_ind, end_frame_ind = temporal_sample(total_frames)
+    # print(start_frame_ind)
+    # print(end_frame_ind)
+    assert end_frame_ind - start_frame_ind >= target_video_len
+    frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int)
+    print(frame_indice)
+
+    select_vframes = vframes[frame_indice]
+    print(select_vframes.shape)
+    print(select_vframes.dtype)
+
+    select_vframes_trans = trans(select_vframes)
+    print(select_vframes_trans.shape)
+    print(select_vframes_trans.dtype)
+
+    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(dtype=torch.uint8)
+    print(select_vframes_trans_int.dtype)
+    print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
+
+    io.write_video('./test.avi', select_vframes_trans_int.permute(0, 2, 3, 1), fps=8)
+
+    for i in range(target_video_len):
+        save_image(select_vframes_trans[i], os.path.join('./test000', '%04d.png' % i), normalize=True,
+                   value_range=(-1, 1))