From d1ea86d351d08969e85d82da7e04f88044867ec4 Mon Sep 17 00:00:00 2001
From: lb203 <62638829+LinB203@users.noreply.github.com>
Date: Mon, 4 Mar 2024 14:06:39 +0000
Subject: [PATCH] Initial Commit based on Open-Sora-Plan-V1.2
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Former-commit-id: 7a3dccf2c738eaf535e62996211c8bc5f3de5683

Update README.md

Former-commit-id: aee63a3a8422e029f1550e47825ca4beb916874b

Update README.md

Former-commit-id: fb46c80c1a2663c09849472e93b2da6cd068a98c

Update README.md

Former-commit-id: 5fd1e6afe1277ccffdfeaf507477d944a6a35f03

Update README.md

Former-commit-id: 1cd7d8b5df34ab9f61d9346b2647b643f78a8f5d

Update README.md

Former-commit-id: e27cc5d49a19e9580ba1a29d948051b918aede74

Update README.md

Former-commit-id: 0ff9314bddef1dfbbf429726ed4240b823b221b3

Update README.md

Former-commit-id: f295f4623448dd4ffa9be747a5cc6733cf48c663

Create Contribution_Guidelines.md

Former-commit-id: 897dbd63dc6f53095e1dda3f10a86dd40e4bd263

Update README.md

Former-commit-id: c652acd1fc5a003b4d6eefcfb83dfb310f5b5e81

Rename data.md to Data.md

Former-commit-id: 27484f7a2d8634bd1cd9781a2f10afa7f304f5cb

Update README.md

Former-commit-id: 351be764d8648ac60cf170c3fe4e6ed9c544c90e

reformat code

Former-commit-id: ce23b280fbb419c08347ebcbf1ed4075bdf9af60

Update requirements.txt

As in previous step,the three packages have been installed, these packages should be deleted here because it can not be installed by " pip install -r requirements.txt" command and will abort the command

Former-commit-id: 0655fa28cf8bb761f1a3d2def62f44197b257b2c

Update README.md

Former-commit-id: 1dd9c945be2a170d87d568bb00e1ec30e7be592b

Update README.md

Former-commit-id: 6ecaba8e03d7c7f99c068a57c923f11c5456a8c9

Update Contribution_Guidelines.md

Former-commit-id: dc87e459ca5a8a090800331111118acb65214f65

Update README.md

Former-commit-id: be7ecc3d56f7289052fc66d535038e7111f7ac84

Update README.md

Former-commit-id: 164e76cf496b1157eebea299ed58ca6c1799d70d

Update README.md

Former-commit-id: 3e3d674134498512b1ef1b0dd72d6e19afd6ef65

[feat]: frame_interpolation

Former-commit-id: 711fc4bba80d5f051f712e3460e2c0685117cf68

Update README.md

Former-commit-id: 1436d8feea7f82d3d747c632a68c47b3ea74258f

fix_readme_requirements

Former-commit-id: eb5a5c69d12cc412c12a4fea7f00baafd02701ca

Update README.md

Former-commit-id: c5a79c7e993d9e6a77c6c70efa1c0ac2bb3a063b

add latte

Former-commit-id: 01ba10e5d65c9d2d957472fbe7958e1b07a971e9

Update README.md

Former-commit-id: 437bc5c68b5c7a85865ad97d544c4818be45cb52

first commit

Former-commit-id: ec3481bc35a8120d519082dc76585428f0896470

[refactor] reformat videogpt, support training videogpt on accelerate

Former-commit-id: 75385ed47fe60008b61b1bcf89f61ea696d095bc

[refactor] adapt to old training code

Former-commit-id: 358d362ecbc100b21ea02895cb0426cf5b386a72

Delete LICENSE.txt

Former-commit-id: b10d8b2c6620e3a8412527bd03ea0adc26c981a6

Create LICENSE

Former-commit-id: 369fe085e91e55b9be6165da4a1eb7f9c0cf5a7f

Update README.md

Former-commit-id: f1aaeca19bc9e4304961c8c903f9bf03cb42fd9f

Update README.md

Former-commit-id: 0f2bdd2bcd819d528f2a41a96ba93f5c3c04dc8c

Update README.md

Former-commit-id: 10ee64c122f28efea39fc1a4542d278683a58ce1

option to use rebased linear attention

Former-commit-id: f1b39b34df5fbdef14ef3f6a0e6d0972a8570445

support for ring attention

Former-commit-id: 0533c73b99217e88da121bd0f68e0949bc88016e

add vae and fix bug in dit

Former-commit-id: 2dd4254c22ce671d5cd5e7d22a2386b7660fbdc4

Update README.md

Former-commit-id: 898d35f15ce0fd124cd48078f0a8af2764bf55d1

support latte training

Former-commit-id: d471a9414842f422d1200c2fd3f3ddfcf9e97fd6

Update README.md

Former-commit-id: 9fd88eb8ac5122b0999ccb0071f195aaeb0979a0

Update README.md

Former-commit-id: 398cfcd7b19cf1ba827eb49fc11223f72fa93a85

Update README.md

Former-commit-id: 673659495b7503ba402fbcb6c46204499761b4cf

[fix] fix a bug when using ucf101_stride4x4x4

Former-commit-id: 9911403c688321db731450d9383bc02923a4be2d

[refactor] remove test.ipynb

Former-commit-id: e73c3a9e0d41e51a0fd755e57739424213a76312

[fix] fix a bug when using ucf101_stride4x4x4

Former-commit-id: f541e723d65ea44ae4560c68908123c7e6781bbc

[refactor] rename

Former-commit-id: ec317605d24223527b8c2d4b4d2ee876f86abbaf

update

add RGT

Former-commit-id: a491111e48cec588095381eb70ffca94600ca4f7

add sample script

Former-commit-id: 9ce847651bb6dba28c51ce5807e222ae2902a78a

[refactor] configuration default values

Former-commit-id: d7ca7685d6a8c6d015082dda241bc98ab156770d

[refactor] added more methods in abstract classes

Former-commit-id: a535e3a06eb766bebe7baa3d9df64b04d04d40e6

Update README.md

Former-commit-id: 59a595a1e31b156500a944f244afec40988e872d

[fix]: disable gradient computation to save GPU memory when reconstructing a video

Former-commit-id: b7bf763e7f8800959ed82bb89f41684db216dd8e

support accelerate training

Former-commit-id: b32671c9991b9a9db66cb148822ccb3e8be7dc11

Update train.sh

Former-commit-id: 6c4957d715eaa489691be21dd0abdb9309660e0d

Update train.sh

Former-commit-id: 66aa3709f80d89a30d0a57c5e033ce1ae46f998f

Update README.md

Former-commit-id: 604a29d1de43e0fcb90ada4dd97f5f080e37caea

Update README.md

Former-commit-id: d097e40998a3396f74c2bfbcf26d6ace8080d0b6

Update README.md

Former-commit-id: 1a7178a61547c33a744d032302fb76ee0a3c6671

denorm fun

Former-commit-id: 05707a07b14295f5e1932d0c10bb84a20a0cd328

Update README.md

Former-commit-id: b3f6bb7f252505efb0ffd8195003c79206fac313

refine readme.

Former-commit-id: 391d9511b72aa9c09c7ad8992ae5df8c89c88c63

complement emoji.

Former-commit-id: 50679cc482d8ecf4e3083c0f5a791f195a58a855

Update README.md

Former-commit-id: 29f24a819fe2d4593c48afe6ca990659784bd776

[feat]: add sit model

Former-commit-id: 2eda854d6c25254b71d0105a3ed092953a9d6f68

fix: remove replicated diffusion module in sit

Former-commit-id: 997cb049185421d303b14d59ab6c0603f1d1a66f

fix: rename sit model filename

Former-commit-id: f3b2cd7dcf644f680616eee0571db90478fede23

feat: Incorporating SiT sample

Former-commit-id: 91e2931e412f8c96ef748cb2322ee956309fb6aa

feat: update sample

Former-commit-id: e7dc492b9b4764db28aa31cc1eede914052b4ca6

attn_mask with bf16

Former-commit-id: d9667a8c15a5344e3c777197fd6badec15dcddbd

Update README.md

Former-commit-id: a1cc0b26a3b156d7708e8df0cfe42aad33b1c946

Update README.md

Former-commit-id: 4179ed075e68d48094655c38dbde784836fd6e0f

Update README.md

Former-commit-id: 36e69275598b0cff38fdc5006926e422cda069b3

[fix]: disable gradient computation to save GPU memory when reconstructing a video

Former-commit-id: 06da17790b53cf593a0c264dda17c76f68d2a541

Update README.md

Former-commit-id: 8e9658ed993380abc43b518fcbac21a02d8ba9ba

Update README.md

Former-commit-id: 8b553a4477c64b5452f8424ca6987c0fb8c111ce

Update README.md

Former-commit-id: 0898e4b917f5e245c4394118511bcb7639973817

Update README.md

Former-commit-id: 4a893ff6177d59219b6755c4a89764ef54fc7f62

Update README.md

Former-commit-id: 89f9fc8eed3ca8cc2434ec227acbb845bfbcff82

[refactor] rename train_videogpt.sh

Former-commit-id: e2ea2caacaf5be046f48cf50a53ccafa7a3cc351

[docs] Modify README.md and add VQVAE documentation.

Former-commit-id: a5bfe70f8a8998db969d9c4938fda123aa153bc0

fixed dit forward

Former-commit-id: 4cd537fab12231f7471107762b02b36217ca42d8

videogpt inference

Former-commit-id: c5d617450a8271ad3b8a5e8f76ee2d77a822e31b

Update README.md

Former-commit-id: 4d83f1feab7d6872aca1c6b41be95950ec3cf465

Update README.md

Former-commit-id: da5507e570a27becab754ba24a0671d669b7e08e

Update README.md

Former-commit-id: 4a04e71636dabccbef446623674f3a514395b7f9

Update README.md

Former-commit-id: 2a286e5aab48aa68e7d3f04df23ec1cdfe71c6de

Update README.md

Former-commit-id: f4c4b3b4b1ce3c1d0b56ba5a109cd6cc57da900d

Update README.md

Former-commit-id: 7af6b82b29c265e53e9c48477083b4cff0e0a3d3

Update README.md

Former-commit-id: 7bc544af55c0a987d0df426848cb06cb330fd52b

Update README.md

Former-commit-id: 2a68a1aa399a53f58246df4f9f64dfdcdd130bcf

[docs] Modify README.md to fix translation error in name "CloseGPT"

Former-commit-id: dd36043df4b0408e592a82f9046f6d654a5375c3

[docs] Modify README.md to fix other grammatical errors

Former-commit-id: f0de27fcc38bb4dab0eeafac225a2924080a3218

Update README.md

Former-commit-id: a881ff5198e1125032418f4f02fd95d2179ffb3a

Update README.md

Former-commit-id: 8a191abc1903ad5998fe83149c9c9ef5c8226b9e

Update README.md

Former-commit-id: 3b4eb14dcd8862b8251b72f30cba523a5dc1def8

Update README.md

Former-commit-id: e458366e47d4b13737a71bbf866e20ec10961ade

init pos embed

Former-commit-id: 28f8076136ec214c66039ead7547c9b25cfe7bfc

Update README.md

Former-commit-id: a1d1f07c83eb0c03d93b023b5c9162a6a3c26cd9

Update README.md

Former-commit-id: deddba1f03bf697fb1252a0ce4fdb5df5bcae8de

Update README.md

Former-commit-id: 076e6e70e717888f24265d624c6a7a8ffa2f5d1b

Update README.md

Former-commit-id: 7cd398721e556db6199f1d06c7a0d1de2a9720dd

[updata] updata video super resolution

updata video super resolution

Former-commit-id: 82397545859791d112d6f0569e33fe23437563df

train script

Former-commit-id: 3fd1ccf099c2deed20af50bf9556e2fc3b7a4382

Delete opensora/models/super_resolution/placeholder

Former-commit-id: 8dd5a6a3c85dc28989ff95e6ad41be7fdb0f0811

Update Data.md

Former-commit-id: 9531e33666f842125b331cb266eb26a3a03a4009

Update Data.md

Former-commit-id: 86be6b2bf5a9be1030bb3af661a2ec85e710a5c3

commit

Former-commit-id: 81359cf67248b6c594ff79f121a8a38558c1cc70

Update README.md

Former-commit-id: c4d7aafa3ceb6e4d48bf85b402ae2b2ce9394e53

Update README.md

Former-commit-id: bca7070fac79d5ae5c37830cd55ca1e81a289197

Update README.md

Former-commit-id: b37da2202e730a4930a194a80a703b591a535adf

Update __init__.py

Former-commit-id: 03a45aaf60d294b7a68d8fc44132d3d671e4472b

Update README.md

Former-commit-id: 97fe914aaa2b9bd7c2ff5d4239381631c47e8356

Update README.md

Former-commit-id: c6b1ea8f1ef01096aa8c3e5272bc9fb255b5532b

Delete scripts/train_vqvae.sh

Former-commit-id: 952a4ce3a3b0da159d599ff3a0a052265f48c2dd

update script

Former-commit-id: 43f539b68af150c45040d79918a84de1dce47c33

feature dataset

Former-commit-id: 9683caae8648735552270b871de22fc540ba2b55

update train.py

Former-commit-id: 844173711342f2dcc021294e87971ca39854bfe4

del use_fp16

Former-commit-id: 5424c0112e04b56d1df30ac76cbc104e848406a3

update sample.py

Former-commit-id: 10d8621051d1510270c6a55abe686eac86297170

fix extract feature

Former-commit-id: d4dce5ec1aaa058e452a41f15cd35fe1713f0b78

Update README.md

Former-commit-id: 70f6ebf421f0830a51dcc2838c2c3783cfb67376

[Updata] updata the video super resolution

updata the run.py and add README.md

Former-commit-id: 1e6d6ace7a0963149e228813801210b022aceacf

add CI for docker autobuild

Former-commit-id: c4e2d5ff1dce8662a42f1ed6a7786855d40173f7

[refactor]: add docker development support

Former-commit-id: fded3baff7915c0cfc9cec3765c271259359831c

update docker scripts and configs

Former-commit-id: 4567e6368f1bf413d8f6319a97a852b8289ca770

Update README.md

Former-commit-id: 99b0066895af072a8887737f2b603ec7fc893ccd

extract script

Former-commit-id: 4139f2f68ae09695302ce49afdb7318df83cabfd

[docs]:update the eval_code for calculate the FVD, clip_score, ssim, lpips, psnr

Former-commit-id: 1ee226c338126fbdc5178020b12271b980a7ce0c

[docs]:update the eval_code for calculate the FVD, clip_score, ssim, lpips, psnr

Former-commit-id: b504623d7ec2ac0b508530d21ffc88978f8a4f71

[docs]: update the evaluation

Former-commit-id: 02e3789b74daa036fbbff6d13dfc7e98372260ed

Update latte.py

revise the xformers inputs

Former-commit-id: 60db66c912192c3fbafe55ebe16157514b884e5b

Fix the memory error bug when loading metadata pickle file of vqvae dataset.

Former-commit-id: 97ab243bd12915cea7aa2deecfac4b99417f78a3

Update README.md

Former-commit-id: 8f7b41118a3f491d1d5b5789c53790c46b023191

Update pyproject.toml

Former-commit-id: 36da5d0e950625e629aebe954f89f92827397d0d

Update README.md

Former-commit-id: 826cc5e9162888bed82413088c6be662c278176b

Update README.md

Former-commit-id: 88e9f21cfc8ec775dbc9a65f740043e91e26bd91

Update README.md

Former-commit-id: 38e0e290e7a5909706c2cce2e4dd047647134fb8

[docs]:update the eval code

Former-commit-id: 937f2fc87d34ab9d4724cb95a6457b54a6bf4ebc

[docs]:update the eval code

Former-commit-id: a8c9e6cb9e39a9e929dab34077280a8d80d85573

flash-attn train and sample

Former-commit-id: e336565bd5ff2aaf4d0b4ad5f42732e89cd216b3

[refactor] add static type checking

Former-commit-id: 12c4240dc28888c61f1f79243627e7cdb445c840

Update README.md

Former-commit-id: 513883df1ad853c9a5e7457c364227261df17a83

Fix einops bug: module 'keras.backend' has no attribute 'is_tensor'

Former-commit-id: 849f8ed815295bac78c3c10d319e6731aa26ca02

[feat] Add 2D RoPE

Former-commit-id: ccae80d339e34fb23706827565b5095ae2c2f320

Update README.md

Former-commit-id: 6630c22baa2340a8b55badd61472daab7177bc63

xformers and flashattn

Former-commit-id: 002cad95253f7da36344237ce8c8410ffb517711

safe dit

Former-commit-id: f407cab7dc0e53257972d430e873f4fd77bb039d

[fix] 2d RoPE init

Former-commit-id: ff5931bae6efa94675b9ef6bc25f1b76217d128e

fixed attn_mask in flash-attn

Former-commit-id: 5e2504f429a342b75bc7a9004cee7b3f75582ee3

typo

Former-commit-id: 9e87366e906ed19726c8905afb27b5df4d4e1512

Decouple mixed_precision and gradient_accumulation to command args.

Former-commit-id: 903d473d1432eb492a87efe184d03439f643b92d

Fix mixed_precision bugs of latte.

Former-commit-id: d70a551972e04d95444d46cdbe2bfb16ad099fcb

Refactor latte.

Former-commit-id: 304534269f105578097bc16ab497636ad59fd9b2

Refactor latte.

Former-commit-id: 7408fd0ee16247afc3d4ba43b8c8e7dfb1091a97

Support deepspeed zero2 and zero2_offload for latte training.

Former-commit-id: c87e95dd357e6fa6d6fb2dcfee781bc49c039e13

Add some comments.

Former-commit-id: d5b5bd7743edde6a63dc74457bded0f8c56be1c2

Fix typos.

Former-commit-id: c10330d373466f3c1ab455bfb391569b9175bc3a

deepspeed

Former-commit-id: 6179a4b0e7d06caed4201a27679abc64615febe6

Update README.md

Former-commit-id: 582070d4f97c4a4398dcb9a6699f4ab6e980c812

Add base config.

Former-commit-id: d37c551a0421d6b73dee26060b517c46d8217e4f

train 1080p video

Former-commit-id: 74625142ef892f9e7fa459b2f98f2f7ae5b36d6a

Update README.md

Former-commit-id: 87636c1dc2275c4da77c6368a0f4425bea15dc12

Update README.md

Former-commit-id: 059d9e863958cfd679dbb8d365de6cd933e284fc

Update pyproject.toml

Former-commit-id: 5c379b6e4515ad75b89cdebbc5cd490c9a821d78

Update README.md

Former-commit-id: 7a62dec060c91b02ca556451266e6f3ef28e03e7

clean script

Former-commit-id: 02c61f63a8539a31516b92f5bb7c2d2e428ca7d7

Update README.md

Former-commit-id: 6d41cfe906f97dbed1cea8c9ed06a1f6dc84308e

Update README.md

Former-commit-id: 64d147863368aeef479ad25f2e5463fdcb76de46

Update README.md

Former-commit-id: 2808edf11005e04aa730971a786d1a4e3bca88b5

Update README.md

Former-commit-id: 1f94f2721f35670aeae8dbfb7c0a5abd06cd2aee

[feat]: Caption Refiner

Caption Refiner for Video Caption

Former-commit-id: 97c16bccb9979acf19fb91ae8d7b84ac4a3b4d97

[BUG]: Shorten the refined caption

Shorten the refined caption using GPT-3.5 summary

Former-commit-id: bb159f0d5fa928258475b2058983397ceb91af9d

Fix resolution typos.

Former-commit-id: 3ad55d28eaf13f9c8fd23912d2691965737527e9

clean dit

Former-commit-id: 117ec4196a6f2cd83da223b993b14e28f7a693f8

Update README.md

Former-commit-id: abb8ed57a29b1f66adec66d854429f36f9199540

Update README.md

Former-commit-id: c2c64665d4f281e264d87c20cc1cfd38e7e5f168

[fix]: correct attention_mode to attention-mode

Former-commit-id: 9957e2eebb2166b02600c6052fb54908f5d12bad

Update README.md

Former-commit-id: 266f9f97ea8f327fdefb2ca959c23d08b7084137

Update README.md

Former-commit-id: dd00e04220ff5232c5c4503a095e409471476a51

Update README.md

Former-commit-id: 5717aab628fceba311e284d1bdddeacc4c43c7c5

Create train_256.sh

Former-commit-id: bf62151617feb49663f8dfe0c78aa01f41e62092

Update README.md

Former-commit-id: c5eea8d3492872255a6a34d422230509d1c3531c

[feat]: use accelerate on multi-node

Former-commit-id: 0b12fc31f1d2a63f9da38821174ab63d9b3a2d72

Update README.md

Former-commit-id: 0d7dd6233811ec3c8741affb436d8004d8c676b8

Update sample.py

Former-commit-id: 6e09fc266a626f465109abc199d9fb452657929e

Update interpolation.py

Former-commit-id: e10e3416433e5f396aa13e7d8beaf451ff47a383

Add files via upload

Former-commit-id: b707b1968fa59d85891040604e7d54bef61b4549

Update and rename Frame Interpolation.md to readme.md

Former-commit-id: 87bdf1ebecb002c086497b3c47acb3f3cb2aaf34

Update readme.md

Former-commit-id: e287f94838007f2eda434a46bcd250099b627097

Support deepspeed for videogpt training.

Former-commit-id: 1f3dfe038845d91b81b0ecc9b6f3cabb06eed5d5

Fix config typos.

Former-commit-id: 455973b57bda19379d41e59aa418421fd269281e

Add hidden_size argument for videogpt config.

Former-commit-id: 1be9f22767fe9c80ad90cb17184751b97aff9012

Add deepspeed script for videogpt training.

Former-commit-id: 3bda2a05cb789fe785b2355a613385b3eb7abbbc

use fp16 temporarily.

Former-commit-id: 93381968e10aae7bbec23ec1ef40e7cb311fddf0

update videogpt.

Former-commit-id: 2e211563c684a8b6ef62dfc0c4f1897224dfb58d

Fix deepspeed zero loss bug.

Former-commit-id: e24af415bc5d6fd38d72a24d6312f4e65d3627dd

Can't fix zero loss bug.

Former-commit-id: 1f32f9d889743b7961b34b823b276bfd03974191

use imageio for write video which has better compatibility.

Former-commit-id: 005119dc9f78ea2a611d22a53772fd59d97d7139

Update training arguments of videogpt training.

Former-commit-id: c9660af891c0dae1f51a1075ef5b2df012f65e06

Update README.md

Former-commit-id: 8ec7524f4f171e650d7a25cc88ed29f7b5017454

Update train_256.sh

Former-commit-id: a2dfc9a8c70845446e731c88acd222e9b2db9122

Update README.md

Former-commit-id: 7e74356a3279b02d062bab9c9f361f61eb756179

Update README.md

Former-commit-id: cf1c56c0058288427c1dfd651c8c0e9a981c5653

refactor dit for supporting deepspeed.

Former-commit-id: e1909745cb05ea2f28f3a391c194f7656593a236

Fix typos.

Former-commit-id: 6add5ee7ef39a973e91b7a7c1b3a4b0980130b88

Update README.md

Former-commit-id: 2d5b6815f95e94a98847b8a0376473f01bb15d14

[bug] fix quantization loss

Former-commit-id: 111fcf55f969ffe4c9be284a317ae5712570a527

del t2v

Former-commit-id: 910398ca297be8a188d59c59c121857666774af1

Update sample.py

Former-commit-id: eb906691abec030d7aa2078b63fc9c7d678662ed

Update sample.sh

Former-commit-id: b149a300f1aff53d9ab6e66a91cc2bcd37900347

Update README.md

Former-commit-id: 7b96fd89eaf3743a4a9fbc995e6a2a791c734f2e

del t2v

Former-commit-id: 5daa2952997a29e74e2b5d235001d6d9dabf9ff6

[fix]: fix sample.sh

Former-commit-id: 09916dfcfb0b80a4217016919e0394a2cb665cfd

update scripts

Former-commit-id: 8b2a794095f2deb9b1e8090486cb8242bf81031d

Update pyproject.toml

Former-commit-id: 7fb921440b608e47781d2f541b2a8d893a48193e

Update README.md

Former-commit-id: 34e529fd02712d4466ab2bc1a63f9bd69976b6c3

Update README.md

Former-commit-id: d97edd0e714a237a54f191dcebb6fd39f73185d6

resume and compress kv

Former-commit-id: 0af7d58f88c1b65b4089da6f3e5c70456614da7c

refactor and fix resume bug

Former-commit-id: 1b1136671b9a7e0b15b70168e2c44dbff04a22e2

Update train.py

Former-commit-id: 783faf5cd86c148a713ea986ceed8527b108f2b8

Update train_t2v.py

Former-commit-id: d25b39ffac4474abd40b551fe0e51c421cb7bf5e

Update README.md

Former-commit-id: 6729039f059a23bf71faa944e2498657ef2c7442

[docs]: fix VQVAE script path

Former-commit-id: d5f5c681b8c50622a84256881d193d5803ff3c28

Update pyproject.toml

Former-commit-id: 5637656a79faba1b397b56918c42022cc209e731

[feat] add casual vqvae

Former-commit-id: 89087970b267908c8667984aba42b1ade2505467

Update __init__.py

Former-commit-id: 7c8af2909c4549e96bc0893d09fcc53d71fb6cc7

Create clip.py

Former-commit-id: e3397c567efcfeef6f58dfef2eb4251bedbbf33a

Update clip.py

Former-commit-id: 1ee53d152630b6acc69369be33287abfb143224d

train with image

Former-commit-id: 209b9a8ad2f19a50e46f9b7ec68dd781898e491a

support attention_mode

Former-commit-id: c72f354fbfd948e8cd4b0cc63c157bc4688c7665

Update README.md

Former-commit-id: 9065d12ee867c76452ce370fb71cc83bdc01470b

Update README.md

Former-commit-id: 95b28706d257d7769049abe30a7121b60e39a53f

t2v attention_mode

Former-commit-id: 4fd8bd87e0d0a98d5d0a3201d5db9a5ba055b894

Update README.md

Former-commit-id: 859c0f60dac1562b8a6bbb72b793a4709dbc3597

Update Data.md

Former-commit-id: d1c72e3209769d0204efb8e7d705bdf4f22b82ea

train t2v feature

Former-commit-id: 7f16b162cae0e3533c8ba9aa03186a846d9e02bf

scripts

Former-commit-id: 047046f684c2831e68e0cf256cb46fc3790cda33

add causalvae

Former-commit-id: ac936276fd3213ee1815dfe096cd64cecb601428

Create causalvideovae.md

Former-commit-id: 98414c2ad7dd39f740440b2d9059140f93492e83

Rename causalvideovae.md to CausalVideoVAE.md

Former-commit-id: 0a78e4193fcaec4662b7c6465a9d80643b354cd6

Update README.md

Former-commit-id: 8d70a890fc4191bfc98c3aca6dd3c8a70534fb3d

Update README.md

Former-commit-id: a41674f5ef71ced59b374226dddfa09c15a99f25

Update CausalVideoVAE.md

Former-commit-id: 1cd665f699632e9dc9357379225c085d9c1278d5

Update README.md

Former-commit-id: 90d6d298eaee6255098b342b31444e7a347fc50e

Update README.md

Former-commit-id: f294a0a28f288ae592e7de004320d3eb5e48e945

Update README.md

Former-commit-id: 02b0422fd745b4b7ec091cc6ee81d1d237245fe6

[feat] eval

Former-commit-id: d2c3cf5ef8d18feb7b0544d68d9796656e90ba8d

fix training bug

Former-commit-id: 09a6283ddd9ab11479f30324de788bd7415702ae

update model

Former-commit-id: e4e24650376d9dce974b290752efb416d5061983

update vae

Former-commit-id: 24298df5d374d5aaa5696f6200d4e23ae887715f

fix videovae trainer

Former-commit-id: 79e1feb412b6eeead3306f7f0d55a01f9b2a543e

Update CausalVideoVAE.md

Former-commit-id: fcde4d0ff783448c2006d6f4751a0e4f9cc23d0f

Update README.md

Former-commit-id: 2d9ae56309717a91e0690bd5c354fe3cbe5d5ed6

Update __init__.py

Former-commit-id: 2bdc29d86ebff3d6f29c38ded57c7b0dd1e6bbdb

update sample

Former-commit-id: 116726b9d565c689496123d4b6e62b98400536a6

tile conv

Former-commit-id: 5133615bcba221c4b86ae4e8b09b82ee5d34c66d

fix reshape bugs in AttnBlock3D

Former-commit-id: 299622e168a6cbee14ef2056c47cae722c20e005

tile only2d

Former-commit-id: 7f31ad03fdbaac97aea51a98bd29dc7327b700c9

released v1.0.0

Former-commit-id: 8df5897a4ffc341ced45faf617687fb8ebf1286c

sample pipeline

Former-commit-id: 8964b16c0dd65445a20050c9a0e82d909b55eaff

add tile training

Former-commit-id: 216079dc925beefe1ec722a4e87e46a4e1ca56be

update train script

Former-commit-id: ac16f77d19488cd46130587369cd7a5ea378052b

refactor vae to hf

Former-commit-id: a938d3f7f962709f7ecff5e9d49c9bf1ce40d7fa

release

Former-commit-id: 52fe2ac6610661d6a13a08390c84a8237f37986a

v1.0.0

Former-commit-id: 93f44f3b425c090d9bc92687a9c50a26b4f0d7c8

clean

Former-commit-id: 56ed87088e13cdd93929e2ceedbfae693e2f013a

update gradio demo

Former-commit-id: 80c2cc0bcc2a7532bc2440df16c1f811b0b30516

Update and rename CausalVideoVAE.md to Report-v1.0.0.md

Former-commit-id: 3b848f8de33cf4f0ed802428e9ad31d2fa26de63

Update README.md

Former-commit-id: e6fb5d864ac8af8214cdc89c71b42f296dda14e7

Update README.md

Former-commit-id: a61200bea2eea4fd22bd8095f192fef4c22c49b4

Update README.md

Former-commit-id: a2ec90a19eacb099cb9f36c1b5dd6e24be307fae

Add files via upload

Former-commit-id: 1ae8512b7846e9faafdb38ca532898d7f1bdcce2

Add files via upload

Former-commit-id: 292cbcfa4e4f384fe6958badc27d3067173119f9

Update Report-v1.0.0.md

Former-commit-id: 92d906b340f5970a77377586d7295f46b5de99ec

Update README.md

Former-commit-id: 037d7d3e084d7ff426347db49c4b210a2ee3927d

Update README.md

Former-commit-id: 7d4118566b77e09de4a3615612a5b368c76c293e

Update README.md

Former-commit-id: 3210474a46845ae3a6c0bde7fb6c7f9845866628

Add files via upload

Former-commit-id: a232ec2d98b245a9b502a2c49fc4f468a86fbd64

Update README.md

Former-commit-id: d5ab411b2e2fd5d6abeda500895e3339167572ab

Update README.md

Former-commit-id: 30f61958e221c3d41c4aa5917f4511e8c0482cf9

Update README.md

Former-commit-id: f8cb7f38751d7a6ef586c293fb769f5d60ad5c59

Update Report-v1.0.0.md

Former-commit-id: bd9aa7991212d98a6224b855a3bfb2c12feb5950

Update Report-v1.0.0.md

Former-commit-id: aea1f4a82a525d3e6e1d1a1bb7631326b294cdcc

Update README.md

Former-commit-id: f709d75e7a7d8a08ec60132aa906f6b880af8e58

Add files via upload

Former-commit-id: 9a207435d9efac0b19771df5d2288160af1bdd6e

Update README.md

Former-commit-id: 9bae31c73ea902581e1ee94cfd9012efdf78d349

Update README.md

Former-commit-id: 9f926abb08f156807091dd044c133788649a6bac

Update Report-v1.0.0.md

Former-commit-id: 7cec40420601f7da57449586508e50119bd02a61

Update README.md

Former-commit-id: ce2bbf82045041495ed448b46d5f4bdba84aa740

Fix some compatibility issues

Former-commit-id: 942d8598b28516369780c3519607929869b427b7

Fix bugs in inference

Former-commit-id: c251a86ac7c91d6f2bb12ae645628e9cf4b83d3d

fix bugs in inference

Former-commit-id: e17c7a7b7704920cb86894419bf573d0fb0c06fa

Update README.md

Former-commit-id: c1492211ced3cdbec93ef28678a1467245e122e9

Update README.md

Former-commit-id: c4a976e9b93536299b2e350454c780fec54fbe95

Update README.md

Former-commit-id: bbf3cef43a129a25a80ce134a8895796e24353c6

Update README.md

Former-commit-id: 5731e810e2efb7ca7eda1aa8954eb7fd260f4250

Update README.md

Former-commit-id: 1b54444249fa91b29334b8fc364fb1978c25fb2d

Update README.md

Former-commit-id: 9c0f82ae9953b56f2540e9c7ada021e4e1bc6165

Update Report-v1.0.0.md

Former-commit-id: 38b779022418cc41cdf9fc508e461f2f6299fd60

Update Report-v1.0.0.md

Former-commit-id: 6ebe2c62c0904142026d3796027e95ce43afe16a

Update README.md

Former-commit-id: 2873d0eaedfad392c627a714448118949d7b7e13

Update Report-v1.0.0.md

Former-commit-id: 400f708a080f0799e02697bb8e66d8430ffb1f56

Update Report-v1.0.0.md

Former-commit-id: 37f181f34f64ca1141a1345b7c1f2af283be1cad

Update README.md

Former-commit-id: dc923fe2413b3be31d8d88e7dd83a931fcbf9458

Update README.md

Former-commit-id: 828be4267f0795969293e4a321e4e0bfa8e1f1e6

clean

Former-commit-id: a7f2c1c8b2243587cedd6e152e64298842287095

Update pyproject.toml

Former-commit-id: bb037420b75303b58486ba79da9180b2bf63cf88

Update Data.md

Former-commit-id: d9a5c8fbd36a34b7baf720bc27d69604ae5c9d57

add causalvae doc

Former-commit-id: 4e6d9e038c8edb81f8402834bd01eeeb823924f9

Update Train_And_Eval_CausalVideoVAE.md

Former-commit-id: 589657b089930379bb75f3d49063ebb72671efc6

Update Train_And_Eval_CausalVideoVAE.md

Former-commit-id: 156739c35b5dda594138b2e3b1fd6ec99556af23

Update README.md

Former-commit-id: 567ade18047a9cc81648997c5d9287fa8900614d

更新 README.md

Former-commit-id: 42607786f3f3d0eba9336f774b0a2ae076cce793

Update README.md

Former-commit-id: 88e7f2062c494c199caff468f5f14f3165f763de

Create Report-v1.0.0-cn.md

Former-commit-id: d484754e080bd65f6b35619b840040581fa6b69c

Update README.md

Former-commit-id: e386ec112337271bb9f55ea6ec872eaba94aff50

Update README.md

Former-commit-id: c437a18dbec7ac6e02287565a84d75f96664ee99

Update README.md

Former-commit-id: 0cc27a5b065b61f37d44dc2c90ba2782075eefb2

Update Report-v1.0.0-cn.md

Former-commit-id: d5cf6aed7c1673562a6d2022c3fd7f7b16006cfd

Update Data.md

Former-commit-id: 15defc1a3846335fd68e18fa9b8ffbe38e3e7100

Update Report-v1.0.0.md

Former-commit-id: 3c8eb75dc2045154e8654f860b1fabd7e597094b

Update Report-v1.0.0-cn.md

Former-commit-id: 0edac24fcaa447a534755bf3aac1f2027a491c06

Update Report-v1.0.0-cn.md

Former-commit-id: 59edce70edbaed957ba40f1d95b4a5050ece12bf

Update Report-v1.0.0-cn.md

Former-commit-id: 83124efeaa5a8cd528a15760bdf88acc403142bd

Add files via upload

Former-commit-id: 943dd9a063baaa05d7f0d0bd52cbde35cf2b3d98 [formerly 5542359738243a8e948a8f08ed4e8dedf6b4b818]
Former-commit-id: fbee07f6fef2ac86bc15ff66f2e2a1a68b64ed01

Update README.md

Former-commit-id: 6fd60c73b6ae69d6ad9dc14c7371bd9bff086b44 [formerly 2ed5d5b7bdcd88044da26ca94237c50af7ab47dd]
Former-commit-id: 443c333353963bc1ec64a4946afbe5e17f9b5a87

Update README.md

Former-commit-id: 54bad83bf054a2ff574199931121c686b4e28008 [formerly 06bb190a39a9b98c8de5e385fc304a86cc8c9fb2]
Former-commit-id: 23e8a235d6828d9ab7c02ef59d40fe13a46fa4b1

Update README.md

Former-commit-id: c48eba0bb15d594007809d4cc605dc46ba029cec [formerly 7deb8a0decf7a9757af4475e3029d415d78d185f]
Former-commit-id: 7f43bb2f1d9308bd08154171d78e9b36e00a837f

Update README.md

Former-commit-id: 68eaf2e6b00266cf0858b33f9b4d09681efbd2bc [formerly d8dcad250ed8357964ab8ee570b9973b45e29fcb]
Former-commit-id: 75f51b3d84203be3672f6b07bb3c70ba0ae5556c

Update README.md

Former-commit-id: be00a17729ac889493f63fd1d649054d7980ca42 [formerly 9b4d8dcfc8e6e2abe4ca727e01561d27d1ba3968]
Former-commit-id: f1b6f5013b99d0d288af0f05c46cb8300d79a8aa

Update train_videoae_65x256x256.sh

Former-commit-id: b2f5cfe72452dfb369aba166e4a101dfcf5758c9 [formerly 557b78ab38a61d7520e91044137f818e50ec24cb]
Former-commit-id: 0e46c9f261e7f5fd4d08e3d7735c17c7975b0e72

Update train_videoae_65x512x512.sh

Former-commit-id: 8e40db106efa2fd879cca42bd9ce2d9b79d777f8 [formerly 3e11781b045a838a5aa98ea75e274198f2b26a44]
Former-commit-id: f3b0804a7cf063e7ff04ee57dccee0320c1f32a5

fix using pretrained bug

Former-commit-id: 2dfbb0683b2f3bd1a1d619e5e4fabb5bd4c92d79 [formerly 27f352e822f732cb351d64ac5726c5df3b62c1bb]
Former-commit-id: 722d60ae5ef28499c47292e0a4807a87b9658b01

Update train_t2v_feature.py

Former-commit-id: 4a4f7ea2366f1dff5669d6d4834238ebd790e6e2 [formerly ca7d3158fb01f8e7228f833cea8e6741be127c2a]
Former-commit-id: 389f0e7dc66826026046b8e12d7a9f8692d32e42

[docs] update docs and train.sh

Former-commit-id: 1d9358e84d12863681b396c7621ae7244caa4dc2 [formerly b801491889b3abad1211225d7f07af3ae0389895]
Former-commit-id: 9e6a59ae2f246b50ace0a033cbd6548660241204

[fix] fix bug in #185

Former-commit-id: d74efbd0195380ca383598c06d3d8185dd9cc9bc [formerly 0f2687266bc2271e11b2f06edfe514f466e57b0f]
Former-commit-id: cb6bb234befc49b3870c5690653469a0714890c5

[docs] add evaluation detail

Former-commit-id: c15a717cc61159e5f6adafb718237ab578ac81e7 [formerly 9e92772776fa4dd18772c7fec939a42ad6dfbc0a]
Former-commit-id: 98c2a38af5064f27888d1b102c9fded87d2420d5

[docs] update train.sh

Former-commit-id: eeb00ac046692bc97a064ec1b47c994436798ab1 [formerly 6d0fd85e1b13a7824b8181ba9ccebec2a3c76d40]
Former-commit-id: 1e8e669ad9967aeeeb7622bc04a0759f348cb26f

Delete assets/we_want_you.jpg

Update README.md

fix typo

Fix typos

line 24: `github -> GitHub`
line 54: `a -> an`
line 85: `re-organizes -> re-organize, modulizes -> modulize`
line 87: `opened -> open` (modified according to the changelogs of other dates)
line 94: `a -> an`

[fix]: Fix variable naming errors

Update README.md

[docs] update inference example.

[feat] add time chunk inference

[docs] fix typo

[refactor] fix hardcode

Update LICENSE

Update Report-v1.0.0.md

Update pyproject.toml

add 2drope and dynamic training

refactor dynamic training

read image from folder

img training

img training

support newvae

vae temporal tiling

abs and rope

compress kv and rope pi

fix rope with compress

update dataset

mask loss

multi-data

fix dataset

train vis

update

fix mask

5.15

prepare release

prepare scripts

fix demo

5.27

update prompt

Update README.md

Update README.md

Create Report-v1.1.0.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update README.md

Update LICENSE

Update Report-v1.1.0.md

solve typos

Update pyproject.toml

Update README.md

Update dataset_utils.py

Update modeling_latte.py

Update train_t2v.py

Update train_t2v.py

Update t2v_datasets.py

fix wrong code

Update LICENSE

fix train bug

fix vis

[docs] update CausalVideoVAE docs

Update README.md

fix the bug

Update gradio_web_server.py

Update README.md

Update gradio_utils.py

Update sample_video_513.sh

Rename sample_video_513.sh to sample_video_221.sh

Update README.md

Update Report-v1.1.0.md

release v1.2.0

Update README.md

Update Report-v1.2.0.md

Update README.md

Update pyproject.toml

Update Report-v1.2.0.md

Update Report-v1.2.0.md

Update Report-v1.2.0.md

Update t2v_datasets.py

Update Report-v1.2.0.md

Update Report-v1.2.0.md

Update README.md

Update Report-v1.2.0.md

Update Report-v1.2.0.md

Update README.md

fix sample

Delete opensora/train/train_t2v_diffusers_lora.py

lora bug

Update README.md

Update Report-v1.2.0.md

Update train_t2v_diffusers.py

release Open-Sora Plan v1.2.0 i2v (#389)

* gitignore

* inpaint

* Create condition_image_path.txt

* Rename condition_image_path.txt to condition_images_path.txt

* Update sample_inpaint.py

* Update sample_inpaint.sh

* inpaint

* fix bug

* Update README.md

* inpainting

* release v1.2.0 i2v

* release v1.2.0 i2v

* Update Report-v1.2.0.md

* Update pipeline_inpaint_sp.py

* Update sample_inpaint_ddp.py

* Update sample_inpaint_sp.py

* Update sample_inpaint.py

* Grammar Error Correction

* Grammar Error Correction

---------

Co-authored-by: LinB203 <2267330597@qq.com>

Update Report-v1.2.0.md

add 29x480p link

Update README.md

Update train_inpaint.sh

do not set seed

[fix] fix the path typo for google/mt5-xxl in gradio_web_server.py (#405)

* [fix] fix the path typo for google/mt5-xxl in gradio_web_server.py

* [fix] fix the issue: the cache of pretrained mt5-xxl weights is inconsistent.

---------

Co-authored-by: guangyi <guangyi.liu@mbz-h100-029.core42.ai>
---
 .github/workflows/docker_build.yml            |   38 +
 .gitignore                                    |   49 +
 LICENSE                                       |   21 +
 README.md                                     |  453 +++
 docker/LICENSE                                |   21 +
 docker/README.md                              |   87 +
 docker/build_docker.png                       |  Bin 0 -> 64398 bytes
 docker/docker_build.sh                        |    8 +
 docker/docker_run.sh                          |   45 +
 docker/dockerfile.base                        |   24 +
 docker/packages.txt                           |    3 +
 docker/ports.txt                              |    1 +
 docker/postinstallscript.sh                   |    3 +
 docker/requirements.txt                       |   40 +
 docker/run_docker.png                         |  Bin 0 -> 88368 bytes
 docker/setup_env.sh                           |   11 +
 docs/Contribution_Guidelines.md               |   87 +
 docs/Report-v1.0.0-cn.md                      |  135 +
 docs/Report-v1.0.0.md                         |  136 +
 docs/Report-v1.1.0.md                         |  256 ++
 docs/Report-v1.2.0.md                         |  193 +
 examples/conditional_images_path.txt          |    2 +
 examples/get_latents_std.py                   |   38 +
 examples/prompt_list_0.txt                    |   20 +
 examples/rec_image.py                         |   67 +
 examples/rec_video.py                         |  131 +
 opensora/__init__.py                          |    1 +
 opensora/acceleration/__init__.py             |    0
 opensora/acceleration/communications.py       |  133 +
 opensora/acceleration/parallel_states.py      |   57 +
 opensora/adaptor/__init__.py                  |    0
 opensora/adaptor/bf16_optimizer.py            |  458 +++
 opensora/adaptor/engine.py                    | 3568 +++++++++++++++++
 opensora/adaptor/modules.py                   |   23 +
 opensora/adaptor/stage_1_and_2.py             | 2522 ++++++++++++
 opensora/adaptor/utils.py                     | 1042 +++++
 opensora/adaptor/zp_manager.py                |   31 +
 opensora/dataset/__init__.py                  |   92 +
 opensora/dataset/inpaint_datasets.py          |  116 +
 opensora/dataset/t2v_datasets.py              |  422 ++
 opensora/dataset/transform.py                 |  592 +++
 opensora/eval/cal_flolpips.py                 |   83 +
 opensora/eval/cal_fvd.py                      |   85 +
 opensora/eval/cal_lpips.py                    |   97 +
 opensora/eval/cal_psnr.py                     |   84 +
 opensora/eval/cal_ssim.py                     |  113 +
 opensora/eval/eval_clip_score.py              |  225 ++
 opensora/eval/eval_common_metric.py           |  224 ++
 .../eval/flolpips/correlation/correlation.py  |  397 ++
 opensora/eval/flolpips/flolpips.py            |  308 ++
 opensora/eval/flolpips/pretrained_networks.py |  180 +
 opensora/eval/flolpips/pwcnet.py              |  344 ++
 opensora/eval/flolpips/utils.py               |   95 +
 opensora/eval/fvd/styleganv/fvd.py            |   90 +
 opensora/eval/fvd/videogpt/fvd.py             |  137 +
 opensora/eval/fvd/videogpt/pytorch_i3d.py     |  322 ++
 opensora/eval/script/cal_clip_score.sh        |   23 +
 opensora/eval/script/cal_fvd.sh               |    9 +
 opensora/eval/script/cal_lpips.sh             |    8 +
 opensora/eval/script/cal_psnr.sh              |    9 +
 opensora/eval/script/cal_ssim.sh              |    8 +
 opensora/models/__init__.py                   |    1 +
 opensora/models/causalvideovae/__init__.py    |   32 +
 .../models/causalvideovae/dataset/__init__.py |   65 +
 .../dataset/feature_datasets.py               |  213 +
 .../causalvideovae/dataset/t2v_datasets.py    |  203 +
 .../causalvideovae/dataset/transform.py       |  573 +++
 .../models/causalvideovae/dataset/ucf101.py   |   80 +
 .../causalvideovae/eval/cal_flolpips.py       |   83 +
 .../models/causalvideovae/eval/cal_fvd.py     |   85 +
 .../models/causalvideovae/eval/cal_lpips.py   |   97 +
 .../models/causalvideovae/eval/cal_psnr.py    |   97 +
 .../models/causalvideovae/eval/cal_ssim.py    |  113 +
 .../causalvideovae/eval/eval_clip_score.py    |  225 ++
 .../causalvideovae/eval/eval_common_metric.py |  226 ++
 .../eval/flolpips/correlation/correlation.py  |  397 ++
 .../causalvideovae/eval/flolpips/flolpips.py  |  308 ++
 .../eval/flolpips/pretrained_networks.py      |  180 +
 .../causalvideovae/eval/flolpips/pwcnet.py    |  344 ++
 .../causalvideovae/eval/flolpips/utils.py     |   95 +
 .../causalvideovae/eval/fvd/styleganv/fvd.py  |   90 +
 .../causalvideovae/eval/fvd/videogpt/fvd.py   |  137 +
 .../eval/fvd/videogpt/pytorch_i3d.py          |  322 ++
 .../eval/script/cal_clip_score.sh             |   23 +
 .../causalvideovae/eval/script/cal_fvd.sh     |    9 +
 .../causalvideovae/eval/script/cal_lpips.sh   |    8 +
 .../causalvideovae/eval/script/cal_psnr.sh    |    9 +
 .../causalvideovae/eval/script/cal_ssim.sh    |    8 +
 .../models/causalvideovae/model/__init__.py   |    4 +
 .../model/causal_vae/__init__.py              |   29 +
 .../model/causal_vae/modeling_causalvae.py    |  628 +++
 .../model/configuration_videobase.py          |   44 +
 .../causalvideovae/model/dataset_videobase.py |  107 +
 .../models/causalvideovae/model/ema_model.py  |   31 +
 .../causalvideovae/model/losses/__init__.py   |    1 +
 .../model/losses/discriminator.py             |  181 +
 .../causalvideovae/model/losses/lpips.py      |  120 +
 .../model/losses/perceptual_loss.py           |  389 ++
 .../model/modeling_videobase.py               |   53 +
 .../causalvideovae/model/modules/__init__.py  |    6 +
 .../causalvideovae/model/modules/attention.py |  239 ++
 .../causalvideovae/model/modules/block.py     |    5 +
 .../causalvideovae/model/modules/conv.py      |  127 +
 .../causalvideovae/model/modules/normalize.py |  101 +
 .../causalvideovae/model/modules/ops.py       |   40 +
 .../causalvideovae/model/modules/quant.py     |  100 +
 .../model/modules/resnet_block.py             |  137 +
 .../model/modules/updownsample.py             |  385 ++
 .../causalvideovae/model/trainer_videobase.py |   26 +
 .../causalvideovae/model/utils/__init__.py    |    0
 .../model/utils/distrib_utils.py              |   42 +
 .../model/utils/module_utils.py               |   17 +
 .../model/utils/scheduler_utils.py            |    7 +
 .../causalvideovae/model/utils/video_utils.py |   10 +
 .../causalvideovae/sample/rec_video_vae.py    |  277 ++
 .../models/causalvideovae/utils/__init__.py   |    0
 .../causalvideovae/utils/dataset_utils.py     |  160 +
 .../models/causalvideovae/utils/downloader.py |   18 +
 opensora/models/causalvideovae/utils/utils.py |  627 +++
 opensora/models/diffusion/__init__.py         |   21 +
 .../models/diffusion/opensora/__init__.py     |    0
 .../diffusion/opensora/modeling_inpaint.py    |  373 ++
 .../diffusion/opensora/modeling_opensora.py   |  734 ++++
 opensora/models/diffusion/opensora/modules.py | 1383 +++++++
 opensora/models/diffusion/opensora/rope.py    |   98 +
 .../models/diffusion/udit/modeling_udit.py    | 1073 +++++
 opensora/models/diffusion/udit/modules.py     | 1173 ++++++
 opensora/models/diffusion/udit/rope.py        |  154 +
 .../models/diffusion/utils/curope/__init__.py |    4 +
 .../models/diffusion/utils/curope/curope.cpp  |   69 +
 .../models/diffusion/utils/curope/curope2d.py |   40 +
 .../models/diffusion/utils/curope/kernels.cu  |  108 +
 .../models/diffusion/utils/curope/setup.py    |   34 +
 opensora/models/diffusion/utils/pos_embed.py  |  243 ++
 .../diffusion/utils/transport/__init__.py     |   63 +
 .../diffusion/utils/transport/integrators.py  |  117 +
 .../models/diffusion/utils/transport/path.py  |  192 +
 .../diffusion/utils/transport/transport.py    |  443 ++
 .../models/diffusion/utils/transport/utils.py |   29 +
 .../frame_interpolation/cfgs/AMT-G.yaml       |    9 +
 .../frame_interpolation/interpolation.py      |  197 +
 .../frame_interpolation/networks/AMT-G.py     |  172 +
 .../frame_interpolation/networks/__init__.py  |    0
 .../networks/blocks/__init__.py               |    0
 .../networks/blocks/feat_enc.py               |  343 ++
 .../networks/blocks/ifrnet.py                 |  111 +
 .../networks/blocks/multi_flow.py             |   69 +
 .../networks/blocks/raft.py                   |  207 +
 opensora/models/frame_interpolation/readme.md |   17 +
 .../frame_interpolation/utils/__init__.py     |    0
 .../frame_interpolation/utils/build_utils.py  |   12 +
 .../frame_interpolation/utils/dist_utils.py   |   48 +
 .../frame_interpolation/utils/flow_utils.py   |  122 +
 .../models/frame_interpolation/utils/utils.py |  297 ++
 opensora/models/super_resolution/README.md    |   25 +
 .../super_resolution/basicsr/__init__.py      |    6 +
 .../basicsr/archs/__init__.py                 |   25 +
 .../basicsr/archs/arch_util.py                |  318 ++
 .../basicsr/archs/rgt_arch.py                 |  757 ++++
 .../basicsr/archs/vgg_arch.py                 |  161 +
 .../super_resolution/basicsr/data/__init__.py |  101 +
 .../basicsr/data/data_sampler.py              |   48 +
 .../basicsr/data/data_util.py                 |  280 ++
 .../basicsr/data/paired_image_dataset.py      |  113 +
 .../basicsr/data/prefetch_dataloader.py       |  125 +
 .../basicsr/data/single_image_dataset.py      |   69 +
 .../basicsr/data/transforms.py                |  179 +
 .../basicsr/losses/__init__.py                |   26 +
 .../basicsr/losses/loss_util.py               |   95 +
 .../super_resolution/basicsr/losses/losses.py |  492 +++
 .../basicsr/metrics/__init__.py               |   19 +
 .../basicsr/metrics/metric_util.py            |   45 +
 .../basicsr/metrics/psnr_ssim.py              |  128 +
 .../basicsr/models/__init__.py                |   30 +
 .../basicsr/models/base_model.py              |  380 ++
 .../basicsr/models/lr_scheduler.py            |   96 +
 .../basicsr/models/rgt_model.py               |  127 +
 .../basicsr/models/sr_model.py                |  235 ++
 .../super_resolution/basicsr/test_img.py      |   48 +
 .../basicsr/utils/__init__.py                 |   30 +
 .../basicsr/utils/dist_util.py                |   82 +
 .../basicsr/utils/file_client.py              |  167 +
 .../basicsr/utils/img_util.py                 |  172 +
 .../super_resolution/basicsr/utils/logger.py  |  213 +
 .../basicsr/utils/matlab_functions.py         |  359 ++
 .../super_resolution/basicsr/utils/misc.py    |  141 +
 .../super_resolution/basicsr/utils/options.py |  200 +
 .../basicsr/utils/registry.py                 |   82 +
 .../options/test/test_RGT_x2.yml              |   94 +
 .../options/test/test_RGT_x4.yml              |   94 +
 .../options/test/test_single_config.yml       |   41 +
 opensora/models/super_resolution/run.py       |  138 +
 opensora/models/text_encoder/__init__.py      |   76 +
 opensora/models/text_encoder/clip.py          |  222 +
 opensora/models/text_encoder/t5.py            |  197 +
 opensora/npu_config.py                        |  442 ++
 opensora/sample/pipeline_inpaint.py           |  904 +++++
 opensora/sample/pipeline_inpaint_sp.py        |  936 +++++
 opensora/sample/pipeline_opensora.py          |  820 ++++
 opensora/sample/pipeline_opensora_sp.py       |  841 ++++
 opensora/sample/sample_inpaint.py             |  264 ++
 opensora/sample/sample_inpaint_ddp.py         |  357 ++
 opensora/sample/sample_inpaint_sp.py          |  347 ++
 opensora/sample/sample_t2v.py                 |  236 ++
 opensora/sample/sample_t2v_ddp.py             |  343 ++
 opensora/sample/sample_t2v_on_npu.py          |  295 ++
 opensora/sample/sample_t2v_sp.py              |  322 ++
 opensora/sample/sample_t2v_sp_on_npu.py       |  273 ++
 opensora/serve/gradio_utils.py                |   88 +
 opensora/serve/gradio_web_server.py           |  310 ++
 opensora/serve/style.css                      |    1 +
 opensora/train/train_causalvae.py             |  583 +++
 opensora/train/train_inpaint.py               |  876 ++++
 opensora/train/train_t2v_diffusers.py         |  957 +++++
 opensora/utils/communications.py              |  122 +
 opensora/utils/dataset_utils.py               |  329 ++
 opensora/utils/downloader.py                  |   18 +
 opensora/utils/ema.py                         |  328 ++
 opensora/utils/lora_utils.py                  |   82 +
 opensora/utils/parallel_states.py             |   44 +
 opensora/utils/taming_download.py             |  145 +
 opensora/utils/utils.py                       |  471 +++
 pyproject.toml                                |   49 +
 scripts/accelerate_configs/ddp_config.yaml    |   11 +
 .../deepspeed_zero2_config.yaml               |   13 +
 .../deepspeed_zero2_offload_config.yaml       |   13 +
 .../deepspeed_zero3_config.yaml               |   14 +
 .../deepspeed_zero3_offload_config.yaml       |   13 +
 .../accelerate_configs/default_config.yaml    |   12 +
 scripts/accelerate_configs/hostfile           |   16 +
 .../multi_node_example.yaml                   |   18 +
 .../multi_node_example_by_ddp.yaml            |   13 +
 scripts/accelerate_configs/zero2.json         |   25 +
 scripts/accelerate_configs/zero2_npu.json     |   25 +
 scripts/accelerate_configs/zero2_offload.json |   26 +
 scripts/accelerate_configs/zero3.json         |   28 +
 scripts/accelerate_configs/zero3_offload.json |   39 +
 scripts/causalvae/eval.sh                     |   21 +
 scripts/causalvae/prepare_eval.sh             |   24 +
 scripts/causalvae/rec_image.sh                |    8 +
 scripts/causalvae/rec_video.sh                |   11 +
 scripts/causalvae/train.sh                    |   45 +
 scripts/slurm/placeholder                     |    0
 scripts/text_condition/gpu/sample_inpaint.sh  |   18 +
 .../text_condition/gpu/sample_inpaint_ddp.sh  |   21 +
 .../text_condition/gpu/sample_inpaint_sp.sh   |   21 +
 scripts/text_condition/gpu/sample_t2v.sh      |   18 +
 scripts/text_condition/gpu/sample_t2v_ddp.sh  |   21 +
 scripts/text_condition/gpu/sample_t2v_sp.sh   |   21 +
 scripts/text_condition/gpu/train_inpaint.sh   |   74 +
 scripts/text_condition/gpu/train_t2v.sh       |   66 +
 scripts/text_condition/npu/sample_inpaint.sh  |   34 +
 .../text_condition/npu/train_image3d_480p.sh  |   54 +
 .../text_condition/npu/train_image3d_720p.sh  |   54 +
 .../npu/train_inpaint_video3d_nx480p.sh       |   61 +
 .../npu/train_video3d_nx480p.sh               |   57 +
 .../npu/train_video3d_nx720p.sh               |   57 +
 .../npu/train_video3d_sp_nx480p.sh            |   59 +
 .../npu/train_video3d_sp_nx720p.sh            |   59 +
 scripts/train_data/merge_data.txt             |    2 +
 scripts/train_data/merge_data_on_npu.txt      |    1 +
 tools/get_img_info.py                         |   51 +
 tools/merge_imginfo_to_anno.py                |   42 +
 263 files changed, 47841 insertions(+)
 create mode 100644 .github/workflows/docker_build.yml
 create mode 100644 .gitignore
 create mode 100644 LICENSE
 create mode 100644 README.md
 create mode 100644 docker/LICENSE
 create mode 100644 docker/README.md
 create mode 100755 docker/build_docker.png
 create mode 100644 docker/docker_build.sh
 create mode 100644 docker/docker_run.sh
 create mode 100644 docker/dockerfile.base
 create mode 100644 docker/packages.txt
 create mode 100644 docker/ports.txt
 create mode 100644 docker/postinstallscript.sh
 create mode 100644 docker/requirements.txt
 create mode 100755 docker/run_docker.png
 create mode 100644 docker/setup_env.sh
 create mode 100644 docs/Contribution_Guidelines.md
 create mode 100644 docs/Report-v1.0.0-cn.md
 create mode 100644 docs/Report-v1.0.0.md
 create mode 100644 docs/Report-v1.1.0.md
 create mode 100644 docs/Report-v1.2.0.md
 create mode 100644 examples/conditional_images_path.txt
 create mode 100644 examples/get_latents_std.py
 create mode 100644 examples/prompt_list_0.txt
 create mode 100644 examples/rec_image.py
 create mode 100644 examples/rec_video.py
 create mode 100644 opensora/__init__.py
 create mode 100644 opensora/acceleration/__init__.py
 create mode 100644 opensora/acceleration/communications.py
 create mode 100644 opensora/acceleration/parallel_states.py
 create mode 100644 opensora/adaptor/__init__.py
 create mode 100644 opensora/adaptor/bf16_optimizer.py
 create mode 100644 opensora/adaptor/engine.py
 create mode 100644 opensora/adaptor/modules.py
 create mode 100644 opensora/adaptor/stage_1_and_2.py
 create mode 100644 opensora/adaptor/utils.py
 create mode 100644 opensora/adaptor/zp_manager.py
 create mode 100644 opensora/dataset/__init__.py
 create mode 100644 opensora/dataset/inpaint_datasets.py
 create mode 100644 opensora/dataset/t2v_datasets.py
 create mode 100644 opensora/dataset/transform.py
 create mode 100644 opensora/eval/cal_flolpips.py
 create mode 100644 opensora/eval/cal_fvd.py
 create mode 100644 opensora/eval/cal_lpips.py
 create mode 100644 opensora/eval/cal_psnr.py
 create mode 100644 opensora/eval/cal_ssim.py
 create mode 100755 opensora/eval/eval_clip_score.py
 create mode 100644 opensora/eval/eval_common_metric.py
 create mode 100644 opensora/eval/flolpips/correlation/correlation.py
 create mode 100644 opensora/eval/flolpips/flolpips.py
 create mode 100644 opensora/eval/flolpips/pretrained_networks.py
 create mode 100644 opensora/eval/flolpips/pwcnet.py
 create mode 100644 opensora/eval/flolpips/utils.py
 create mode 100644 opensora/eval/fvd/styleganv/fvd.py
 create mode 100644 opensora/eval/fvd/videogpt/fvd.py
 create mode 100644 opensora/eval/fvd/videogpt/pytorch_i3d.py
 create mode 100644 opensora/eval/script/cal_clip_score.sh
 create mode 100644 opensora/eval/script/cal_fvd.sh
 create mode 100644 opensora/eval/script/cal_lpips.sh
 create mode 100644 opensora/eval/script/cal_psnr.sh
 create mode 100644 opensora/eval/script/cal_ssim.sh
 create mode 100644 opensora/models/__init__.py
 create mode 100644 opensora/models/causalvideovae/__init__.py
 create mode 100755 opensora/models/causalvideovae/dataset/__init__.py
 create mode 100755 opensora/models/causalvideovae/dataset/feature_datasets.py
 create mode 100755 opensora/models/causalvideovae/dataset/t2v_datasets.py
 create mode 100755 opensora/models/causalvideovae/dataset/transform.py
 create mode 100755 opensora/models/causalvideovae/dataset/ucf101.py
 create mode 100755 opensora/models/causalvideovae/eval/cal_flolpips.py
 create mode 100755 opensora/models/causalvideovae/eval/cal_fvd.py
 create mode 100755 opensora/models/causalvideovae/eval/cal_lpips.py
 create mode 100755 opensora/models/causalvideovae/eval/cal_psnr.py
 create mode 100755 opensora/models/causalvideovae/eval/cal_ssim.py
 create mode 100755 opensora/models/causalvideovae/eval/eval_clip_score.py
 create mode 100755 opensora/models/causalvideovae/eval/eval_common_metric.py
 create mode 100755 opensora/models/causalvideovae/eval/flolpips/correlation/correlation.py
 create mode 100755 opensora/models/causalvideovae/eval/flolpips/flolpips.py
 create mode 100755 opensora/models/causalvideovae/eval/flolpips/pretrained_networks.py
 create mode 100755 opensora/models/causalvideovae/eval/flolpips/pwcnet.py
 create mode 100755 opensora/models/causalvideovae/eval/flolpips/utils.py
 create mode 100755 opensora/models/causalvideovae/eval/fvd/styleganv/fvd.py
 create mode 100755 opensora/models/causalvideovae/eval/fvd/videogpt/fvd.py
 create mode 100755 opensora/models/causalvideovae/eval/fvd/videogpt/pytorch_i3d.py
 create mode 100755 opensora/models/causalvideovae/eval/script/cal_clip_score.sh
 create mode 100755 opensora/models/causalvideovae/eval/script/cal_fvd.sh
 create mode 100755 opensora/models/causalvideovae/eval/script/cal_lpips.sh
 create mode 100755 opensora/models/causalvideovae/eval/script/cal_psnr.sh
 create mode 100755 opensora/models/causalvideovae/eval/script/cal_ssim.sh
 create mode 100755 opensora/models/causalvideovae/model/__init__.py
 create mode 100755 opensora/models/causalvideovae/model/causal_vae/__init__.py
 create mode 100755 opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py
 create mode 100755 opensora/models/causalvideovae/model/configuration_videobase.py
 create mode 100755 opensora/models/causalvideovae/model/dataset_videobase.py
 create mode 100755 opensora/models/causalvideovae/model/ema_model.py
 create mode 100755 opensora/models/causalvideovae/model/losses/__init__.py
 create mode 100755 opensora/models/causalvideovae/model/losses/discriminator.py
 create mode 100755 opensora/models/causalvideovae/model/losses/lpips.py
 create mode 100755 opensora/models/causalvideovae/model/losses/perceptual_loss.py
 create mode 100755 opensora/models/causalvideovae/model/modeling_videobase.py
 create mode 100755 opensora/models/causalvideovae/model/modules/__init__.py
 create mode 100755 opensora/models/causalvideovae/model/modules/attention.py
 create mode 100755 opensora/models/causalvideovae/model/modules/block.py
 create mode 100755 opensora/models/causalvideovae/model/modules/conv.py
 create mode 100755 opensora/models/causalvideovae/model/modules/normalize.py
 create mode 100755 opensora/models/causalvideovae/model/modules/ops.py
 create mode 100755 opensora/models/causalvideovae/model/modules/quant.py
 create mode 100755 opensora/models/causalvideovae/model/modules/resnet_block.py
 create mode 100755 opensora/models/causalvideovae/model/modules/updownsample.py
 create mode 100755 opensora/models/causalvideovae/model/trainer_videobase.py
 create mode 100644 opensora/models/causalvideovae/model/utils/__init__.py
 create mode 100755 opensora/models/causalvideovae/model/utils/distrib_utils.py
 create mode 100755 opensora/models/causalvideovae/model/utils/module_utils.py
 create mode 100755 opensora/models/causalvideovae/model/utils/scheduler_utils.py
 create mode 100755 opensora/models/causalvideovae/model/utils/video_utils.py
 create mode 100644 opensora/models/causalvideovae/sample/rec_video_vae.py
 create mode 100644 opensora/models/causalvideovae/utils/__init__.py
 create mode 100755 opensora/models/causalvideovae/utils/dataset_utils.py
 create mode 100755 opensora/models/causalvideovae/utils/downloader.py
 create mode 100755 opensora/models/causalvideovae/utils/utils.py
 create mode 100644 opensora/models/diffusion/__init__.py
 create mode 100644 opensora/models/diffusion/opensora/__init__.py
 create mode 100644 opensora/models/diffusion/opensora/modeling_inpaint.py
 create mode 100644 opensora/models/diffusion/opensora/modeling_opensora.py
 create mode 100644 opensora/models/diffusion/opensora/modules.py
 create mode 100644 opensora/models/diffusion/opensora/rope.py
 create mode 100644 opensora/models/diffusion/udit/modeling_udit.py
 create mode 100644 opensora/models/diffusion/udit/modules.py
 create mode 100644 opensora/models/diffusion/udit/rope.py
 create mode 100644 opensora/models/diffusion/utils/curope/__init__.py
 create mode 100644 opensora/models/diffusion/utils/curope/curope.cpp
 create mode 100644 opensora/models/diffusion/utils/curope/curope2d.py
 create mode 100644 opensora/models/diffusion/utils/curope/kernels.cu
 create mode 100644 opensora/models/diffusion/utils/curope/setup.py
 create mode 100644 opensora/models/diffusion/utils/pos_embed.py
 create mode 100644 opensora/models/diffusion/utils/transport/__init__.py
 create mode 100644 opensora/models/diffusion/utils/transport/integrators.py
 create mode 100644 opensora/models/diffusion/utils/transport/path.py
 create mode 100644 opensora/models/diffusion/utils/transport/transport.py
 create mode 100644 opensora/models/diffusion/utils/transport/utils.py
 create mode 100644 opensora/models/frame_interpolation/cfgs/AMT-G.yaml
 create mode 100644 opensora/models/frame_interpolation/interpolation.py
 create mode 100644 opensora/models/frame_interpolation/networks/AMT-G.py
 create mode 100644 opensora/models/frame_interpolation/networks/__init__.py
 create mode 100644 opensora/models/frame_interpolation/networks/blocks/__init__.py
 create mode 100644 opensora/models/frame_interpolation/networks/blocks/feat_enc.py
 create mode 100644 opensora/models/frame_interpolation/networks/blocks/ifrnet.py
 create mode 100644 opensora/models/frame_interpolation/networks/blocks/multi_flow.py
 create mode 100644 opensora/models/frame_interpolation/networks/blocks/raft.py
 create mode 100644 opensora/models/frame_interpolation/readme.md
 create mode 100644 opensora/models/frame_interpolation/utils/__init__.py
 create mode 100644 opensora/models/frame_interpolation/utils/build_utils.py
 create mode 100644 opensora/models/frame_interpolation/utils/dist_utils.py
 create mode 100644 opensora/models/frame_interpolation/utils/flow_utils.py
 create mode 100644 opensora/models/frame_interpolation/utils/utils.py
 create mode 100644 opensora/models/super_resolution/README.md
 create mode 100644 opensora/models/super_resolution/basicsr/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/archs/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/archs/arch_util.py
 create mode 100644 opensora/models/super_resolution/basicsr/archs/rgt_arch.py
 create mode 100644 opensora/models/super_resolution/basicsr/archs/vgg_arch.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/data_sampler.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/data_util.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/paired_image_dataset.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/prefetch_dataloader.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/single_image_dataset.py
 create mode 100644 opensora/models/super_resolution/basicsr/data/transforms.py
 create mode 100644 opensora/models/super_resolution/basicsr/losses/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/losses/loss_util.py
 create mode 100644 opensora/models/super_resolution/basicsr/losses/losses.py
 create mode 100644 opensora/models/super_resolution/basicsr/metrics/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/metrics/metric_util.py
 create mode 100644 opensora/models/super_resolution/basicsr/metrics/psnr_ssim.py
 create mode 100644 opensora/models/super_resolution/basicsr/models/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/models/base_model.py
 create mode 100644 opensora/models/super_resolution/basicsr/models/lr_scheduler.py
 create mode 100644 opensora/models/super_resolution/basicsr/models/rgt_model.py
 create mode 100644 opensora/models/super_resolution/basicsr/models/sr_model.py
 create mode 100644 opensora/models/super_resolution/basicsr/test_img.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/__init__.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/dist_util.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/file_client.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/img_util.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/logger.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/matlab_functions.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/misc.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/options.py
 create mode 100644 opensora/models/super_resolution/basicsr/utils/registry.py
 create mode 100644 opensora/models/super_resolution/options/test/test_RGT_x2.yml
 create mode 100644 opensora/models/super_resolution/options/test/test_RGT_x4.yml
 create mode 100644 opensora/models/super_resolution/options/test/test_single_config.yml
 create mode 100644 opensora/models/super_resolution/run.py
 create mode 100644 opensora/models/text_encoder/__init__.py
 create mode 100644 opensora/models/text_encoder/clip.py
 create mode 100644 opensora/models/text_encoder/t5.py
 create mode 100644 opensora/npu_config.py
 create mode 100644 opensora/sample/pipeline_inpaint.py
 create mode 100644 opensora/sample/pipeline_inpaint_sp.py
 create mode 100644 opensora/sample/pipeline_opensora.py
 create mode 100644 opensora/sample/pipeline_opensora_sp.py
 create mode 100644 opensora/sample/sample_inpaint.py
 create mode 100644 opensora/sample/sample_inpaint_ddp.py
 create mode 100644 opensora/sample/sample_inpaint_sp.py
 create mode 100644 opensora/sample/sample_t2v.py
 create mode 100644 opensora/sample/sample_t2v_ddp.py
 create mode 100644 opensora/sample/sample_t2v_on_npu.py
 create mode 100644 opensora/sample/sample_t2v_sp.py
 create mode 100644 opensora/sample/sample_t2v_sp_on_npu.py
 create mode 100644 opensora/serve/gradio_utils.py
 create mode 100644 opensora/serve/gradio_web_server.py
 create mode 100644 opensora/serve/style.css
 create mode 100644 opensora/train/train_causalvae.py
 create mode 100644 opensora/train/train_inpaint.py
 create mode 100644 opensora/train/train_t2v_diffusers.py
 create mode 100644 opensora/utils/communications.py
 create mode 100644 opensora/utils/dataset_utils.py
 create mode 100644 opensora/utils/downloader.py
 create mode 100644 opensora/utils/ema.py
 create mode 100644 opensora/utils/lora_utils.py
 create mode 100644 opensora/utils/parallel_states.py
 create mode 100644 opensora/utils/taming_download.py
 create mode 100644 opensora/utils/utils.py
 create mode 100644 pyproject.toml
 create mode 100644 scripts/accelerate_configs/ddp_config.yaml
 create mode 100644 scripts/accelerate_configs/deepspeed_zero2_config.yaml
 create mode 100644 scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
 create mode 100644 scripts/accelerate_configs/deepspeed_zero3_config.yaml
 create mode 100644 scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml
 create mode 100644 scripts/accelerate_configs/default_config.yaml
 create mode 100644 scripts/accelerate_configs/hostfile
 create mode 100644 scripts/accelerate_configs/multi_node_example.yaml
 create mode 100644 scripts/accelerate_configs/multi_node_example_by_ddp.yaml
 create mode 100644 scripts/accelerate_configs/zero2.json
 create mode 100644 scripts/accelerate_configs/zero2_npu.json
 create mode 100644 scripts/accelerate_configs/zero2_offload.json
 create mode 100644 scripts/accelerate_configs/zero3.json
 create mode 100644 scripts/accelerate_configs/zero3_offload.json
 create mode 100644 scripts/causalvae/eval.sh
 create mode 100644 scripts/causalvae/prepare_eval.sh
 create mode 100644 scripts/causalvae/rec_image.sh
 create mode 100644 scripts/causalvae/rec_video.sh
 create mode 100644 scripts/causalvae/train.sh
 create mode 100644 scripts/slurm/placeholder
 create mode 100644 scripts/text_condition/gpu/sample_inpaint.sh
 create mode 100644 scripts/text_condition/gpu/sample_inpaint_ddp.sh
 create mode 100644 scripts/text_condition/gpu/sample_inpaint_sp.sh
 create mode 100644 scripts/text_condition/gpu/sample_t2v.sh
 create mode 100644 scripts/text_condition/gpu/sample_t2v_ddp.sh
 create mode 100644 scripts/text_condition/gpu/sample_t2v_sp.sh
 create mode 100644 scripts/text_condition/gpu/train_inpaint.sh
 create mode 100644 scripts/text_condition/gpu/train_t2v.sh
 create mode 100644 scripts/text_condition/npu/sample_inpaint.sh
 create mode 100644 scripts/text_condition/npu/train_image3d_480p.sh
 create mode 100644 scripts/text_condition/npu/train_image3d_720p.sh
 create mode 100644 scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh
 create mode 100644 scripts/text_condition/npu/train_video3d_nx480p.sh
 create mode 100644 scripts/text_condition/npu/train_video3d_nx720p.sh
 create mode 100644 scripts/text_condition/npu/train_video3d_sp_nx480p.sh
 create mode 100644 scripts/text_condition/npu/train_video3d_sp_nx720p.sh
 create mode 100644 scripts/train_data/merge_data.txt
 create mode 100644 scripts/train_data/merge_data_on_npu.txt
 create mode 100644 tools/get_img_info.py
 create mode 100644 tools/merge_imginfo_to_anno.py

diff --git a/.github/workflows/docker_build.yml b/.github/workflows/docker_build.yml
new file mode 100644
index 000000000..3dea684aa
--- /dev/null
+++ b/.github/workflows/docker_build.yml
@@ -0,0 +1,38 @@
+name: docker-build
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - "main"
+    paths:
+      - "docker/Dockerfile"
+
+jobs:
+  build-Open-Sora:
+    runs-on: ubuntu-latest
+    steps:
+      -
+        name: Checkout
+        uses: actions/checkout@v4
+      -
+        name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+      -
+        name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      -
+        name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      -
+        name: Build and push Open-Sora image
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          file: ./docker/Dockerfile
+          push: true
+          platforms: linux/amd64, linux/arm64, linux/s390x, linux/ppc64le
+          tags: ${{ secrets.DOCKERHUB_USERNAME }}/open-sora
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 000000000..391b25624
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,49 @@
+ucf101_stride4x4x4
+__pycache__
+*.mp4
+.ipynb_checkpoints
+*.pth
+UCF-101/
+results/
+vae
+build/
+opensora.egg-info/
+wandb/
+.idea
+*.ipynb
+*.jpg
+*.mp3
+*.safetensors
+*.mp4
+*.png
+*.gif
+*.pth
+*.pt
+cache_dir/
+wandb/
+test*
+sample_video*
+sample_image*
+512*
+720*
+1024*
+debug*
+private*
+caption*
+*deepspeed*
+revised*
+129f*
+all*
+read*
+YSH*
+*pick*
+*ysh*
+hw*
+257f*
+513f*
+taming*
+221hw*
+65x512x512
+runs/
+samples/
+*validation/
\ No newline at end of file
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 000000000..04d11366b
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 PKU-YUAN's Group (袁粒课题组-北大信工) and Rabbitpre AI
+
+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.
diff --git a/README.md b/README.md
new file mode 100644
index 000000000..eeeb4cb6d
--- /dev/null
+++ b/README.md
@@ -0,0 +1,453 @@
+
+
+<h1 align="left"> <a href="">Open-Sora Plan</a></h1>
+
+This project aims to create a simple and scalable repo, to reproduce [Sora](https://openai.com/sora) (OpenAI, but we prefer to call it "ClosedAI" ). We wish the open-source community can contribute to this project. Pull requests are welcome! The current code supports complete training and inference using the Huawei Ascend AI computing system. Models trained on Huawei Ascend can also output video quality comparable to industry standards.
+
+本项目希望通过开源社区的力量复现Sora，由北大-兔展AIGC联合实验室共同发起，当前版本离目标差距仍然较大，仍需持续完善和快速迭代，欢迎Pull request！目前代码同时支持使用国产AI计算系统（华为昇腾）进行完整的训练和推理。基于昇腾训练出的模型，也可输出持平业界的视频质量。
+
+<h5 align="left">
+
+[![slack badge](https://img.shields.io/badge/Discord-join-blueviolet?logo=discord&amp)](https://discord.gg/FkFm5M2J)
+[![WeChat badge](https://img.shields.io/badge/微信-加入-green?logo=wechat&amp)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/53#issuecomment-1987226516)
+[![Twitter](https://img.shields.io/badge/-Twitter@LinBin46984-black?logo=twitter&logoColor=1D9BF0)](https://x.com/LinBin46984/status/1795018003345510687) <br>
+[![License](https://img.shields.io/badge/License-MIT-yellow)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/blob/main/LICENSE) 
+[![GitHub repo contributors](https://img.shields.io/github/contributors-anon/PKU-YuanGroup/Open-Sora-Plan?style=flat&label=Contributors)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/graphs/contributors) 
+[![GitHub Commit](https://img.shields.io/github/commit-activity/m/PKU-YuanGroup/Open-Sora-Plan?label=Commit)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/commits/main/)
+[![Pr](https://img.shields.io/github/issues-pr-closed-raw/PKU-YuanGroup/Open-Sora-Plan.svg?label=Merged+PRs&color=green)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/pulls)
+[![GitHub issues](https://img.shields.io/github/issues/PKU-YuanGroup/Open-Sora-Plan?color=critical&label=Issues)](https://github.com/PKU-YuanGroup/Video-LLaVA/issues?q=is%3Aopen+is%3Aissue)
+[![GitHub closed issues](https://img.shields.io/github/issues-closed/PKU-YuanGroup/Open-Sora-Plan?color=success&label=Issues)](https://github.com/PKU-YuanGroup/Video-LLaVA/issues?q=is%3Aissue+is%3Aclosed) <br>
+[![GitHub repo stars](https://img.shields.io/github/stars/PKU-YuanGroup/Open-Sora-Plan?style=flat&logo=github&logoColor=whitesmoke&label=Stars)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/stargazers)&#160;
+[![GitHub repo forks](https://img.shields.io/github/forks/PKU-YuanGroup/Open-Sora-Plan?style=flat&logo=github&logoColor=whitesmoke&label=Forks)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/network)&#160;
+[![GitHub repo watchers](https://img.shields.io/github/watchers/PKU-YuanGroup/Open-Sora-Plan?style=flat&logo=github&logoColor=whitesmoke&label=Watchers)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/watchers)&#160;
+[![GitHub repo size](https://img.shields.io/github/repo-size/PKU-YuanGroup/Open-Sora-Plan?style=flat&logo=github&logoColor=whitesmoke&label=Repo%20Size)](https://github.com/PKU-YuanGroup/Open-Sora-Plan/archive/refs/heads/main.zip)
+</h5>
+<h5 align="left"> If you like our project, please give us a star ⭐ on GitHub for latest update.  </h2>
+
+
+# 📣 News
+* **[2024.08.13]** 🎉 We are launching Open-Sora Plan v1.2.0 **I2V** model, which based on Open-Sora Plan v1.2.0. The current version supports image-to-video generation and transition generation (the starting and ending frames conditions for video generation). Checking out the Image-to-Video section in this [report](https://github.com/PKU-YuanGroup/Open-Sora-Plan/blob/main/docs/Report-v1.2.0.md#training-image-to-video-diffusion-model).
+* **[2024.07.24]** 🔥🔥🔥 v1.2.0 is here! Utilizing a 3D full attention architecture instead of 2+1D. We released a true 3D video diffusion model trained on 4s 720p. Checking out our latest [report](docs/Report-v1.2.0.md).
+* **[2024.05.27]** 🎉 We are launching Open-Sora Plan v1.1.0, which significantly improves video quality and length, and is fully open source! Please check out our latest [report](docs/Report-v1.1.0.md). Thanks to [ShareGPT4Video's](https://sharegpt4video.github.io/) capability to annotate long videos.
+* **[2024.04.09]** 🤝 Excited to share our latest exploration on metamorphic time-lapse video generation: [MagicTime](https://github.com/PKU-YuanGroup/MagicTime), which learns real-world physics knowledge from time-lapse videos.
+* **[2024.04.07]** 🎉🎉🎉 Today, we are thrilled to present Open-Sora-Plan v1.0.0, which significantly enhances video generation quality and text control capabilities. See our [report](docs/Report-v1.0.0.md). Thanks to HUAWEI NPU for supporting us.
+* **[2024.03.27]** 🚀🚀🚀 We release the report of [VideoCausalVAE](docs/CausalVideoVAE.md), which supports both images and videos. We present our reconstructed video in this demonstration as follows. The text-to-video model is on the way.
+* **[2024.03.01]** 🤗 We launched a plan to reproduce Sora, called Open-Sora Plan! Welcome to **watch** 👀 this repository for the latest updates.
+
+# 😍 Gallery
+
+93×1280×720 Text-to-Video Generation. The video quality has been compressed for playback on GitHub.
+
+<table class="center">
+<tr>
+  <td><video src="https://github.com/user-attachments/assets/1c84bc92-d585-46c9-ae7c-e5f79cefea88" autoplay></td>
+</tr>
+</table>
+
+
+# 😮 Highlights
+
+Open-Sora Plan shows excellent performance in video generation.
+
+### 🔥 High performance CausalVideoVAE, but with fewer training cost
+- High compression ratio with excellent performance, capable of **compressing videos by 256 times (4×8×8)**. Causal convolution supports simultaneous inference of images and videos but only need **1 node to train**.
+
+### 🚀 Video Diffusion Model based on 3D attention, joint learning of spatiotemporal features.
+- With a **3D full attention architecture** instead of a 2+1D model, 3D attention can better capture joint spatial and temporal features.
+
+<p align="center">
+    <img src="https://s21.ax1x.com/2024/07/22/pk7cob8.png" width="650" style="margin-bottom: 0.2;"/>
+<p>
+
+# 🤗 Demo
+
+### Gradio Web UI
+
+Highly recommend trying out our web demo by the following command.
+
+```bash
+python -m opensora.serve.gradio_web_server --model_path "path/to/model" --ae_path "path/to/causalvideovae"
+```
+
+### ComfyUI
+
+Coming soon...
+
+# 🐳 Resource
+
+| Version | Architecture |  Diffusion Model | CausalVideoVAE | Data|
+|:---|:---|:---|:---|:---|
+| v1.2.0 | 3D | [93x720p](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/93x720p), [29x720p](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/29x720p)[1], [93x480p](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/93x480p)[1,2], [29x480p](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/29x480p), [1x480p](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/1x480p), [93x480p_i2v](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/93x480p_i2v) | [Anysize](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/vae)| [Annotations](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0) |
+| v1.1.0 | 2+1D | [221x512x512](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/221x512x512), [65x512x512](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/65x512x512) |[Anysize](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/vae) |[Data and Annotations](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0)|
+| v1.0.0 | 2+1D | [65x512x512](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/65x512x512), [65x256x256](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/65x256x256), [17x256x256](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/17x256x256) | [Anysize](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0/tree/main/vae) | [Data and Annotations](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.0.0)|
+
+> [1] Please note that the weights for v1.2.0 29×720p and 93×480p were trained on Panda70M and have not undergone final high-quality data fine-tuning, so they may produce watermarks.
+
+> [2] We fine-tuned 3.5k steps from 93×720p to get 93×480p for community research use.
+
+> [!Warning]
+>
+> <div align="left">
+> <b>
+> 🚨 For version 1.2.0, we no longer support 2+1D models.
+> </b>
+> </div>
+
+# ⚙️ Requirements and Installation
+
+1. Clone this repository and navigate to Open-Sora-Plan folder
+```
+git clone https://github.com/PKU-YuanGroup/Open-Sora-Plan
+cd Open-Sora-Plan
+```
+2. Install required packages
+We recommend the requirements as follows.
+* Python >= 3.8
+* Pytorch >= 2.1.0
+* CUDA Version >= 11.7
+```
+conda create -n opensora python=3.8 -y
+conda activate opensora
+pip install -e .
+```
+3. Install additional packages for training cases
+```
+pip install -e ".[train]"
+```
+4. Install optional requirements such as static type checking:
+```
+pip install -e '.[dev]'
+```
+
+# 🗝️ Training & Validating
+
+## 🗜️ CausalVideoVAE
+
+### Data prepare
+The organization of the training data is easy. We only need to put all the videos recursively in a directory. This makes the training more convenient when using multiple datasets.
+``` shell
+Training Dataset
+|——sub_dataset1
+    |——sub_sub_dataset1
+        |——video1.mp4
+        |——video2.mp4
+        ......
+    |——sub_sub_dataset2
+        |——video3.mp4
+        |——video4.mp4
+        ......
+|——sub_dataset2
+    |——video5.mp4
+    |——video6.mp4
+    ......
+|——video7.mp4
+|——video8.mp4
+```
+### Training
+``` shell
+bash scripts/causalvae/train.sh
+```
+We introduce the important args for training.
+| Argparse | Usage |
+|:---|:---|
+|_Training size_||
+|`--num_frames`|The number of using frames for training videos|
+|`--resolution`|The resolution of the input to the VAE|
+|`--batch_size`|The local batch size in each GPU|
+|`--sample_rate`|The frame interval of when loading training videos|
+|_Data processing_||
+|`--video_path`|/path/to/dataset|
+|_Load weights_||
+|`--model_config`|/path/to/config.json The model config of VAE. If you want to train from scratch use this parameter.|
+|`--pretrained_model_name_or_path`|A directory containing a model checkpoint and its config. Using this parameter will only load its weight but not load the state of the optimizer|
+|`--resume_from_checkpoint`|/path/to/checkpoint It will resume the training process from the checkpoint including the weight and the optimizer.|
+### Inference
+``` shell
+bash scripts/causalvae/rec_video.sh
+```
+We introduce the important args for inference.
+| Argparse | Usage |
+|:---|:---|
+|_Ouoput video size_||
+|`--num_frames`|The number of frames of generated videos|
+|`--height`|The resolution of generated videos|
+|`--width`|The resolution of generated videos|
+|_Data processing_||
+|`--video_path`|The path to the original video|
+|`--rec_path`|The path to the generated video|
+|_Load weights_||
+|`--ae_path`|/path/to/model_dir. A directory containing the checkpoint of VAE is used for inference and its model config.json|
+|_Other_||
+|`--enable_tilintg`|Use tiling to deal with videos of high resolution and long duration|
+|`--save_memory`|Save memory to inference but lightly influence quality|
+
+
+### Evaluation
+
+
+For evaluation, you should save the original video clips by using `--output_origin`.
+``` shell
+bash scripts/causalvae/prepare_eval.sh
+```
+We introduce the important args for inference.
+| Argparse | Usage |
+|:---|:---|
+|_Ouoput video size_||
+|`--num_frames`|The number of frames of generated videos|
+|`--resolution`|The resolution of generated videos|
+|_Data processing_||
+|`--real_video_dir`|The directory of the original videos.|
+|`--generated_video_dir`|The directory of the generated videos.|
+|_Load weights_||
+|`--ckpt`|/path/to/model_dir. A directory containing the checkpoint of VAE is used for inference and its model config.|
+|_Other_||
+|`--enable_tilintg`|Use tiling to deal with videos of high resolution and long time.|
+|`--output_origin`|Output the original video clips, fed into the VAE.|
+
+
+Then, we begin to eval. We introduce the important args in the script for evaluation.
+``` shell
+bash scripts/causalvae/eval.sh
+```
+| Argparse | Usage |
+|:---|:---|
+|`--metric`|The metric, such as psnr, ssim, lpips|
+|`--real_video_dir`|The directory of the original videos.|
+|`--generated_video_dir`|The directory of the generated videos.|
+
+## 📜 Text-to-Video 
+
+### Data prepare
+We use a `data.txt` file to specify all the training data. Each line in the file consists of `DATA_ROOT` and `DATA_JSON`. The example of `data.txt` is as follows.
+```
+/path/to/data_root_1,/path/to/data_json_1.json
+/path/to/data_root_2,/path/to/data_json_2.json
+...
+```
+Then, we introduce the format of the annotation json file. The absolute data path is the concatenation of `DATA_ROOT` and the `"path"` field in the annotation json file.
+#### For image
+The format of image annotation file is as follows.
+```
+[
+  {
+    "path": "00168/001680102.jpg",
+    "cap": [
+      "xxxxx."
+    ],
+    "resolution": {
+      "height": 512,
+      "width": 683
+    }
+  },
+  ...
+]
+```
+
+#### For video
+The format of video annotation file is as follows. More details refer to [HF dataset](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0).
+```
+[
+  {
+    "path": "panda70m_part_5565/qLqjjDhhD5Q/qLqjjDhhD5Q_segment_0.mp4",
+    "cap": [
+      "A man and a woman are sitting down on a news anchor talking to each other."
+    ],
+    "resolution": {
+      "height": 720,
+      "width": 1280
+    },
+    "fps": 29.97002997002997,
+    "duration": 11.444767
+  },
+  ...
+]
+```
+
+### Training
+```
+bash scripts/text_condition/gpu/train_t2v.sh
+```
+
+We introduce some key parameters in order to customize your training process.
+
+| Argparse | Usage |
+|:---|:---|
+|_Training size_||
+|`--num_frames 61`|To train videos of different durations, e.g, 29, 61, 93, 125...|
+|`--max_height 640`|To train videos of different resolutions|
+|`--max_width 480`|To train videos of different resolutions|
+|_Data processing_||
+|`--data /path/to/data.txt`|Specify your training data.|
+|`--speed_factor 1.25`|To accelerate 1.25x videos. |
+|`--drop_short_ratio 1.0`|Do not want to train on videos of dynamic durations, discard all video data with frame counts not equal to `--num_frames`|
+|`--group_frame`|If you want to train with videos of dynamic durations, we highly recommend specifying `--group_frame` as well. It improves computational efficiency during training.|
+|_Multi-stage transfer learning_||
+|`--interpolation_scale_h 1.0`|When training a base model, such as 240p (`--max_height 240`, `--interpolation_scale_h 1.0`) , and you want to initialize higher resolution models like 480p (height 480) from 240p's weights, you need to adjust `--max_height 480`, `--interpolation_scale_h 2.0`, and set `--pretrained` to your 240p weights path (path/to/240p/xxx.safetensors).|
+|`--interpolation_scale_w 1.0`|Same as `--interpolation_scale_h 1.0`|
+|_Load weights_||
+|`--pretrained`|This is typically used for loading pretrained weights across stages, such as using 240p weights to initialize 480p training. Or when switching datasets and you do not want the previous optimizer state.|
+|`--resume_from_checkpoint`|It will resume the training process from the latest checkpoint in `--output_dir`. Typically, we set `--resume_from_checkpoint="latest"`, which is useful in cases of unexpected interruptions during training.|
+|_Sequence Parallelism_||
+|`--sp_size 8 --train_sp_batch_size 2`|It means running a batch size of 2 across 8 GPUs (8 GPUs on the same node).|
+
+> [!Warning]
+> <div align="left">
+> <b>
+> 🚨 We have two ways to load weights: `--pretrained` and `--resume_from_checkpoint`. The latter will override the former.
+> </b>
+> </div>
+
+### Inference
+
+We provide multiple inference scripts to support various requirements. We recommend configuration `--guidance_scale 7.5 --num_sampling_steps 100 --sample_method EulerAncestralDiscrete` for sampling.
+
+**Inference on 93×720p**, we report speed on H100.
+
+| Size | 1 GPU | 8 GPUs (sp) | 
+|---|---|---|
+|29×720p|420s/100step|80s/100step|
+|93×720p|3400s/100step|450s/100step|
+
+#### 🖥️ 1 GPU 
+If you only have one GPU, it will perform inference on each sample sequentially, one at a time.
+```
+bash scripts/text_condition/gpu/sample_t2v.sh
+```
+
+#### 🖥️🖥️ Multi-GPUs 
+If you want to batch infer a large number of samples, each GPU will infer one sample.
+```
+bash scripts/text_condition/gpu/sample_t2v_ddp.sh
+```
+
+#### 🖥️🖥️ Multi-GPUs & Sequence Parallelism 
+If you want to quickly infer one sample, it will utilize all GPUs simultaneously to infer that sample.
+```
+bash scripts/text_condition/gpu/sample_t2v_sp.sh
+```
+
+## 🖼️ Image-to-Video
+
+### Data prepare
+Same as Text-to-Video.
+
+### Training
+```
+bash scripts/text_condition/gpu/train_inpaint.sh
+```
+
+In addition to the parameters shared with the Text-to-Video mode, there are some unique parameters specific to the Image-to-Video mode that you need to be aware of.
+
+| Argparse          | Usage                                                        |
+| :---------------- | :----------------------------------------------------------- |
+| _Training size_   |                                                              |
+| `--use_vae_preprocessed_mask` | Whether to use VAE (Variational Autoencoder) to encode the mask in order to achieve frame-level mask alignment. |
+| _Data processing_ |                                                              |
+| `--i2v_ratio 0.5` | The proportion of training data allocated to executing the Image-to-Video task. |
+| `--transition_ratio 0.4 ` | The proportion of training data allocated to executing the transition task. |
+| ` --v2v_ratio 0.1 ` | The proportion of training data allocated to executing the video continuation task. |
+| `--default_text_ratio 0.5` | When training with CFG (Classifier-Free Guidance) enabled, a portion of the text is replaced with default text, while another portion is set to an empty string. |
+| _Load weights_ |  |
+| `--pretrained_transformer_model_path ` | This parameter functions the same as the `--pretrained` parameter. |
+
+### Inference
+
+In the current version, we have only open-sourced the 480p version of the Image-to-Video (I2V) model. We recommend configuration `--guidance_scale 7.5 --num_sampling_steps 100 --sample_method PNDM` for sampling. Please note that due to the addition of frame-controllable fine-tuning, using the other samplers may not yield satisfactory results.
+
+**Inference on 93×480p**, we report speed on H100.
+
+| Size    | 1 GPU        | 8 GPUs (sp) |
+| ------- | ------------ | ----------- |
+| 93×480p | 427s/100step | 81s/100step |
+
+Before inference, you need to create two text files: one named `prompt.txt` and another named `conditional_images_path.txt`. Each line of text in `prompt.txt` should correspond to the paths on each line in `conditional_images_path.txt`. 
+
+For example, if the content of `prompt.txt` is:
+
+```
+this is a prompt of i2v task.
+this is a prompt of transition task.
+```
+
+Then the content of conditional_images_path should be:
+
+```
+/path/to/image_0.png
+/path/to/image_1_0.png,/path/to/image_1_1.png
+```
+
+This means we will execute a image-to-video task using `/path/to/image_0.png` and "this is a prompt of i2v task." For the transition task, we'll use `/path/to/image_1_0.png` and `/path/to/image_1_1.png` (note that these two paths are separated by a comma without any spaces) along with "this is a prompt of transition task."
+
+After creating the files, make sure to specify their paths in the `sample_inpaint.sh` script.
+
+#### 🖥️ 1 GPU 
+
+If you only have one GPU, it will perform inference on each sample sequentially, one at a time.
+```
+bash scripts/text_condition/gpu/sample_inpaint.sh
+```
+#### 🖥️🖥️ Multi-GPUs 
+If you want to batch infer a large number of samples, each GPU will infer one sample.
+```
+bash scripts/text_condition/gpu/sample_inpaint_ddp.sh
+```
+
+#### 🖥️🖥️ Multi-GPUs & Sequence Parallelism 
+If you want to quickly infer one sample, it will utilize all GPUs simultaneously to infer that sample.
+```
+bash scripts/text_condition/gpu/sample_inpaint_sp.sh
+```
+
+# 💡 How to Contribute
+We greatly appreciate your contributions to the Open-Sora Plan open-source community and helping us make it even better than it is now!
+
+For more details, please refer to the [Contribution Guidelines](docs/Contribution_Guidelines.md)
+
+
+
+
+# 👍 Acknowledgement
+* [Latte](https://github.com/Vchitect/Latte): It is an wonderful 2+1D video generated model.
+* [PixArt-alpha](https://github.com/PixArt-alpha/PixArt-alpha): Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis.
+* [ShareGPT4Video](https://github.com/InternLM/InternLM-XComposer/tree/main/projects/ShareGPT4Video): Improving Video Understanding and Generation with Better Captions.
+* [VideoGPT](https://github.com/wilson1yan/VideoGPT): Video Generation using VQ-VAE and Transformers.
+* [DiT](https://github.com/facebookresearch/DiT): Scalable Diffusion Models with Transformers.
+* [FiT](https://github.com/whlzy/FiT): Flexible Vision Transformer for Diffusion Model.
+* [Positional Interpolation](https://arxiv.org/abs/2306.15595): Extending Context Window of Large Language Models via Positional Interpolation.
+
+
+# 🔒 License
+* See [LICENSE](LICENSE) for details.
+
+<!--
+## ✨ Star History
+
+[![Star History](https://api.star-history.com/svg?repos=PKU-YuanGroup/Open-Sora-Plan)](https://star-history.com/#PKU-YuanGroup/Open-Sora-Plan&Date)
+-->
+
+
+# ✏️ Citing
+
+## BibTeX
+
+```bibtex
+@software{pku_yuan_lab_and_tuzhan_ai_etc_2024_10948109,
+  author       = {PKU-Yuan Lab and Tuzhan AI etc.},
+  title        = {Open-Sora-Plan},
+  month        = apr,
+  year         = 2024,
+  publisher    = {GitHub},
+  doi          = {10.5281/zenodo.10948109},
+  url          = {https://doi.org/10.5281/zenodo.10948109}
+}
+```
+## Latest DOI
+
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.10948109.svg)](https://zenodo.org/records/10948109)
+
+# 🤝 Community contributors
+
+<a href="https://github.com/PKU-YuanGroup/Open-Sora-Plan/graphs/contributors">
+  <img src="https://contrib.rocks/image?repo=PKU-YuanGroup/Open-Sora-Plan" />
+</a>
+
diff --git a/docker/LICENSE b/docker/LICENSE
new file mode 100644
index 000000000..c4bc32c80
--- /dev/null
+++ b/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.
diff --git a/docker/README.md b/docker/README.md
new file mode 100644
index 000000000..461c24dde
--- /dev/null
+++ b/docker/README.md
@@ -0,0 +1,87 @@
+# Docker4ML
+
+Useful docker scripts for ML development.
+[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 Environment
+
+```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 [requirements.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>.
diff --git a/docker/build_docker.png b/docker/build_docker.png
new file mode 100755
index 0000000000000000000000000000000000000000..a5dbf2c1426ecad3e2a18e18f59319c0de86f766
GIT binary patch
literal 64398
zcmb@tXIK;M7A}m6A|fK7ptOiIfdGmK(nX{<rFW2yG?CsC1r-4W=_Nprj?$~t7^?K%
zA(37}2tAOH<iyw4z0dysoIl@ikubR?5A&>kuX|0zQ*9O6ODvZtC@5&vRG;WlP+Tyk
zpg5C$@dEji8sU=X<hL_kx+;$;N(b0B$S0JJ3R(&j6csTvr_axm&#7Okns`xAF!B8U
zJ)@?}y+c8P=Tdv3pzm+FMR?}H+8f9h2m`d4nkUtCs2({0CDV*qsN<UbUDdAHaRl|V
z81GixPrnF4-#p(_8#?^x=gZ4Nf>V`W&|KC@R3AQGw$r?oUiBraAF^}X`n56Zhj*WV
zoIhP|8P(7FoiGf}jvI_y%#(k{?V_#`V8_pesx<mHH#d1tWGoCwP>p0^dh$K~=SWcu
z{e91Wj~@&Sz`uY0eoy*redhl<FCSy1PD<_)deL-Z*oc12FZac-gek<I_z5IUzAC*c
z+52=U=3tk=*f0JX;j_PTFXR<&JEp|%Ah%HD+@E#k#}6rcCx%sU!za5_fy|eJAIb|%
zqq;;ACj;jt%)AeVPdHC)qY(1^o4x7D7*fi6;s&$WsSa5A9Vu%Xp%IyLFm84FvGlY}
zq;P$#{#d8XdA*GExy5uWp~i=EW<i*?p^M#n1^mgmXO+j6U07EhrL<x-I!1(zY!P8Y
zKfgN!C0~0%S{y6FZjRr7i4acmKtAg|ksA*7(u2+Giv){EVx@ML$6^4;QAPC8mB8an
zxk{`Qs0fL08CWoHl5s%tr^Xx*zm)bV24X-MePif1FPSdp-A37e_V+0AfU<WoOsxEv
z>R=q&lGBPMiCnKG0R&u6=~#~CTYT!ANu5Iy9%<WjB}v;wBT+c6E!MU0O+NB(fo<T)
zCkMPF|7i_?8$dLmiz(#Se!74ORR}>spDoadfo&ITs@poW%DpQ?U^s)F=K4vo13&qk
zGFVJeS3fmwGnth5p1AOUHVSr=7%ra!ho#Ua1>wLWli+<>!jpLnQYNpAm~|ECSNc_~
zx;u1%q<$YlGOKD1x>~4{E$OkOmG*L0&?yjO<rLhkX4E;jh8PW6LnB9r-p}y9XeQEp
z+WI+3GC6^0h<#rr^*=$wK1`D0Kt|h+OCsl1YD#5ulQFDgHqb7hfAR4#@e9#o<Q6=8
z8lj+C{v<cXTH#+)%bJa21WtC&+o>;%iyRV{N%<PEeLVs)-vm*>i&vNoEF%s3zX89!
zJ{dCyXooi3<o<q(i>3uP4&8r^l|MDY;T&^L4>y4!LUF=v;AXZHRzYUBw;g~}P76;)
zinR>uXH#Qe!qDGz%8x;N-x}etfXzug$DJTc_g~-;43AYdac;w}M{aO*899Q9dc9k-
z2OCb_VBj`pSRXLc4DG<0;pR`kfO(_v!<hiq6cVl0g&BpC<hZ+@<}jn{Ww88V3(}z`
zHPcDfIehe&H<VGd+{QMCb6EXY*S}l*W!}MO{613uIKB?o9=Z69bn4Y0cFahSV{l~;
zc?Gfy$)y{#%^zv_CEsevQNRWm%EN$pm9L%hU+?+iwqk!TdvUOED59uUI5H~9dgGQn
zeD6NVC&?Ah8vEAkGV$~hxVcF+43Y2(Il2PXIR@`z@{4wQ$%dd5`tDO44u>I1TM#?M
zXy>k#5w?n;=_inn@Coq%osCIw-aJA}wY+^aQ-TM+$Vd5AKgDGvB^h`zdgP~~PDY&|
zSdGVXi*dFgFyd`D0@g5E(*HBJn=B};Iw00Vl-t-)ey2!9{sCerKGo%5Sr(7I*GzW3
zP|7HWASzDy2roA|SpGg5IK|<jPfuXg)QKN3BrDvDUAk)ghnic;Wn#Vuvo|p+#<5Hj
zNWC`5*#*X$XutK@p~B?FR)J>87b7YwMdLaEh<9}kU1n*EWOlx2Nozq2Ev`Jz%89g~
zUFeER@#U}_O*$+F{dyCcYb$m*dJnf*MocJuo3GGt8Z{KWr*3V2ojMLoy1$E)&_BKo
zZU7$xh=SLNZFfR=NbDT3KQ@19cPwMDd}&1~7l@PKwONxP98#!ZRJ2=)w2UF9wBY<e
zCXx@P5@ACzp#rD0NRY{e9(*<X*~{^)=j6faMnr^1IP{bN&MO0#ig@TvBpOnU9@XS(
zmEV=Huk4ou*sxT-f#DEANVwHys$SMO!nb1`aPfRg)96nRZqp4TG2ze-VmzMoPPz8P
z1@u|2Ph^Q;v7qC$J}o~An`VSNcERQe%((eP!Y^c?8U@obQ9{4)^gT%Dw7)M?j|d0v
z_iUa!BxRPmo*FiBu*AYQ%W*LWL5%|~bLil}Vi!6~ygeHNgo0(a9+UE24v{0(cSv|M
zu!fOx7*{|`+9}T`cOuRFW%k+QsucclWvfx0$8PsbNVK#{`RlJrPkPFngS*Nwu=x|m
z2Gf)Hr}bCS0bWjwz);b}6|50KTwKm04qjmbl^FFcUzZ;FmF7Q)<_$q7#7t;#Yu19%
z2j^<h65FWL8|fkLA8dEy_*bc!THHJMa69HFete-_%{$H9wO0T|C(i|q#m^B#!ewtQ
zPPSclhb^8hfzfEgJ2&y55YBS>!J6=D#urgBpRm;}Su$@K;;m0QB^~_IigmPR=t?{s
zj|iIFSEoQG%}9x3{N3V#!TZ_^j-!C4${rVC3tethwUq-k(~_5K5i?HYAH#ekvCarg
zF8j`<S51BWQyCFfwmveMoP&?3C?A}(R5O2RczO{(^&$dbr2vhdNZd;U<V3DQe#NZ$
z_#J#Kbv>Y-^r^TA9^+f1vKa}juhbrptJ!y86&M(qRKF{)sU5alsR=w?1`y|H8c;;E
z?6g&jMl$2(-BZ7qrjVeB^Nv=;K;}poDW<8|$cRG|aBKzM)PAs1`U%c(2TTgrKNggG
z1=sp<;3EAtk4J*d)wJdj;o5?8wQEn{$fos7OOg#Ea1(deJ;MVRvH&3J6&6d1z!wi{
z)_o5%N3#8T7trXFKm$ssA5nY<cfH<H$(7<nq-gm<Dp(^YegGClymKaFBXE>t#zTF#
zmOir~0zV&+DK>O|6Xt6fuW)S5@ElsoJGS7rW-n<E3xbWekhX;BD!tSnS=jQ*>^Zi|
zw|sERV>fPnPN+jvyw&2R#4HV%*@P-(?5R(*tv20z-{s(k0M5X(G2d4~LOl-;Q;sFP
zoh{?7`6EK44guyTj*hITQSg&C;B+g)ODr*#r1hRMN+LAxB-*en!lj;(PWqq~8<!v@
zhON^1)~$k2#%}0_VLJE@TJre0o=l8uq*ZEZ<vTF|L9g-@JX&aZjrw5F2uh75+2+xS
zMm4<2_O3WVApz|h+HuTgw)u`o^N-qxgm9D^y8_pVC@ZIUep79eBZ-+1=bZu#vx99j
zf^YIMeP{5CmV!f%lb99AH%!q=3nEwIRRds6Wu?&}DjFeE;Nf`6#|&*+(8(fA%@v&A
zO~WeUN9o=8?Hx0la_Rc%-T*>h2_*u|R1WqcNVbITUM|b<7=V)jKj7%}93H+QuLSAF
zg1659q#}hI=e(uX3A>hbjhi{*e2ZJCFhS-3Lpp*!(S1$T0K<*DR1gN8cbILg<k|De
zKpQbVpMCZjXi7=jCEpff9v*w)aPrb%l031GfW1qTLORvSQKMM%sTZju5@xQr0=`x+
zzt?k2EwA5$J%QT!91EU+kfW12_Y(~A#ue=K#Fbn_^Ke5!F<)S>qI15l^O+NQZSQtM
zLbsU@;;pn@dS&K3cw?+nkFQ%WM|taDPPsnyrIc3e%5S9k<49P;Osd)n>+h_O5U2LP
ze2^q21qbnY{c)Y2Tco5?`KvisE)8z<qyh8+K|3m^K0K-en(h&l??eS`eZyidQ&Vk@
z*NEMw{;=n9irtG|V`aIq#wA=TFiJYXui!Z{Y<=+WFU63iXS<Fg3o&`H+;4~E(BzOX
ze9eeRLmPE|bs9<1R%<Ec1rn#H%ZgPJFWraby0C&o-?P&|;vy3>3M1&nwIQ@}@ji^*
zg3l1lU<A3^NIrds@H4UL+-ovngY*Yfy`YL}S>jD-_gScUCk(%ht)I^gJtjqr!}3rA
z;Q6COa6Bixm4Ra6QQ+N^*k^RXBF|7ewy1L@BKv#2uVCbYw!{?i)KeoE`8?ty-+cs1
z0v4c1d(^-X3|`l+=S5gm!}DD36Dum-c7)DhMbQ$|F2()nwz0@O<CA1S;R&Z4Irkbx
zZ>&hX0;kVo?+>?>>Wl@$w0kGhC!=VpxfrhW6H;O&)Le04cU+RJE2c};g<ZU&(cUo9
z*?+9rTaTFmh>ZN9;$5x!`k4M`_l}f3g^u|;sPKIRAmk0;)3*b7i0BR%dza=Y=Puvg
zflk@J6gOYQvT8={YMLPfO&gfhrGGq#+o-;+$w0HWLE3>e!w4qUf%(u=QL#_TyR6_R
z;TOH}5%1&+;F!cBNP^u3R=3+>8^C9rDh%bZ0(%&fTe8CN5gqieX5l<;ppA;f9kp=I
zB(_Q*>u23?3P>ft&7?KGUUz-OSNWHgdWiTKr(Nvxc8hqWQ-b{I$CgWFB-`6p50ctP
zaqc&7UN%&*dudg#gQ{{P-Ix}?1)WLGVORzAhy4nIt&)z;&JTdfPDf~!tM>F<NcYPD
z4{rYAaY?v!@=<hCA;VX*#pBjfs|dWbf?Bzz2kAXXlu=#o&RP87Oyo2cg*v}!>kr3)
z^IW`6AC%3ry@|a>Z-9%?9C{M1#9KDzBOk1`CL~1Gi2l|f<fLLcU847QPD@_71E?`K
z4ONQX^NbErPl^REp}TIwZ_7|?0Xp)-A|aQ?RLPdHdv=U$58QX9d^4h_?!<mnct!lx
ztAZ<MK`@nKyl}swm7PZ*Pe;N6gq7(6xDxC~+E4J-ArBHLdnL?N7BSII8vAm-$U(ca
zS!x(V0_AIs%uy^qy3}?Q@MPMZ*SK%6qGTIu^^ewHyN~c0x3p|&<MYT7a7v4w7ev9r
zYwQh=T-4^OEICFqHG5RIkH&L;0l0SXVcJLE_3^!UAU6bATi{!;zkSEV<}S~|jgQy)
znrhU8$-v2jT;~wwII}XY9VaP(Zu--aX7RJo^GElOX)1nvUqr#yW|xt;h<8JG-afT?
z17KwB*rm1bwOWgs1|Mq3slVkjJ9eUzQ!)nEDq6F~>s$G}DCQ=os<?HAYY{#d@7+Pw
zDH)f$ol-_25@A=nvy{M3q7E-6E$S5QOuiaE#cb?>duu^7G{gqAcTzN0NI?xv0xBG*
zG@mp<ipdV3e&$VI6<{G4k8CQypF?@7BF-XZ?kx!wZoiDD7F}o!9Z(;VF6a0l$a1W$
zakg+QmRB@@0j7!ve?fXq4qe0%N{~Z8gI@~Kfnf{n@%b%%$5tahm0!{_C8mS=JovQs
zxE`gJ3aW9xhQ2~D%!()UmCJanb~$(<fV1#wjK5y$I|io)%BY#-3~*2v{~(x(0a&^Q
zn8TdkOnB1p`g?3CSj5xr${3#-&XVa$OB9$&`PZE%gr>Lej{Q&>3O2nOox#hb@sH*U
zt}ti~<k5%cqX1E;1Y|!=e`1ww7UfR*1H9Kur<As*TMu2&1X(=h`4Je4&jkcTz6F#@
zAm1v#m(d4LIj@7tuQNOOtjV%VBuX}&lV7!$jo;YhQ`($<{{}<)EKKUh#vaOumgK}f
zsMztc2Rxo=KTDDpzpYSO2JZzQPJ(WpkPx@A(IA)bAL6`4Ey;GG0ynKLBJzuI9MSDk
zZ*L^HiNIUqoXCpNPi4}U6^A@9mP$1D{Q2mp-=O|OW5eg~B5WVA1@<^caG>CLVfw2L
z32!AzxZ@^C6SmeQUQ>&q$}o9tLLEqQ=lT9Hc#Q!_dOSICy3x2e$#?Sp(!Q2}8$rV^
z<jSE310(SE8%|X3jC$QfxH_+KX>sr5NZeI0|3tZ(com1tz}SA$RBg{Xo47#7UQ;UI
zcd&M)YHP{i^^YS&jU;z`;8{Z!w;&$y)YYY0FF@M>8-z9uee)@LmpAK~ixNYJ7@lbJ
zvtx@M^6bNbc4p?GD_Vu=w>Fuz%Wr7GS3X$>ER|dz!DX@1JK_mjr;NZ@;^&j5>Ga9m
zO7XRaIi-jHdfIlRR=GLLX~ghjZ1H8mZ}Tdb$ujYlH}_-N!Vqe-mh~?~Sv=MwL`hY<
zG0JNy3pgsG_ctjW>9PW;f(CZWGIEO`0u<yRk5^<lQ0Ci-1;mP~z43swhl-lsl6sH?
zJO1976uECIO`TNcj-#2TBbsYBUiS*t1XTZarMU7P<=s;YoHwty{W;yutJdU{y~P4#
z5k2W-xEm%fp_)J$C#VezG95~G=>-fYI+7lf3@P96O|<D*$YPZPH9IbRu>Xa(a$4p9
z^uHn?TsXXu_y@G|t7lv#exiA7z4f$C;OjR}EN?Obxt8F|!x+PniL{gVEoo#hDDk)z
zpg<3?<|p)M-L7zRzr8U69m4#4(uR;9*?GH{DY~=$nK~)diJ53RB4*kHYKvO#5`E}F
z?I!*7L?t$zp34yUpD_CK0i0np|CLyBI){jD=t@MR<p<IuDw6pmxK4Ha98zPC`l-kV
zAKh%EUz5+%CPHLUI>>b>dbKe&{r1*zp7CePkGCSVeilA<Dfd@P+Q~lI77qzI9FT`I
zY|~Y``g0FTXyd;}e4lFY5rErxMN9uO$NyX8`;qZ?3i5>iI=Q_t!E0zfAa1CPucKa`
z;YP5{OD(Dxdy&j=tGuKymQ5!m*f1>uJ?}YQfP8Tvc#vr_7<pA^xJZ7Aqa{aq{{upX
zl;kjVyP&dZpUa5wd~oiXp3IJt{sSExpXt$Lni|AUrW>oX(IWyFF(N5Y-$piV#H1V-
z$mfL${}?K258d|wzYzBE#NOkvCaig%I0hq+sFRk9ir<q=HSY}<ps!AK=JEiEi-S!t
z*q%|9*sYk8f;O=oC`pG5<TD)vzo~x<>&bV$1flHBgA(}cuZ?28Pu6-&MM>^ymtbV9
z|Fhs@ccT3o&QYO&97^1@h(d%G)?3kLPcnATA(woZidVZM%I1B9KpHTkwKHCFdvLl_
zDrOB5Bd$CDaxuAP>yLQGkje-gaOr^MA`cmO2V6D}bXH(_yN<AIcm$~zjP7w^Njwv5
zzx)n{1m{A{>G3xRP<4|D@hu;F?yl`J=OH0h?YU!WrUwdbGeBoff3Q#@I1gz?fa-rf
zymJ~<w>2K<WGWHApYLMk2(~f;aS<mfxUen(2w-&oYLiJ#m=bb$FU}_0a;qlN*o^mq
zGf%u{K6O&~LFsT7Rs{MT$v%gn8t{9~H|elG>qFni#ASK_hGtQ<Brjjjvb-KapL8|w
z=jKbd)Cr^nVWU5R;9ie{J}f)WRBzN%)iJPw&CJ&P`OwYt^nGUUUbM#~l27;*xajyT
zBC6;|s~S@y#_@#LSQFRAQ-6n~Qqa%jKg}oIb?VQEnkb@-nwK<`4d9DWfpKVy9Q@71
z>e}UZt>{L6kamTxixa7S_xJ`*z(*{z`2>hqug&Ji32D;_H^=svyUk3^;M<we-gDkP
zf$~><#sIE+R813nYbwyRgMbv7LBA0?1<#Mr%7H^+{xbm?oc56g$3-c@?oW}T6)!jB
zj}^ShpFIh8wowrNv3&s!BiXXNsw4Q}aBIEZCfr(}Jqc3mh2^u$Av0ItmAg%LAAN;b
z68At%_HLWj239nN)J$DJtWIyzh9itt#<9tw^$a3uYz%M?tm&jku4{q1q78d=Y6AT^
zV(9u6lL<4(yXS9SJwur7KgM9k=^4=6##5-xAhZD2(*$}RDT&05+7ac4g@=o3JRy>j
zhl$D7&sPAX`&xqbXPcb!ZxGV%$@zGHS_sgaF?4jm0yg+~<m;dNAlUM|9DES2Go0a+
zQAU$M>ZBV@@`aP7neI>C##$a*#u>d141lJ^VtuQli|7#;36Tzc7pdYF_}g%a$B$KL
zTf!NEeG-EkVVadqdV`RDY*jELFc416;!SH^F<m&)niE-y`B^DIE|O)l(UkK}8JX*_
ze`+QQi-b(;ZfYoQfh)wn^pwItNWOqTNBbg#K}S6j@S>zs8nYokzR7B6n1qj&t1Q?s
zG}>T?FmN7(p%;3wV@~I>y*(Hi{es=AQ({z@Q2Y`1>D7a1kzN_h*Tp-ur;Gd+CwO_I
z%qg&zA{bk6a5xyIurS9!)3*q`PQ(p9x&8LJ&+*L=0y2~<MU#={)jj+dm@cC!HLjXH
z)PAYEQn)~<1s5gnJF{4mY<+{=v@H+>6T%M+tV<a7E5&`vOUpQCDWgzyOdF^lC&%Xr
z?Px~g2PVRdutLMI1ReGyzM1_dJp1G0?}mnF(JU_*TIn1FQT}fy7jMiLAgP5Y!Bs?7
z7Je$`6c1_m;_*y?1+Rp|RN(Dx6!=Lp;54n%4}{W+=|t_sh?z#1EfcLakncC~(q7<e
zcY2lgFU#T`Jt%+nAv`v);vl849MUXJ^*3s!)o<NaBjZC-cVVzY7Zw(}3&AECR7Rks
z2@_u7{CR?evy*O7_(njKTKU`@xqXCtR65FN6NsRUx2=-)`1oAf`vLqmzR@{!g~1}|
zr*h<5=HxZ#zK0?0#My{Z8l=Xn%TKM$Lb}eOxAEZIPP~uloz5*RWbHtBtg`}{!#2sW
zJ@jlpN7r)Hy@8TuiN{;+nU|GX<qfcWFF)L78W5?2c(`Vq<Tdkc7a_!_ViMrfD+pgK
zm+E@Fd67xwzaG3{8rl&QH*j+AFgJt0I`HW;&GW+*MelIKgQNtiQ__OtK_Xk;*`WgA
z-r*Yvv1eAi@@KN5bJr~%dN9wOC1pOZ{vq|p!YCeHnAa$O$Prnvh!ddagjSz=Yl`@|
zgi;u#Uyqt1g}x_k%a0ftR0v-%11#g^bIC5ew2q`Oe_}+Guwd4ivVtYoJ_-gmKBZ<#
zuZz+FmX~HsT>N+GtaA-%*3XaiP2Kt;5E@POt%_bAVVQMDk<lGu&nRAUnpe{UDR%NM
zmOhAEUBZW_h5Y&6kK{rWSx+mTQPF-l>q{|<jDg4(!*~M%+j16sV(1uxY4$OAPl_m^
z+XEC`wbY=b5SAf#ZnGpf;eP-o#q`0G|6zuH-TvF#{`WupKbWGolz;jzf3AL=x6N&A
z$uSxuPrTuMs%tedr6F|%v)}a(lj4f`zER!5XE-?_9jH5&+-)W@b|%}Bqpomw)eK!*
zD)_%F)1L(&K@~&udda=>@k7pLQXxnD(o#_14gTaO)mnYXyvZ{Ivt*75n_i8O9d`H7
zAd5jxa;>o<2kE}Z>IiU$xenv`ltOO)S&RH?^eVVAhohw+W~+H7TEfrFadr!`hAAAz
zaWDROyuFsgSb;`Nj8pdRYV_A$Sx^fTXS~qx<nz*K7cV<nW{mVH?_ikggwI&G?_!b*
zo%32@x5sBPqJPnpbNFt6@|f8WNw!!n4;s!vIhHkugci2h3}YK7nC5#t<?k_fC71+V
zcM{Y`_`;A636VT6<k&}hJ;Y{9oN*26gr^_t#bqt6jE+ur&y@sXuFu8zWq!g6USoD5
zScI8E6X^JS!Twe1G>97R&8vR%k=WB?ON3tgAxoD_c*+rOEW0yYqI%{F>R-K5WQC}f
zPoc~q*P-C2HlWD4!SE2*poJCs8aW?6VaRy~pqg%INk74!fL!6_dCH1n^RNyG_^zYH
z)JK<bj+tf$Id=M~XtC}fpS}liX${0FBHh)EXgsupZZ*^kD@Qxs;E^)i2%8Om-2MI|
zUZ@mMPhF(Le)1zy?S?9v*O={7281Y=K(JFU=dMX@(@DuSuUnLrPO{#$z9jLlI@YiM
zxx-!K6TTNMLE85x5A63GnPx%9nS)zllp1j_pWfGcAC(<Xo%BeytbZt8f>y5QOT7-w
zi<7xjpVmS-)$(i5A-rK!G*YI(-jnhB$ujMOCnI_Fu&V@CMU=|z`n?JTjFsFM4|iWo
zH%X{Y7_iai6~a!YHdy-#I#~lnz4}8+{P)dFNkfZe05>I~vccRUb?hp{`d4!Xp+!c!
zw5t9iL3)>XEB>m`i<)#&En>K`L3YEEOvuDa;C1@3)x;9$_j>XDx0WEMr%M(P>ZCE>
zAq*D6C$gcacd+3bAM@Ar*vb4+vvy7D2hI?wY#;Q|z6}#zF|?yyNxS^}lQly|V1>O8
zyGm2%5$8>dJMS}Iu6;-YJm3<ooPf)UPb5vIhX5KCtB;Yv*%QM_qSxbo<ZA^Yvtm5C
zv_D3V_i5eBFkEF_-IIP%0gGY1mtBE-UsDE_fy#-mw;56~<yc+KF|wId4rm=r^}0V7
zIp&j!sGw|(={gYI2^`W^#A4mH{R}SOb9K~WPBtM$-&HSX8---N-$&|rSC?{!#%(L{
zW7$D5_v&Nv*pdqs3mAaH5}e(rN$tJAB=aqrs^i=?2xwlU&BhHgX}fsyRb)_Juw5iR
zC1rq|<wdEXY1IHzy_CsG<XXYRJ+#a+9Vh-t-+Q!Cr@T@C7LHCyeUjvN+=g?$|Kmno
zs?mf<PFa~9lb4`yF}bZW`y#ogJTi>cof)s0#D?8s;x>+R+T4)K0JTFm=``sb#w*%g
zwo-Bvt_<W^T@Dsga6;-%NPV!yAICb<wVnRBn`}(gxo+J5ld`w#`=Pona4YIib>eWg
zt!Y@89&PojN!Zm}&51>DP|7NuI7$^*I!!!1_s#4vlsV#a48Txvc4n?@GC5uRuBPrr
z9@aysjj~r`evzn#%}&5a_vMqOaGxuQT0NfKsBiJqOySAFuF|`I`F87phjK+Sx4tr(
zK#fbOiuqYTAf0Ta{_SdadQQX~+|yyYF9Ry-+<4NLziTP;nQbDMU=N0+f7H4s6v#Y1
zklN)<ixQl@^$x(Ohq`$^g|<uJiu-qqt@}3_fGxRM`P#{hz-NFBe?BoA&Y0-yvF3)_
zq-DgN%s&4z<moEygDI2=Dy#M!4X?<)>w0CnK*?crOQQXX)BBjN>ixG?L!0_~LLM`@
z1OB6Shd5P$;o<D)4rnBcuonC#2=g!_L_EfKefCl7qq!&3s&w~0G2LI%?mqZ}rb(ad
zs2>J>LxMI&&Y~N;P}NJ}J~%^7`&C{-)5m~d=-Oz^hWh~;oxGryqI0lN4{E$QOgE2>
zifpz@;WbVN@8^Da1tiO`8XvvS^N4WQO%hNlq6DQ=mlN*xre%Ve_xK2+JiI;rjwU!k
zwSq-=cD|;AEG#nLbW|`cw<I?okMl8w$KShQNL-<M0s25DgS!)#eoD2yBGV>jDsR3I
zFxC!91<962SFlu_Jn6@xg89Q1&!Rt2M!56FhU2xcvWv>)Q`M=SA@;1fJGM4-8x~bc
zDju1mb?S3^jrz(1o-fJvb(_Di`-bo9+_xGn)J(k9ue}G@qcnsCV<gjDuQG%GF|}oL
zRk(9?`%kJv%#-OPsfO_S+BT1i)N6ec@z~qW@ty~l+Gc@oy!wb~>^5)w3qHlPDK5<%
zl?-3ALTHH@YC>V%v-0jjiGBF%%j=LM_uy=RG6JH}Y~_Scnk#Ucd@0R6AH{%0(a?3<
zDLxL<>*N6>|2O%Xd#~<xuN}bIKaJ!{e95KSH!o)1>OaVCm((pdF|v_Dy0S;Wq_c(k
z@jKHm4!AUq&vk6J4K3b~d$2M8!bwV)-ga+EeAU&a!S^sx*G4r$%yT7u>6dDGit42o
zqI<Vhy|k&BYBF8LI`<{Inn0Ur*CqXc!bQbX`LshW!f#(ENKLUpE{{n|lqBEY@0;1m
zyU}n&j<qY8D!iec&giPJCKn?RYN}Gb->#F=?@zN%<w52MZTgue|7~1;vD4*MTN7FT
zCmS!Ccup2DJwU*6{aiz?qqX!ADRC+|yXi7MbeeeF2;CAl@8fKAk9<~#YG!UA-?qf&
zYJBoR3mh-G#7FEaEmIBUsbSYW!nczapO0t$jlikg9d)g{S9U#o?A{8niO`eWo9`P_
zO_SSq*Y24Y$fYeul74uJ%DdaKR8>45b*TME5RZ3fg#wXI8-P|rlDK$g(vh6Q;iJ%w
zC%RAdx<rS%mlLgHmlISGGwlD$<7nf<1lR0{Ffl>Om;=)ISt<xtJ_Sk5Fu3<>Wpr2D
z`%Q>O6I9Ger3z60UTOI|nFMIq&x{`D@0gJjA5OYJ%_LiKqj~$Zw?^e`j;AKimJ$7|
zp`kABIa@EeD`20=?KPRL*1&B49dmcm5E;dt-b6~Vu*%Rb2f~&nBismoJ*ASq;!@6W
z?y=uKxhr-xI{5<_mH#Vp$41ISwbX57MD9RnNxA#<d#rO`=X)*4^n{UmG-o}>SGq)(
zLa@tCL2={&C)|c?^F^a=)Jz`z)X&7zl?baRRQE_0L1e=h5S$!#8i7FIZU~t$6O^^r
z@Q0|~(JrQt1Evhz696klkLj#LR-fbXJ%_J?W;$~;Nf}*E^LDSLDxgw6aj4#4Dbva1
zd+m7F9@W?zdu*u9mtF-6981GVA$UHtQAZioT3}E%(!F!l^yw72me>yewhynViGL~5
zlVbhcE-n@`xrN(w3hgjE5Sk63`d7i_j*QT28g%7UP(}Njl1fv%8uXL?#5bf6^T*29
zBRBljnEFM}$#5{2lv#36Gc~4^T7={sq_D=I39A3}_$n_%9i=D%1<M{jWB~HEtKD_a
zwNx)Je;4m7ruQ~L&o*R-f!b0Mdw3+plx*&6Fy}w@Zs)U5x%cm8#J*C(%vZ&&lSpEQ
z#PxC9Ss~yZm$7Qpj6kzX4`ktM%g0W}lWX8UnWzS(lT?q%4qAogZuwx9PpkCJ0Ystx
z@}ukaj_HtaH_})vD#hPRBi|FZ*c$IMBfq+4As?yy@ewufXJV@pOkkG-7%@EY;WorY
zA(j(<doU;`_U<09Jq*J_8%Wxt7P<qF>VP1Rqaq~%zQLU0Rj6-_)J#MoppZWvQsU?9
zM*Xj_hAu#kmy;SZFGg7vM$HMXj8C~d<Fy0j#&mhC*8?%fzJ!p)+t}nto%QAvQ^9n*
z1v=g;sdA;IQJd6h24L?%Q2A1X+2)V@{;f^HyskNa;JmY$_)?U=n{vh#E5a-}c8Dr(
z9!0GTnM%lBq)y__&J2X6;{-E|>wkusj$QIgZzqVV-<`_e+LL?WV{h6uw*H7+lL3f0
zpVc;T;GgV9%amHL$v`fyHagtb?3knlyrb>q$879r8M!>$#|(gXIHVJe%-Q80IFr=8
zb02#~A$LqMx)Tu<a~`h?{pQZ%R{Pv5ORI?}wBpwNbZgHhPw16jO)bxxk^XLFaE7i%
zy8JOt;iA9r>oYUtI9K9axb_|4wH;;e^+Hcu3*ZSEiz@J%Ad1wI!n;vvl#7xk48Rws
z%CLvc1`w0+NfM-eDq+b%cUE*0xq6){d2vS3sFuq7Bc;ZL6CUHK)>X{LN}JPB@lPgY
zBdPD{OuZ!<6Vdk5)MX_tOU#*#Xs!o{veC!Syy)!<!_-1obUhR0K43`)YS8ab?|;d!
z<T1XiTz=Ozt;~=11;505odMdd{ug32r#mb7UypaR!N*fF+3Z^2pP_s~axV&iUgsG&
zq3MC-iVw0@UHq}|S-qY?(KPX465b%dN89)jQWMvYh@SvAhy0|gSf@4Sg&GIoDWkkf
zWP{Z(F_0d8|H~LU55w`3M|#m3O^%KXz%4?|eC&fa*uMevAz$aN+UjqBUR@rjTa2D*
zW0!gX((kD<MmPQ%p(KO!9L*xcoRJM*+Ef<bQzW_jqjB&O8lxP+3U4*8A!pdUgsFsb
z9Pf7>jwPF_OG+g1<cHrEA#2(}ls4tN_KmdEOzU^Z?bv2&VQJ=~pnZ+{6qu8!BC2{g
zzBYgMRbf;;s@Ct(v&+<)xYvJgih^QYMYa49RW~1Ings~=QFjOJzx7>C<G9$#7kEc`
zz8X?(?6+3&@n+x@wJ2(E!l}a4_`itt*COggC{GqTtI^$*o!fX|QWPzSGIgqWAp25V
zx7+7=D~-qDWe-FTQ?rg)&fkKLf+CBKjGr@n(45V|gR3Z-{@5<2@1$}?I@*Z?itT^E
z_*jihpz1Dlvq#N(vxC6H|J?0s5e2{h9ld0}@xSN)-%-r}=J?->yyp2kc=cyVAURIk
zv?yS)@>Fh3BgsCO6Ik^<d-PFMGw9Fbf@xQv;RHPv*2%<o`>$KvMask*o$%)$(*5n3
zQGDj(Hr6!eHC7h4lcXfY%Vt-&<!Pm&6}2S!V<@BC*azYtB8m^`3BG}rRV9ie_&~my
zx_;je<I@BthbvC`Amv{EuCz(QLN;1LCx30xrWe&PbVEK@daId14m)7BR$9k#eEDH`
z(1U*kuzZf1i7^QtBa@jT)CbTDs6(E;i#}iL|EqXDON0T4R>h#KR736eK5T#Oy5oE_
zkB;T7#m3P{oc*zAZtiOJGBht|lJ!6G8nIqUamPG_<#c!~=LN=>4kpytmt6ZQcABN&
zcQgeO8BlEE#qq**QzWy^%p}iK1l>BmXDp50r}%jBl?i37RJ`@G=}AtjW~g;Vd`(D@
zJiHaZcyk`qls}Kqz!5B8#SSP1waig!fY!ciK+mX_|JlSNx*L|VhYW%WjsqSOAJvjt
zxmZlxMjl&(Rm%@@;az<0RE67YE_HS1sgt~3Q)*me0GgNa8Mk|)S6*Ve6~IUh;fZ>u
z!Q?dS?tM1Djr$bTOn=s0ukGfFrIh@_>SaJQ{i5Dsz~KdhLC}83D~XI)2@#Evxvzvr
z5`I%8F8D~5Fw&(#@68<2?U{3Sp-_CpvmHfBhq)_U?dP562o=!>dMw3ag<DrX*$fn(
z7IO#`Z0UTXpR41NO8BN=l!3jJ#0Zy^u&z~YtKk*;5UB5$CYFDNAX-u;hD%T%H1(bo
zTpN2{pd=v}lWi}1^rnQ|)Jd;3brx{@5#8@xR^Q=d((C(c4c#a0VX%goBMq`_?c>~H
zFL{NFp$@6nno3~N{NI{ZM`|LJtmy7acd>g=!5v@@lb`KgnKQXL_%z?#p2S(&1Vz)9
z(aZ)#Jc##t?uqfb$p%q<R9+w%Ie0wEF+P9c)3qq~J!lD|u02%4ZSGqizN|wwl3ynK
z+ZX&;>nbO0i~#h@Ju*S%C6_#fN*gq<a`|~=ewV!kaaQxlEeL$~*F2bZd;H_Lqekj`
zj<=hrYV@Z)zqCq8Hr7A#Hoy3)_Cu=rah-o+>Euamk>sI|VjjTBL8J5LGUhPXzjPim
zWjRzu%>>eoU)Q082=viXO3oJa+zc0f7XYX07mrbTq_XAlTt$9sQ(KEt?Mc1ei>KZr
zG}}tO8N=QIy8MFY%)-T1ExDfy2_!x-8O1e-p|et&nf<i;&YFGQ*C`)kh(T+Y<R{cP
zhl_qH#=jZ87$r`9bq%d;VV9<uG-3Iq$5#R#rzdTUp4?%HJF|8T6l7X$(l`lmhfQD^
zAYQJX7RKxlPSkJd-Io=A7A}2gXgz$m=Mc|#>#hA^+I@#nu-o>t=<HFZMqK3^yGY*O
z(<o`)CAqvtua4Pz_5*2D8&#&+Am=?TZp?r4>B_$d?qr%As9wpYb6oaz{UfHAXNrd@
zqv+l`Mp<(iPqlqs#6HjeuJ!s<ixn{VfW$s4jT4413N+`LNH}>bS7D_-&ojVdmMNnI
zcvN#JrN13=2}JJ|C{ABwlXzIiZlTfXWD2dfmH1^aFP9%nee-f^F8sunBYj3vw<qdH
zQWWiQ+?jNxh6cFw9{T!jdzC>*hAmR+3xsaMsr(K!&}Vb!5hV*s{*JwY^P7^weeLp}
zUQg$JXR^3w>h5YN?oAoreW{;R=M=tI;osV<mKR}yIee2n8e?KY7VXo4*^!AaMr&`>
zZWg;L^(f<wW&$cCJ^1bSUr<2^X_kr;XRHj=rR~1UzjRn@1x1!G6HG0bQ)7e_CMcr>
zPY{m5uIlVenI*${5sDu4M|q!&^k3p;gpM%4F2`Xodi$AyKtyim$V4m>X>xMH)o=)u
zBeU&-+pW8eX(AE<!B<uSn$PDZ)JTcbh+|(Esw@o}54SXGhcs6XfcPwjY1h*6D{KMm
zf6bF(5u1aByo!C#A=WhYUM1wlea78JY4J>#UFq8JgYj$6&vZy>p?BiJjy3)7UfPBY
zdg_LX*?=|7Ufg;cprGg4L)VW?rBn)dk9k`fq+^ADZ+KFB&jYev9wfv9Igl>?Hun;0
zLmQJ@QRI=}51XnKSGc=PW?OE@vtIN*Azyrw+RS%O&odi+D`u(FF3?19u}P>{slalP
z0N>DAv7UA!sY{E9uL;$=TB5D1tQuWZE76VoEBCED91hoALp$~w)_*p0N2sxW*v|~a
zR6P_=rI?DRrrPB|ARRDmtAxot^Pf_XF*UXFYdu}t%*Edi2VUXX_PmeDW+@Co6Rm%&
zTC!_Qc4dyxVuObGwfMA$Q<p0KAe<lPZ-MmOn^yHy8h6i^wsKIWGu~HT$+W&lHb@QJ
zLwKb!<2E&Fo-#`7uC(kOT*i7<_vx)qN<Ri`Tk@nnxuri<N(CSo@gRP;d`~i<AGY-5
zdijL)eCQix^eu#*bYK4>+XnaHtor@$LW)qa*-~`@h>`Tv*5naOoa@YMHGb%R{huxV
z40g@eJhGjE0;7e*THElmJFt=%9k!Tfq?*Dq9H#3+j?Bs~I(@Ie$eYW^ToQC8z!e{<
zb&^|Yf}y)$Q%0k5Y}r2*+^Nm>bb1z~ywcAl(Ep+1lB^`h|78XB*}C#algrmU-)sgV
z$lax3%oHM&zg3g_ynMCpS*Z$q6shqbt8go;`;|`lN5IY_^YkgKhV%VTrQF8i$QQMe
zcICPkL_W<bG&%Z5D;g#J%DMdod1cP!e4YJzyu+-4+8+mqtlX-hR4p$AmPPe-IrP}c
z=M%aYr}oGL6^B+)I04d8eyHXmJJQTU_68x0CcCBBH8d*Kx<Xc*regBls%6H;ec#yk
z0^wpcK0O7CUX%rXy}MI?LddVPSMNtXgb}{Z(h6w3w>iJKAfFR1#@LO7I`7U^BP*Ko
zM40Q>OiK4Hmn{<+h9c&zPXDo7$VNeUMCC<SlQAQ0$I9xHr1uH+dm~7Zyne9(3HWUf
zF|B@a-?_chosP`@+2_xS!>-DDAa0K^xiMS~Lhx?P@CFUuPy0N2C8BbHm8UUt&e~FC
zz)W|#X$I~N^Cf4aSwFp=uFSe>?7eR)j*LLxNu?v(`}UQ9(I99?#R_>IZv&RJX0!L9
zEQcZa^Rx5OWn1Qar$i<bF=XuQ+UKs?vZ(}}e=POZQNstsw<>U_dX!jG5A&#%A=BU>
z`aiqh0fV&%!;|Ij1kL8Jj!jlwg*BEe*5k+<&n?TpD5!tj+KBZt(@(pqkgNTH7Sygr
zy>53=l*#Z|uXua0!{+&V`I{iEx+(*dYwS$*a-5QNZ*~h8Y5<l|>1PNHEUS3;p;KIY
zzM%tVr4DgU_Xt%Qcs4j=Wbvp$y~8w$N-D9S&uvEGf}F{Zz7VZ>$HL<4o(T7nJ8JQQ
z$Y;;y63*f?mS~)jez?h<JB_})&O`N?O&!HHWK@rc>s=kcB95<Ve9t}G0@0OgBfwpy
z9bbu1t0T(#ePRhynyFh|r@sSZJ&*Agu57&Cn)>Fwg3bBEIx%3c^dxbmzO7wpTg2k)
zuCK2p4WG$x2cJ5+4%zFrubKx45C5_S4cgQEMAgpDeJhc(?V@w?MWv|Uhqs@+)=Upt
zyO00Z-yeZ*yaYiSD&X@>pZX*C>y0%Pm{s(1iL10^SuQQsGr??w)R5f9P4PFYcHGTu
z!a^@S<INQP$glCa;^5qT^j18z_P*zQ$c&`2HSC0ZUDUm397x<Pe>j#LBiJLgA^cG?
z=PNFk0TL=-gr|6!l}I?f7uvFW741z<(r-J@s;SCPu?QPo!Ky*TtN6*mfSi<9S@uZ0
z4_-sQ^ncu@6UQfNZu5|3^<B6)T8J^8U_6J{j1lUhj4~kANHSj$ztkCWM?r-!!XPOZ
zUO8ZLQjuoYz2WsUe(;b<r}=#)`3abNmu4{Z<bgT_Srlf=0i&;QiYC0{DZsAZPo$9#
zLL0Q=wg8(-!e!7tFHwA@308&soqu`Zt@)x2H9lGM6XNuzL%t0BijTU)e;TJqW~Jgx
z?tIU_08UY2v2@lAz}_De{(2T)Lw2bNzx6l!L>q*h!y9`JKIq=BPo@+3WWSPyw_clE
z4|MW;?i-BA*15FgYt2j6N;tG?6T#m<Vfhr!{W-;MtOzB5q6zM&Gz{NF+)R|X(taoL
z=}APL*isV|*jOI0p5!$vi3+H4L(0s&shBC-c))#gOhgfK)54<`EtreYaK*X=T1>H-
z@$%v;E^k8E%rdE&I`c`Ub2;WqWgDv#Me?^hPg)u3Rn{4RACre9bIb%^#PNg=a_Y!U
znl<e_unPkJqI+HLB6Q@;(Hm6L@7Vk6^!d7}d`@Ep0Z#na^Rl2Hkna{!9;+prFX)2z
zXSj{oe94BQXr)D;1ewaZ<*m0?aPi;g)qV^3ZWv`0HMeo#|4^PoLgX&!kQOIHroj1u
zH!0e@|0&pif8`;$>Q8wTnx}H>FQNMT)8;dOi~8SZ6cieNvBRHJ<H!F)x&FT%#v>!Z
zM;FemTrfD2*Q)_#&0k>#4BY#(68V*z%-04gKVN<Au<2OcL`vP(m<^w)oENX;MER@R
zF5yI*QP;JgMzhg?5%wcX*0O`y#vj$tJm`{&SZ&D`mxKI&RvLnSwywW*X}?P!LC_rC
z$Dx@K%k=XG=f-6@%3G##SX+KkMk&c_$agW{h3oo1F@n}OJ=d8y`PjosB_?A!mINqf
z8%47h(6;XijYGUc^4a)}Qi3%k>a?444SYl#VEeSYEeeSlO!LJumy-U1>DONNEHssg
zlb|34vj1MXuj1`Oew1rl26fk1+DF<7#%p}+^yYp8!C$M9U23COt9*HK@JVGJnekYX
zA1bl%J!NV9q<XS2u7^<+S?{Z-D+0XEmfv`bWsSFvbuZHoTlw1({biog6}q)tgR~r$
zpiUy=wC#wmzi}EgTn06Sq-LT|N1zb(_{!iS>LeNb^m)~C^RL6kgtz$|$K82v$fhkA
zCtPR5M5v5css1rAOQfDpVN{EGoEZ7-l*12f<kn^+Rx_V9(&vY8IhV&Q#;J^@+%~xl
zXYel}7p8P+)$K=SBsiG;ue1lk4S+u%TjJ52ArF+MR)Vzr)S>uBApx!Rf~IaW71Zaw
zYKa0lPT{j_?stP?-@9me2xM()$9J~2l`2DZH%y|o14qw?<nzqy%eu=>NnMlAgR6)2
zY#uxw+=N$_iqgtnq{Y8Nkd&wCeq2bUD#&T;X7L^`GAd8klAcPqerUIn_g%`Y?k3pd
z(44ZG{`;Dn?`n<7cmN%i#D(|XEFaO|>=~S$N@wslirB~>lM->|xKQvel3c5$A9<eJ
zOhe`FuNjMOI5Mq0{83B?VTb4-bOccFBUJ92AR$nCk7OKIh=ag_Z}G^X6jJ;_K?{e0
z{|~IRg`wpnkFoh;@a_M&%gb#kdYX|6WSN_JJw8wz-pVQ;0a|bt*Ggs`@~PEH^ieSM
zvM5trrR?SWRL?F}tdtiA{6@_!=z!rrWI8mS8m3!o(4*nCi*#{}6;0PQy6La8KxM;M
zbPRm{&?I04S!W<!J?-Fm#^;;3pmc~(X+LGoEpSQ|(nr1`u4&NJ2K-&`8^1lfQo;OU
zK14U3KaDa9YVJHz6PIS*y})LMt!iQbig)oA`AB|2S;}t#Q?2V!8A}h+tk=Y4M#5z+
zelVD1BNkAWJSv^~ZRjr=<=0%qDm$U3KdwFhtLwEAg7N1}0_f9Xk8D1MJYF1^ezf}0
z0b@Vte)yQ3yNbW{$<VZqa@`A&c2tzN0p$+C`8CJz)CBc>aJh1r(g`Jjn|3pQ^BIHe
zIFA0hN<hGocXMOmVQs_q$HKPld!PLA_Gc!IylqiGMI2d=Go^h|2J^m0(wt$#+97K(
zCm`R15VM*vEfWLvNAvmp?N94BwZlq#Y#aA7B=oljY-b_KPsyX{-x_9fN7{?AxJSbq
z5O?;|hF&6hnN3p$pQi=xHq=qiDv%ogAe=Za?dE-W&mds9#BrCJDaqWEIeN1AB(_(w
zR=d1#B#!rMLZxBXDr@jiPXMJxYN=bL_`^Kb_`L_iwFp`3;SDO<hF(j(q6jgSV^Do{
z(VQapO2fq|>AVefj9kQDnGL}zJ3>dMw~Z2Xq6D!iEorxbS;hrL?4Y>qVqfQlgDVqo
z0eIzPc7+;@v7v{Sqgbb3EHb3<jSc-Zi<wJQnpQ*CAW8gJ^+D<qbG8m=^hVns1>S-<
ziX7_f(v|GF*)BZeS4MxX>F;9zD!cLaB(>>am@W68d^qfnS~}zL;HB`p!HS=}n7Rbf
zmM;bftOF*Z{k5hkce+SSoHc(!%si{p9E;LPhHW?R^Hy(RxQ(CdnqxU{oc|ko^=vl@
z%YFX-xfZxV?Uka3e9G&2+{QPO%=LTQd5z8aycMNu{;lpQD3<xmfOj&`psM!7#<Cbj
zpbO%5FL}twyYvcqo^O=mRgtzWE1&+eiJ_fB{8feF^`I`wM7}G9TWHkB9*By5npzr;
zPv|p&9l4Fx8?-{f^Vxy3eet@E5OD}2mzHm|^OI<f`aSkf#k>4@0;#xYe*k)?e&=2k
zxT^%r8Qq~_{1YrMe#7$?BT%|PM!JNxez9ktO~aW0?8^_<HJP=Q2Dpk{XEmt0cJlI?
z+qc587($IryJQLP$FDbUfE8@bf&&{iS6#r6?S_{)TzPOHr;OdTzbbq+0RpOuw1G;N
zdgjCVTN%u!$e`tH>k^gB4!g>HxLasPjjCdzV%m(<ik^*R_2}dJk+ah_jm2gkX&~11
z`uX3s#k5!|+>(@=V#KZ4<s5b~{dMX$l{7<@<b6y(NQADG)Ny>2n@(AiOwt`NdHWkx
z)pjMpIPbm6t1RV~m$6AmtZHr`$7H*88KboA+5<`m+s(X#g(=kI=;iKOYnHz^YByeF
z{M2}YF!Ngpf1t>=@GE4Ilq#87+p*>~<697-<@cI##>Vo)vrTiKS7iR&OEVW#Ol$#w
zgb(zLm8L0?Rg{~l2~xAwsj9w0ZqOkM53%8!^Yb#f3(tJ${aK@IG7_910a1NHKA$BT
z9kSB73ob%3nSD}a7Ey)mJa+n)q3`Uo2_2_gQLEB>(p_=R;6&FRzT_I|S{b1cgV^<k
zM#D*0N8XMGpyPmKGO9!MRX<bw6IHFtsA`RJ%_sTQf9sgnr``fg*Crwe`5oll{J|Gz
zXa2z#KUq@Tbu4}E6j!HT>8cDNhHgAtRalGlnRqkp6k33AafN49(8>uJSOb!}L#In@
z1^OxP$DF&b=5?4yRY|q|U_S6S1lg(19q&*t-&vmK*Xk3}eO~S5^7Lo>MUCwGvg=><
zi)-4m58+Zf4;~xizy4;H&T&1r97;&Jzq^eyV@X=cfUN!j1tvrDvY@q^mvKBWFJK<F
z>eg>ifR66(%HQ3bx{ZBT;2jbLU3=9J+xEGTPDIIZ{+jFRTy4LrIRR*=GBt%<{y4F|
zDFsUg76CvFsBChQ9SC96d;2%td#g+s5B;)jX>KTq*)T>+=#jGvxl&%t=u0_YJ*Zgf
zbBFM#&O{~hX}C%%RYvo`Z;KSZvpGVUX#dCp^)bJ9!Ij6B;V0<H<UIY&duPEAy-=lU
zMVWZM!<;Tt|Di>RYo9g|)cwVF<}kkeL46c$+H)?NF<W6NbQVzTdV6IbJePv?{1PvN
z&ltpl8@~8-f>%J^Fdi5{4jEZc`-fe-V=kksQ>QRb6e?+0<g<rjbg0H6)(BkvY>fI-
zQTNhY-@*yH(^vfVK7`AE4Rs()3HXz1(d)6;l^)}*zk6!UUumeherYwMstNV<-=hCs
zoDjVAvaFNtG$d~GzEJt)g=IaCi-uUv&}?P$pfI^uF3fcFo}8|$d;RRKTIf5I+=c4}
z1C7ue>WCl6nd&g5>C14xD)At7#%+CdsI#*~FsS~#eQs=G_qjBmBQ!7ZNZYygDD;vi
z+QK{sMGjHahY?fz;>0v}_Y54DP5Qn$;^9SYmwb)#yOxjFKDw6_swK$`RM-s-il#65
zZ`O3Da6mkE6u#MV8E=s*{6dfK@(GUtg=xNpV-T@t5!&Yp1kL7c39Fx#Z7_6yxBpqb
z&G>UD?L+H|zVGzF^b-)899brf$VxYtiQittxJH(<#t*ZehurRI{$k9-@b`NZ6t_Ic
z6}Wfejfe`cy2RS!eHtbgZ=bLjsluamwy^(RW+gyoR~o*m=jGN&rpz-XcKU$-%N0?~
zhms5Ny6YUns-H^wJ(-LlwgK3;4n|fe(!T)h_d*Sq{{Jogqasc&%p*x`)w}9e6Zf9F
zKkeYxD*-swDH@hJNd<Ep={UQZ8!=_cevkfZ`4qRP$^12t{KN-qC0i(|n_NlMj^y_p
zeaTNxPzmj*y|Sjp=3(@INPEk;sM@#P9|c52K|n!4K#)%922qg)2?az#x`rNlKtKc}
zr5mKXJEgm0XryE4hGF2p2F3gSJ^yF#XYbE`<qMfvYu38fwT|mJzsGsOtOLM*Dm8hz
z`(I#>{$7RPrw!UKDOJ)|I43#w?DuM)yGk#q@AG}UyGg}ukMaaF+!<1dN#OUlVaR8e
zj;lJ1a~i$x0B`uZOqG!0Xn3hY#^8>}t?6+c_-4LW6#Snu%u6PGMu$f^akjPCr+7XA
z)^Y_I8?Ydjf7P{~x;AB}YiLug`lYR~Q>)$yG}JAXz;Z(XeFs8rktI}+Bvx7HOwRxi
zSnmxyM)`j&u3q#~2foHz{9@Mq?$pGS{-Lr%H{!yqED{dXP`J~2f;q~Kh5LI*jxK>(
zO-|Qw^j4_Q)yzn1+@-4>s#O?i+&N7K6aOuyxB6+9{W|;Be^_;1ln^yQ_>*2$0wUTW
zUyB!)$IfHim*b<!pnHZGI9)<#y`N?1Ofp$6^B_mmuGUJH#P)+t@gD7mQEPcTMp?C7
zjxa8*w&fD&ly*su9{<k1?Q&ds*1p}=WrB`-OAR)`=FbPD10!$M7)?nBcQUNvT7k8Y
z?4_D^zDXV|;20bO8*4dECyh8V2j!D59F_8N_2%?k-?3j(4GTF29EzB2!C5x7z2!1J
znfb)@h!^Z?Y`16^%p`WFkjsIVJlklS<ny>RQ}fkvG%Vpvi31vUnTN}_jN&#Gm;7Eu
zG`>zsF2)S+D=A1OvjwM9-`q+1$i~^gCedPm`Kq!rL{jx}q+kl`F$JcHd?saQ^e8md
z2%}}Aw!&IFc<MIaC7Kk|?%39_v7~B@E$vWyRl=}$fnmso>AKbsHc*3r+8*kmc07=A
zZ;{L6So<I$leahw1?B0C>Ju9@6n)w=&&?x`LH@EBuD`D`mL?*Zg0X20!_YfG37K^H
z;cl|lIX%{6JKw++0_eOpXh7svZ4y#Z52{6jQL79%obpk}nQ)(QCd4yPz0IeT&H7z`
zXdAyBt9Nu}Ys2(7nXYj;U|P!t43B(6ty^Fdd}}4p;g{ozI33H)VrH(_Fi(5P_q=b3
zyFmS+q~<^C58v$8cM>~Lmkt4&C@Yj}h-VwkR{&!KlzvY#34&G3K^f?eUM!~P;A=}u
zR=ue*4N@@D9(95wt{&*TA5D^BE6(sx=U2+4ym=|5CgxJFyR;IxI$C^!h9#9(;oH4!
zY?kP7c=$O~i0-TGdpjbVccXKib~F7Q{DEOOB)%r?c@U<6o~(cFd)L!>DUF(|3c+d~
zvAIG9j-&yFHmGp-qkQ1B4$<y>W`v>E{KM(7_34UtfayW-HfJc5KYh}ms<<kRMTb5e
zt!pqGPuH$Ck&df2aCsX5Fa(<HUP5bWMj@}VohKv9Wo*W33H)WwNJih8?Vwm1Qu@n4
zlAGZz;WB`|yf0JWExqL;RqS8NA}m~dtow^kwuEHSaw<}qC5CXHPo>O5Iy&-+$8W(L
zV?PxwS$TMfQVgr8p;%50y_4^tAc^Q_t@~;Uj{@SjrGGJP=f|&|jTSf)KFs*bScc`k
zqVM&2oD6u|0;$Ai0+*p__z}^oOpsVD)nv@$FNbFCR8{pS^-^*L&NO_6MQcTj=D0QG
zx#ri>IdG@B?%EhFQJ%W}WZ;e4i&}Yp@8e-?>tY@;P+_Kr9{-gNmY}}U!LRa=I=C<&
zBjuojmv!A>CGyRpIN+8Jv=b$K%}$!J^8?ZvF8Yq9WuU(^zveAmO$DX(wu|Um3Vmhu
zJ|Kn>$&*;59LP>@PAYFiePCf7*zn*JwYEpHrNYhP)`s_*j_#ti{koMjDy=y7K?1TH
z@YIs1H_CX7G3K0;Pv(A>TLPzex6g95==vN1iYxuso!vJDHHV`f12kfy5Y=%s;PcF3
zYwP+uf`W64LA#u~<8`fV`jh(Cv%`+wefFG<;>;ytrRcjYyYw=4F$1W>KDy5IO!3u!
z^*5~8u|>#x8Um?o7wof&JEl0!3jLX5SukS2Ci=*n9Wne;o<PLF_iQiotI*KkvKo=X
za;KDQJ@OftMf0s0;NYA*&x}c1IeodYyjegyO6wEPZSH30%hl1EhHbm~bCk{!D(@rp
zC?k}vaTt9OsdBAgup4v`1<HZ&GP%_}Zpk&jEi>fgSM{{M>4d1rRMS_d*o6K1XPmG}
z4{{!p&$CCz<FL5CJt{-aat7q{TJLwKM2QGsBrN$IrCv@_z%YNGu_Iczy<L{UfqTcI
z>=_pi5e?OIlkWvVcNmzC&8Ix`TW*ejLOqT)Gsqqtrx-QrNOi_z{I7ocNr}YncEq4Q
z!0~H+H&H1#dPhi>uTVHTFL7HB&0B`QJ&{+cwIV+gemAVgc!#wfw-CJ-JPAaxhOo^V
zUvy8}+aZ|lb?x=~!}q2>AyF#kk&{!4D-Rs_-5m%w8bLWDZL{xplIg$&6<;?dpyy*7
z6MvOn-2>G9dqN*{%SCmpwI!Je0`WTh$4*ovyIIAX8|LRL7ce-g!P2Qwzr>^5qGWS#
zgap<}6xh?S>&tJHMN+D~j<WB3PVK9eo&MYXmiH19OoDBRe|hAFSPs$JWp3qJhmQlp
zO$;zMz_fgzG{(IL&1^L;{S#Rv+P$FCnZ(Cs1tA^0l9auriJNUeLdg*0tG~M?egfaR
zKWng-C9Mfv7+JPNxHH^QpkDX^X~_wO-y87u7yb7|IY0`>7<u=M|G9fsC&eAF<N70Q
zpL>XO%eTVu>|G;Nj*v2YfV1|8IUb{_ae<l|#smoDdoO^V@n-S+hWDbsP0b%aAv~Cc
z;!O-P^xG~0w_H)-r8s;szy0f(P>BSiKL&Q*e;U|xj{ge-n><x`Obb(d)Euw4<LY-o
zPfV&R!RvwEmS*3U{5^gteYnaSJ?UK?{RWuy`T1>%@Fqoo*PvNBOr72VcQCc%N|!-v
zEa8Bo{AlQ3uPa|~V$TKn)RU+r1)Tt{-#VN20`kT#L_0t*5m<|L$(UL^sLI?2Bx_`}
z4dJ{2V!KuM;FjyDMu_YEr$&DzTz`3x7ONYYlpNp`SBYV}38cnY|8OAp<pe6@(2~rf
z)v2KC2Lq*F2w-?t)9$_eRf47aR?QO)%jQ*>qy#MLIbbLLFu%@&pd2s(fA_Ru@L7<S
z8!6!5KI1Hl?@-ke3|X4Sk5-4|Uj8*DP*#cmrKx}e{ofnbf6wj@vc&d?mOG>9GAip3
zb<_5r|6B!h`HMc<?kD}y2*2=S;US0-{H<EqB)}t)AeF7=-}w~<<qT+rLHPR9MPz4n
zf{~*}J)$v2AgS6|TXcN@4GTcH*E}yK5BZQX%7E#^SO^0C$C$nw9%GJ{f8$5dP*O~o
zSVPg@qm!_y@TN`jCuGvk%S#<nt0&NzbWOCI)Ym2g<IywfSyk!Cw3b~w7|^O!arv6~
zjLxsyf8#;egla4Yvw5SR&M~S=VKh=VZHXI-STPCKjwCr%NUsl&D1(%-Hz|=awwUJC
z-h_GnDiH4dwSa&8`kh@s>15xlc<|~rTGun8Nn7Tu+O0VLb|_&2V0|PPo?*%e?%;c)
z;+q@2r&Kq~XFSsrb*FmU*T{HJ&~Pw#{|s@4a{;?P63S+@KoUBE-Gnee56#cbAYEyJ
z8Hg*l&J<TZ;wT7k&o=gI)A&I{2w+u1{c>n4HY8w_d`ZGq7IKDoC{qmO(%`+U2XLS?
zZW4K<o}K%od@W|_e1(RB&4?(<NlkOVCk$$(is*sAB^w-SsVXDEny(PvZC9kOWPFCl
zcx%W_)GLH@HA@i?$uYJCbPVeqol=(IGBfsznf1!sh6Y(2tejz0+nzm5tHVk;*bFkM
zdFM6$J#-DOtp{T{aFz&{fG&<A%Vje)m3x8qSw7a3Y`-4Eezu!b7WJOJv2migJ7u!V
zSAmeLVP1bCgXUBB?VkXeO;D`5VHzH(yh17rTn93ig-rV)TtLW0B7pI@;_mXCwX}?L
zsM|FkWm=sykFPnB<aji(salBZm|@&Qs@g8&VDxX}PLj2vf&{h)O_f{2?wC>YL^v*D
zbYK$U2cXby95y+HUDZa5qP>)Q^K4tg;vS&YK^kb8=B~D$Ou>V=1>cxx$$RTgKA5%q
zrO}$B$RCy8!|NiVxMf)wQ;<MQl(V~WJUD@M5e+7xJp7o@56;K2370OI8YwY<+MtjF
zhzwyaY23d9Kd$%<5ZNY(SNHbk_=J;Oy<2E_mnZH&%w)%vf0@Z^{;!LEBF$uqc4s3%
z$fDE<KO6u?*fm5Wh9Nze2;5F8arIm>dCSlU1Gy8HVrcmY*x)sZZJ<%}BufnT8Xy|^
z_PWNYRRFE*_^vz-A>00<T5<MNr!Ij+Z}$DBmg0M=n>%zs7NYdOHI$W7gT;xMctP8x
z#4%>Z7v5%q`}mASuD~}{g@;-d2^YK`${5Rk=laVJs2cOyN_z`YLZ9Pa=C#J+PtrBs
zA}g^XZBR%viOC_6TWd|(bNr}axE-3tP#cqbjeCFg6>k=+wcnut9*F*F3sd@?f-=lW
z1W{|4lk^|0JB%$XVMj4^J1*pZ;TDS+<L;<W1QU83xTlRBymkA0q*yga%8qp_hWeI5
z@I&nQ+TlRTM<13ZTN{%V{3ggvVxp*2{3fJBnY_fPA7~6QX2%@%NCipce6d}6&E-&^
zI^H@Faa=p@M6vd1gp4QT!o)71@}(yYIGiisq?AD8xyGj$t&W%{ApA!p&CNg|gvREX
zzzEs(WBfv&La1stT5fu~6|=H)l>FgJv#E?KK4UAq2TyscPDr*OXW&cZe7X~6c<%^c
ze+Lu2C6ojcP4fKsR<!(8ECj{7&Q3b`)BNYonY!ML8i&`ZQ7K0LgFpwN-OU-H!?doV
z8X10-vZyx32JK;e>o&7sU?DC2%<7;$Y}-)RUoPTZ$AsN$k}UyQ;=`vCPt8C$m1I$d
zMVmMQ2Nxzz0aRc|&cZE8t?K*;;C)?!L!4Wj0@oTW%^xEBs`ph`Am^{)@{}EySm{a(
zx6hMrWgRlIJU%iZQQi`g#tr|~D2uBf%=k$7*W*T>bs;vi;c0Ts-!Aab8{VwKWGvu>
z|HB1Vk|)A^#-`xIx0=qQ=)*PUoISpLAKt;;01$*!>6@X4CKSmc4{C$89^*YzribO1
zPA*I;BOTt1zoy2E(8w_T1ew}dzzn7hS25bFP+NaB^-1;E1RuDtbsrbw4Jxau+ac-e
zPr-Z*ji{x7<+Xq9`I0A+mWVfL3x3zDJh6Xq=5kWmFZbyoxB-umV!X5j^3tw$;I)B`
z-G0zlF$DTlt_*3ViMP@M9u28Ov59fnsp1!zKHsXvU;veAi;dNX9d$;lm@kK#Nv}@k
zz%j!C*v(~m^myxoda?R*@dD`#ld|jafTGJqH;SoW+|&;KPF4BUA{@hozi_ss!s_@C
zz`r~to9@xh0t6@>b$fJq&H7svBVLjKJ@~n3OLBJz`ZHJ#s%sj5wp=ZA%w0r>*!e@&
zPa(z-^kTI(rD#ljQVjjz<#k}u1k9&ZhgqZJH&=}$vVFs}=G-n7qa)Q)8VzZnP<s(Y
zZL_9on7NUPtnl*G6muWRErQAwSOQWHS6=C1->K{&1wEh&exO-xy|FQr=~Ycz7dw#y
z2v2>me}`d=`YRzWf8{oUN)~}sm|Bg07w2X4&|s<gr|_?_PFQJw3-F(1i6hN-1}Lfl
zJPlvbs&Fh(+^o6f1qRvv<_HK<(`j5EH4n_>+?WL0x?46ZCy^#3x;s2C?@`ClVAcHg
zZ-2oCsZ^R9*FP9=sAC&2L>-`EF<4KN{J8&Yr<|8c<u!k^Vy~xY%tJt}-V;r(vU|zS
zFb>`RBypR2E<JfB^qVX{EtU+r=O|a<8cr~PaPhRYlyku9NT2LlIe3=qaaxm5S~70g
zA{yabi>pzT8cx(l#f#v-OD&X*OxP1Iq{@Bj=AC_KK!seVa8}xKX$uyEh7}VI8y$7C
zopk7^#`cUtjm}K)#Cry0DZ2YTsKaXMGKGe2cHNg1S`sk>lZD?nY&dOtE7-T1S8Q&6
zBl&g+8dfo&fMw(CgvIdj-0Evp0Jc6Ney|yPwr$Yt<lG9YQ~a696*)mi<}=Tiyw@sX
zb5=M{q&~~@Lo!3|b+1}Dn(wpPOE&72ie)Nc{cX>Xd+%7^MY;$Q6!Tel;Cvt<ZWuIk
zPpxX}9Vhp1)_C=I3~t`Hth6vKcpxl=Pi*mX-qiA$3eP@bqV0Ea0ky)O+DJD}&k~rX
zbr0S&q-9!})4u7T`1^4b9+oTX^@*h2znxd-|JHf!Oc<meHS?8RsH(r_0|5PK1`rO3
zlYk4D(mG1Qi<QG)RO9rrI6ME)_1Y)`j;0Ri-hU4I-*NJPg)s1ETA2TnOWV@y?6I}&
zialYU$^t`**ZI)QQJj)V@@!=e5+NLTItN!C0~Sh9PQvyGz-$W5ULV{UcoYhx3;aym
zt!`h?IK9@Hl#F4U^oq7p08&Sje6ci^HG+x>W^47f*Rt2CXcofTyP~0<o>|+cH0_)_
z9-gq#7VcboXdR<|dmi3?P@{jX&13nI!2T`3m~AXa8ne?bSXm_zQ@`ARML^Zk`I)Z4
z#1DhKA05xCvX8hk#rmv;`2z<Yc{(v=Z6ypSRc1`bNBsV3lNj~LzA*uOMo%%DLsegb
zqeaNtZ^29?h~wE&ZetA_mcmm$p6~j=pk#C+;;~h6D1PEso?*_?Ui!y^;+qq2F_3vO
zNU~c^YlPAyItYLjjorA@6P<LwhEV1Fut#O9HN2dCUVDiKRkC9egh5F5MkX=C_hBJ%
zMqV0~U+wD!lP3eAC6hlA0B?uC_jHDL&=e?M*KOOR_b_)Gm}<<5wYny%Rd39s6i(g_
zBZR!BQ5bdK7j?NZgdwfYt!g9Xgp0#}9>F*(=K5lw+GHO{Kgoi}SgW_Q>*J)I#osou
z2@-K%dh)aZ4;Us;Ie_?VJF=>f6WOLCtCq{|n_+;JAEz+S+EoqfYdQl_9RJpJ8ji!C
za_896HL0q*EyqWlNeV{~WkB2+9Q$#aWe4umW3{>dRq7hNmY4BViRL3Ni7|`=1n|e^
zinAI)axuSGayGIv^~Ie?IeW1j^8q(mB}i*cFk+i?I#X&3=}!za{E&Kc;~p-L-Zo1M
zdImu5CpO3#6$~a8@olWD3M$EtC%otI;JI3+P~fksWO1)%?Q1kxQFt%gE{YQ6an{ah
zj>J*>_s}*l{U6!SvCGst<JjjlaSkiF>*sOsOVyik3+Hb*1vS&=ixuOYr$GJUg{#!h
zR8fjM^FrQErXIyHLfG<cpO0_z4nqTsgt<2j2%v@chD_>gUg#Y;P8p)&GHPmIP17_7
zDVF*Za^bvsC2P&c72{OeH6M)=nPL(n>!8lYRDM({$TbA(v8s>C|53taUH7Hu;hNPi
zp{0`Xs-pM4GBJ9U!ix@|=#PX+y+iw=^xY(bIWb?fWY><2SNRm%x5dK1)hzj_#~++8
zZ)rMW;n<I_6}dQN4dQSVQw1s?@mBZQ$DA224fNSedX`Q8c2S-hfG&IS?(}ZwQW*%_
z8FR9jWzHHPS1j`(z1yqAdY5jc8biRkfK)&(>2jba?asn}McklL<tw?LjyLWh@qQFZ
z0Ppw0s6%yKs8GZgp_ONjB07&ca9<m+GU&xQkx0$Cd+}sYMet*|(Lav3=(%FY{I-I-
z+aJWgr&h#X4>a*Z`WejSGkar_jKKUGaQcjuhk%vlQcs!_d&s4Rh7<WkYr(e4aR<Pr
zTGUsLLW9ox#ph-<P0ylQz@|+n(S$cFSNLjJA5VBCSG2}U-v%-8fdh!Gh)%S6*2ZjB
zDxN`|+ZVs{j|?=f53QC>O+c5&CTeE1k0)C>b^G=0%1#YOc3dJtogSMPvq$%rM4eAr
zOm?4L1TB;rIjXnYX22%rjwDA|VDfz(nCj&fW{ZSt)e?|t@&x0o)*~#nqF&%lWpwBt
zCMJvp3gMf*e%_z}{GPUYeqnheKTXi(L`kV)pVM?nvWz8+GX6Ul&@m_V$Ccmpm~_5_
zf&0iCK9GY3{%y?T-rT@c<bK4=AJb@f77XM9>@}`(0a4OdxqulXRJjk|eg*o<w2^z<
z^0gTu$xFS9D9XyW-cX~l*>51aN3hh6k&cb?TC4uTL*acG(yxot`o5A_wOJw}QG(PM
zi-m=s$1@ns2e9WQpzCHKlK90lFx2^Fogr#cs*L&_(;oE~lFT`Ei&Hgsdh5(9;RGwa
zS=$!+itCFb*hh$*g+op~0CR?2jr$lAjo9)|V|$s$^5s9ZR{<PGYaa>xL&tV|kT?tg
z4`q9_5%hPM2zk4bdrRD08m-G()ai1Wl^8cCh6WkMFo8%MS9G)#G&Gs&b9l<ySKS!$
zPYJ91j=j$FWPD*SH~Di&Z_ScLHelr5eWudsgITw8T$0*zVr?RKz@FWAvUo<+QJJuG
z?eE5oy3MNY9ZzcvxiWDzP`F5KlMZ<kDq6UN$kZee_ih86{>%VSAw_=x<Dg3Gt-!~W
zR>c#p2th5*3PlMWn;m?A;&q4>t*z)`c+pFF(&V9dmagda-y{X5?p`Max<#rp0P7!E
zio}#Cx-)$H@y^8sz+--@<QVd|FL=*n3{Eh*)yQvlY~+naC|~}n+dRWI7^GQP8o}Px
z<5!<G9XzW93F9BKRR(}m)c1R9!)poRM?1zboGyQIlEYJbDmM-a(HCucd{Im`4qiBv
zn!m8p>~g6cpUtieMG%#ZyX*7Ui}6CsxO;+CXoZLM0ppoHnqDmZzD&LFqCIE=m_$gj
zzBfnrtQ1By5RtXuaw{;zGX6STA12fo=wlRGz=ttRFi}haVcufrKa)jb@t_M~^e;G%
zwTn3e`v(BjY2IdBDgR|7mbNZThO9O06X0x)ph?-d?z%>ho}bF1Jqtc{b1Qio-FTfT
z@I-eCJ+^JBOMvvt-n!jK(?Bu7u(f6alp8+h|5I-G^qn|tO!;JXs}2OmHi#`~&n97n
zs~W&!-hB`*-!uRne5VYD3g9pogRvG(tj?Y+t)aacy)5i#Z;TCKN^FAG{DgmK93nwp
zmcc?FP}CN`1pFb}$5-Q%-SS3n$~w@85MRW?@`~jtq6h4y&k*Eo0B2BDu)(idh|6lK
zgiXAgF-gkx{KCDp>?G8BL`E++xLcQu6YeLMn{Qa|ZQ-;u)u$9^{@c3!+*^oH&){rW
z2)VH7dI^L&Sjk6gY~th{_Pr(9tJVsfFw$<be?Rb0W1#B)PW~d8^!$5+_wSO%2|6Cl
z5kg@?I++18yr(=Gggx-)0>2{-8L%kNDDL6CQ!H69c0|%k`{}zfzwW9*T(0?hBg9lu
zLg|oL@YqapvwPpnsz=4ZZS8t}MW-bj(;E^iO-CQaotu5k9G}W*Y!V0Am};w~!DR<s
zNTTw7*iJ|XJh}}E8RF>pS!&o+Wb2RTizY=1P3r-FnzUnkgeOP)Zdi2Mzg=D^E|O6B
zNnnG4^U;~{;#6M}-}v=}HjoKCPUW^$JLXBEL|IP;hezoWuqWlGH>YQ1{#_e+A@&u;
zN5UeqJj&%Pm=^vkPy({_fJiFF+(<l%A%8P5t`%rB*L=XhD(Ba=_6ewqgg!;qMV4m*
z<ICNLwzg_62tf0Nzb`<F96;Hto?VXtm_b~&viG5&q+iDU8&0I&!+vif1+1WM8GM_>
z_Ua7jI#{byoGkaIO!zck9t@#6>>1B~PSG)0sSBC%D*pItvsCq5dtoM;IeE{yCGbSz
z`k5Gow6#*GtFJtx-=e18o*!<0yR|P_P@<;k9}8DiqPcz!)Z-9wVS-4?Y3^P8jw*KG
zmo;2~P*wRcXyNhT*e;H&3Ln#rNUcz~Vj%^Mj}1-y4#Rne1XzCZIcbVj=kWN;C8G{1
z-(P{E5V6rXSAM<i9HB!E9V}cRbW;WPbrpvH>Ibfj?tRbF{HmXpVc$v#aK!25JSIy#
z?sgY`PJu^P^}SWq2?qlDbxlzR!26I~N*HDXDVEq?;X-t;^Ob$4@@61#*0(mUaueU2
zhp`xTLOx8MCTh4!HSL`+%M316zCutOhqrN^M7plNF9#>HH}cu_SbrTt{bw$o$wzSN
zLiNkvhstLv8*}XWC76RC)PaU>BdJFd_q16kr3sC#T-&!NN-CI(4oEPFC*msFA#-%<
zv-54rd8aGA!#@psPVsvCEOl4QgVrYe;C|-Cb%j=22J|6+$^kswk6jJM5u^C?X}5`m
z6MxfG1}Q|5yn_+jxDcQQBfA{X(gDDA{cpY5c0;8<hNDQVVN{Wf(<JMkknX++D~o+r
zpm}vAF=C6XR>Q;pGqB+ER7J_6kUHDeo40(#(L;YMB$s-HwwtX%8)}E<Arys-C@8Pn
zK-eX)-K!YSR7WXa(_5;Kp<AB5)a@dc0Tv6Osd<>`moG85p_tBZ^?A8=x!sTAnF5$J
z%1macf}dsrS8Sw?3+R+PTR4_~2ox6e9cehyAQ>+y&8OJt5+2?<_7cdKS%AGP@lc<|
zot#3%^&cwOSC1=-&NxxW`mWBygBwcegv;^T>x;AU{cFxt)2=U$@VJY?<MYuWADIG5
z1hb8cUBqEwn%UJdnP_A5y{xc;4l2>NX0jzeyc8e~WVKWJ-1-IkVe#7?$?<%X?_ZUp
zm1e*$B^aH&$TY)+%I|8>mJ<Ec%`O385q;x!X^sq#g>iECq{|yP?4R%1ILQkTOR(oM
z!_eRRtZ#S1o4pk`Xu661P~q$q1q=;KXTr{z8Yy)PCi44t&pxmVq_-lv=YjFIOG-uP
z91yT^`?0GG0J9?o@Or&;9Tq-+6&v}bSEi%WVQcH6lMvg!8I2>m5JsE-SL!Z9gCH@P
zSx!|h_>^KRsevb~=YBo=6C9PyY)5gx`@I>L5$Kdfe@HE(Wc(eGrI6o43T*k3lAve$
zXY$<-KD2wkZY~&Cf#$1Yaq;|M^4xFB*wyYh0ExqL7v8wHn${}^Sj9o#Ri+$8qKL2*
z#!6IM`#xQufA?oz>q-BZ*p`^^dntkpjw7o-A;3;QpEVkwEB$iyXpZkez&<q2OZ@cT
zM0c_#Bi%9Y{`mLbdy`Gwh$fc~;$~fvL_x-`A$=S>Bo{R^_hfd_gV3(QG-L?<CT2LP
z8?XfbN!Y=+6IW+;>jm1`pMw7Xd-dpgZ4LbUe-n7e{`)y3*!6#(bN~Bsly-nKqGm)L
zaS#~Mj1*4WGplRZ+lOlG#9<tN*YW{lV}>8p1)%w65#}vi)I?l=?nU)2hBg}@#`I3~
ziFxKs3XEM3q)P0tSx$H$<DUB5`gM7kRdm|wlba((w@lIj$i43M(~Bv5MtV5kV)ii*
zNGmEczXb9qb9`HuPS>3~P{xRm@$;u51{{J3f{%lN<b)`YoIr9>G+F<Ji?R?eu=!V*
z<^FSr*I0}fWv(x|m$S}Guc}1LE<ly2TkzjiBK3i*D$%zIwx~zwS)w~*o6lc<d@(jn
zCUy5b^cKTC`t6Q--sHVv$Bft_;oLVCvjM9?*|*feGYl%SXt~p+x&%fQqbca}fylhV
zF!_hH!?5ef!cr{XM#-xFl6*<dO49!jDvws#`5!{%nLFNjUwhlG!r))GAAVh_fI%(C
zD{u8qfv~boI$_c>vCnoXU$kELZ$x=<xnMcU9-a9`EL3UDG0r?DT2Yf*r+)T|k81wS
zM=7M0ADFs(a=YL`wBlqQ+KQPb(Ayb$%bV_~>fzcDz}0}nL12|4{tbIr;wt6C1gg80
z^m6oidhSa>ty8^EMw1jILOFBurG~J9LXzc5Gvi4fz()Dr{Q#TJ90@Jr+I*Y#9W4Cg
zKG4JaY42oQRr_r(#@t%(&`p42Nw?MdChIyovJqk#;#ZW@!n@@IJ3*JJY-<tnJAXZ(
z<)2`^{X`A=w@h4j*My^CUdrjm0HtXPrW37ZbAUUm7XA^?k8jt@C)ttPY{U+eKzWU+
z=Ud{!MlH5(xi&H;e6vHcp!kDu(mbszeVqSKRsDQ(F=#Mr^r*pBWcZj+d7zQ*%Uhpl
zj4Uo9Y2s`HLiZ>d083Gg^RcOBO7?u_$t1(P4<!4HudNRJlAA1NKucZ-v^_j0GL$<B
zlqDrq-+kB-T(7}r+<T%_9>+!NbHvm=7$m?I9!eCp#5D24d814RT>h_c_IQpD`ce&Y
zWdtUT6nG%f{yrr~)xAZ(nm!fC7F%RWQhk5@ae{9z|DNsEGC)kR3pZ$k1DxjhfV$Cs
zC4-P*bA2}JrYBwF^@I15X1@6vjAC|@d1x;_{$<le^z=WVs5>o9j^h7WDOyqac2B5H
z0+Z4g!365)r9y&HJwX5%<(_l}MlH4DxAfD`axH`KNo?~)bwXUcz+uE&L)&vSJD_!9
zsu`BD9GS#;egKJSugH?X1pfF6jCvvP7mP9yjV<>qUF9&V5#~W;K4F~fL6b7ue9nw_
zKhxFR0r{hr?>vkjR-NFPKNgHK9eD9iRLq&;s0^qMscyHnRornhbdv3U6vo-0T{D|*
z)8xOBFlaNcfM1)63=wcNxL5W{tdCshpGhv&M(ObFZ-8SGOF<B+KP98Do`66%y?HnV
z2(s5sE&e}oK@T3$8w_L53zjxbO-M8Tmjtvc|CB>4S9H(b69dEIUd^`PdQxA0|3_+e
zR0Y#s|G2Lwucf8sGs)m4H+;P!vSC2$aV~fq5{0_>9a+HjT{xz_YN_@%Hw@3Fm+CnK
z>hPcFa<|TZi7q!b*e%%J-+mnL+plK{gfY)Hf7ZUYQ2Q!mRHh!L#C2)AlUC!a=McbN
zqPmi=ler<Y*@gy1z4NbtowGJqID~7-*+Vkq$D?TbxBQ5?rGgx|ddiPZav4zEp3%wp
zEKT?5An4!nKGB-lO?6!VY|TGuPokhbGd-MZb})!dS>p@|X(6@!;=Lc4<I8enoPvFX
zMBL_HqJ=`@2Q?34ZS9~40DJrlz|ay<u5qXMfiJNEQla-|`^OF!slV}A%>;KYk70+1
zDOq0B5k>Iqj>Rp^R0Qx=W*|1hI0pjKZFCwn4}??*{xPV7&Q<0Pob69NipsBQ+OKJj
zs>PeK1-43SsG;|_nf}Og)J!DG`VEs=ACFkvKLoh4_am8dc5pGa<|nmTH~(gcEjLOQ
z6@)Ouzk^q1WygOPh+e+dAx<}&p8uPzhPG;LW;1`*NY&15!AI#_#fq^Fa&rpzv(!{P
zeoT)X#A*zZgNvPZN)NbGf4Z!k+}uBoyWkY)%!`vCvdN<{AIKcz8;dPpEUGL;m5TM|
zCEd=ePyOVu6{}MtVE0I6ac23+aVvO(CNay(xOfo9CT4__y4YOdy1I*UVs*Eh+uqdT
z8i2wdYQcTz5qNsEUS{^5s0$H5*1s{z8m8}#^BJcd<7Y1PGcVa6RehmbM#UABbzg6u
zuAz4fL4}s6*U$0tUXl<%Kn=lIZ^C<M8s^nGEQgPa0g5qzQpI%rt<8kAsvjL~#g^wg
z{^r<8<604R(@!)Z)+r(OrZgM`;xEr6ye%pL%Xp$BN(PxAZwS6`w)ak~mRBuS$H7#+
zHo({1XcQvIGrzPaDdk+Dm;1Li;mkXH2x)4xd@a!&fezW2(yrH*A%KjtM>BhgH5_~e
zxUGUrKQj=CVR<PNaN6znZOZO+^W9;`;7}CCLf>i<8TcFCy$OVOfs+Toev27R)iKC#
z@+q<KzagLIrT_oYU4VQ_Ap);gb)GdwK?$>*43uVg?k6j~T%ws7^xJnS+Gt$<q1Fpp
z$yw<sgX|k(7@c!YPw^d!HS=)mgVZA9b-G|L*QdZA7q%B_xSUu722<(N?!a@@16y+b
zO0Aux`@0zy(m5$oc^0PqsJ<*M{#tY0$NW65kG`qK4NOOlndGl#(=}S*6N`JAt)Kss
zk^?4h{eP&Z>!}<d>d}FoK~+WpjyCdf7X-1;!2j6FvX=}{W`M1Z`)^|Dh55gxP2ijU
z4-&Tj9Phua9R)SKs1yl9^eG|P`@gpq<%KJ7GXLB3?LXh!U;p36Yw;Nr0;yG=ySD&M
z=VyVM$ej>#ApcDF?&#blUM1b<Rn@b6FMs6*(Eh#45AsUCv2v9^ErWnyL&G2TKFgc+
zcj`8W$>bTJ=G=jC^*}#hb@aWR$F#%frvlq`+d3DSf*o9P0~`2mv@QQf)7pO5YE7^0
z)rr=qW09n9`9oOI6H9ti_3Y9G;^oR&4SRo2o(LeBCMq8Kj?4BL$bJ1^hG@fqO!4MN
zC8ueAJ0#H0EN{8l4Y^#m4QX9AnP3yqBO2FowSJ?>p#4Bf+MUWm+EbOtu|I|Ceci>Y
z0)gv;y%<}>Whzh^yh_;OY!GBAeO%hzf>%FYC<CPPRp%3A3B@Sps)Qmx9R52lYr_Y9
zd$EjUnWmm@P~3w?mpn})&)VWJViedgKA;PWFI~-Jdp8hlW=#Gsk}3cEaah8Lhn~6~
zk>h=$qotb5t0qBy;suYSF$hRxM-#JUYu%SjCmPanGH4iBT1LjJdQP0OS11zWG6(Mj
zt{Sfj0NL2V9m5J&KCp@rJ$PvaWNOB$(tAxa>CHFr28`@@r$_UMQJ}^bV*AyEGr)eN
zydg+MrRRrdeTZ#ySF{VKtW?Y%NVW7u^OZU<`6mnU)bsxgX>zW0SB$3jl8cSYgEsK-
z0Cj$#^^k3U0YR8#Q4X~#(03?(SZ!DbXjT*plVyqw`YyK)AWdO679P0#vJd{&hxauk
zj0OS3L|?^#!Xvd<LF~ex?lYLUr>;Iph}5R1xStc@gDJ?}&e9c48UqG88njuhGuq-G
z;fs23#v?R7I5lV%=Xs>W!|<!N4&6)v($+wjfb2E0d|qP*R+Hv61a^iS1t61h)PGg{
zb&>MbbZ*>pOdS8{c@;64Z|mow@ogw=-jvRDTOnNMxY`K8>Jcn(P<I+PUdK)7<LNA)
zr~9d*RrEJc%9R%lzW}4|T8IW<N2Z8O55JTQs=5-mEPe4in)>-&PYa}mKOINDsX=HY
zx%vFb%Eo9#{msgvZvJOhHthD#bV`-TVJ`evOIiH)s^|kZ0T%ylC4V5uTtgvDlESED
zv7{DiW_l+R+3KeoXY@~8R-v=w;wPR=vUdG@gd6*M{ri<MazfYHhukW=+ZFN@Wr|<!
z2^oS-(Yo6COjye72{nJsfPkLpwTe}U4RA)ys+yC`Woamy)R*E}AP}p6bS?H1mAy|g
z2yrx0egOGd;s9>X1%*tYi^3!2Emouecu^>;$bU{b?IZonMwU+nBO9$t!?%?hA-hV?
zZC@b7IYfK(<5s5VeLEv5jk|+@<3IEY{pN(++E=G;#-Aq9kHjV3^VlCL2nTB!hM5y%
z1xm6Sd*y)i!?jqA+2zZgkb>UYZw6==)H`6Z(I)ig#Uit@9Fr>Wy*h3uIZDBtGgN~?
z*Ues5<&2$_-oaMf`cd5?(I&2jK?$RdiLBuM#bk?@zRu%U0oZ2(Oi%wl?Vr5^k&S$g
zt1VyV6@kTDcv;mHfqqN7|3qRr$kR^NG~83I!3B~9DYu3_72<$?CmBgOm2(4>C`sES
zA-a*jU@|+9Kfbht;=@!bmAZw=TFAOggq{PyA)2xW0t<?a9)M<}6agc7{?+fld4A^X
z;L>exKBTe@7AX;#awM@5W*g9*Bzv-kfF3=OT0b!Uk@8v`n9b-9b*lL0pq5%M18n$#
zJaey!_JaWKmMC3z(a0gYuS-KHbL5fNKMUms1y_cJj*69xE^b8DhrYewPxYe!3DWf7
z6DIQpkOSy&-$f-Y@K4@)_Ugqn$=*BO6gn_p%Gdz95hg*Q`Y2PYC$KRQ{m_Vv%U#_6
zfquGa((=VdxFJdUH1$zM+o-|ZVK7*l=1g>>GAP-li0Q?x>yy&F!^+TeVUT48RIwj_
zqk50b#ClaU6>UMO^QQiw|6sfCaM=-&bsKk23{BxbsDH^2@-zhvzkh9Ur0ANX*;ikK
z&#k-^_Ix*tE>+;I3mrVdoIZSm54_)dUJ;QV5?YjVIV|RpEBH&r+AsFVCZGiU2L<(U
z_~PFbR1DCfbjHBOn)*X|7Xyt(?H+(Z6?|v?0lGeMlqpJeYqE>qBvfiqdCvv;jO+$C
z6Y{Kod()_$zr}CSGE+X^1E6!x))vL0w)`&bzU~c88#w#4DrCL3mEbPX<gXzG^rKHv
zjFHoPX@k<w_j#U&a3}*=4830(L%GVOy#z1hq&}$Upkd);B@8xCQ9OLZ9}={HjId7b
zTt!%i0(rU1yR5z3M7K`QXmTJUqFdQO?Dd4=80e$@7*Hb!=++Z{pozj9^EZ}sE(T=I
zR9%N{0_cAi$TTsVDY~UvWa3tP3}o=>u}(C@tz}v%%LTX;zdY18YOlok^zHfe&wkGO
zeD+tesI3|Y7fsrtcA?1kDQ{mlY@y}J?8&o7p(8Hr$cQfhE6;RHt!cbqh4Qukviw3M
z%_T=*AIBe1g`<8SnhM;X7X_n_Wg?o!C8>e#E}b9k#AusupM!vo&I8!Hq1K-U^gS_{
zTwV=iY{m6=OCB9Mso*uroUCPVrzQ<ycfM2<`B5qZ>`<sX?8?+@xy#g3m5(7-^HF0C
zq}uZlL*j{VwR-(d9SJx$vS6`P=Oe&`3VI;^!Zwtxd9sGm{6R08-+3D51Qm{%ry=<@
z8Fx>aF0cvgMfGOQqPkaO_l=(6zL!4VO)&}dS4BcPXF`>bVE=8PGtJZTDx9v*XaD=i
zszwgbA6W`WRCsOc{0*y??m_{KB{cB0B``BAJmo{ESo-Ig$G)W+rlDA2l5%hPKI8g0
z#QNiYFj?_&Bj%T^NP1~tO{+>$0l2fh{#uqkI*JqJ3g&M+v`PX1#}0s_wvlj@`SlG#
z`7|&VMoi4u(l^+A>q!BR+SF}=Ug!EFB=QAlTomSb2_jVk)hWd)*2lH-1x8s@#NXtL
zuA9yR-_nZ(xT@+@UW0LiJ(d0E?kz;>)4zP<m7Qz?ME)M+Gs{EdUoFWQJv1p7O|_h!
z{d;v$UX%lkXZk^&VjK!Ouz*<fQ4-Y|!5`xz(6jpM7Y9I4<^5|t41DeXvmE|+%hAm$
z0Eo(iUC{aa!#W!yv=gTizYP&0IsYSH#(W*v%HMnh%I71%D4w6fV+2P}z@<lepwXZ-
zAdHal-G6@-D{i&^sE914;(Zwo=j2&+4>OwOY+O9l)I$hKzRA$@LvPkE15~T!i>72z
z8CBF(e!Ado$Vze#-^UVb=dHs{PwUvct11Ee2YqmuZrr3{aWi6~9ab|(imsUvpdM4m
z+**s12y@)A26jkBMR`#NtTsOF-%$D;RJBAEwRIbBtR;AN+-H~-dtWJ<Cxw3xwe4v~
zlSlLBEpIX)Mr$ueQc9YX>)yM%HOi_^a+z5qCxJF<E6m^4+r_fUvG(SaKh^K@0xMfC
z4EmTfQGar4nAEe(&4x6Y)$hz4yNFC}FR;J`RNQ!8=u~^?tVPqua`A`1x_k5N4|N4E
z(wsJ7Kgnl9)D!@;Re?z{k^e{>+vL3}KGqJ+YK7!%Sh~_Q=u0Qjx=-D+Ef&`Kb?V8y
z`Ay-#1w&z!kqhwmw-_NXiUnRSW6d7kx@w0<PNp7VI#sm`C2nK03bhhYgU$x4vcuAx
zS4xELer2<+G=ZymhQ%nVD>+ch375x1M?EkN?MlskSkFbB&amsl*5iy=p;=RyNz>vO
z+bmq$wLbN{74PPuA^hS#OWzS^6;~9Oi<I5oOhns?Cu<CXNb}XvF-Ge_G&hW&3*7#8
zHelASs5GkUx&F}jQE7kgN{oy^XNl=Ng1iZWhcXckwR8ViZ(^f*-2jLo{Y6FxM$o>9
zV9`Vh`hWya42vhARE9dXwx1Z4Nw63iM{G}S`WcZK2F!>Th|_Rz{cy3&lsAG&D+W|6
zX8Txm!oBLB^Ka9zNCQ+{N6O_<bT8s|7m$D80-6ftN5o&^P4p68QUV=>m@LWsz0nTv
zT~ze$U9fW{5WYX$MCSs%w-rk6&B&^~m93P;uURpwVOW;BdRVvMJ-W$X)jH9bVOl*=
zdkmhh>K|?Zit9pcUpJkW(5H2e^<;Rb3-9H}8y|PLa}Moy;-JTzeHazVpSCt7UpUsh
zKpgm)2?IMC^*NHHZ$Z-_`&ALhPGXKyw6NRt36_`O#~QDAUFv`b`0E;i8&=D*b{wPk
ztKiu&M(z1Tgj6lpB+xG0W_~83@9Ih)wIAc{*Xbx?_!d=8+nG*Ev-ylK4%88_xz>$0
zCwF&I&yoT49v*WC2u(Y@`b{}7?_M!+Nxzvm))SPf8moD*pIRwdGvPB|7SAP}W1|ef
zw~qUUku-LZ{=PW`#fQxBqZmFTwdQ+KHS}PTWEMkdN`)-RIK9O3s$_x6kCCZABwq<M
zSV(eHtEyfsE%#Ex=N||V*gAKav~l3O$nj`qBFYJUSbA1)KM$@c#6W<aF?P+($bd`z
zPn;{h4)xc+__~F_3}V{aM9f_(VUy_%Ar{m1Nv!Tk)1SO;T@nd$y+`XBU;7S7*Z8*(
zEx)Xz>GFK~+BX3NuFDH+Z*Z$NA0H}Og8d%qxx~u_nk8S65@t{yx0iY)%DR38xK$sQ
zW?3Vk>n^7q&v3i1d+n}=)z)VpsQaFM7$F}$h1BW6FBs~Ssnutr@)iydA4<(Nj`6Y)
z6Cce9j!9xJV9#s6TEU%*bMo))!_C8;G|OhqV>kL-`pemmG5nuQMo85o@6x96<xE7^
ziYIzC;^-&vdE{N%FFJ+fmzl8lZbbK`Lzr~WoQo?H#zQ$zWbZRLCZK)hS_02@AgbS7
z@!iNpE*%89$bImLZh$%jC-Dk*^V<4GT)P|5C{b0m7Akl0mW~IkVzvG&VEs|O7Vz?%
z52BH&fwX7VerT|2Tzbd_#)IQ^#A_^Tnr#hhRLh+^!Rz*-_2YD^Sv$&=p`g0++h!0|
zII^ze|BDEvFO;PKSe6z;7v2ziYdI`-*|_0|)Mej=t9#$weP3V0G-#2#c26~Li$}AH
z4JRT5@%Fntb`$KWvu(Nwaf}q!FY1F?06se!TYUTaQsPIJBH{5Rs;fce>_EEGow(Ul
zWb>~SD)`_;_bD5na!QxmT!?H;E`v!aumMK}bBJesWb~2OT?BNo>OMF=o<-r)R**hW
zcj~5_*0`r`aOafwZD9==G694jo>Kmc0hS4|@chO93}fD{ISMXq#~#<=8%A6;!lDb1
zrduo6<d4eNk#%KebP4t48Vy|ZF_?h%Y(WCJCdt>To@aUVfPe6|InVyM=z3Ic?jJ}$
z>_WJJcvmaena*TDKj#$7*p<h!8*K^MTz~pFL8o=q$)~AAShp&{$4X+QdHIRc;E~CA
z-Q)H9aYu=`aTMG`h3ISf-23nVp#NW`*wVqN{yWXqic9_K%dCTym!{Ue;P6@Vy9saJ
zBV<N;5#*I}AIe#D99b+Mf82k3Cfl|R;e}|}E5L#*>&aJ6Yn$`7RfZd~>KuUalgVOW
z)3Y*K_faF^71uF)wNpX^g8a>a5GKlI=Rw@h&j^=$uJGfac);vw1<=NRT}97IZVLmr
z60h?-+n&&8j>!^6;W1Hama01Mx6V?<z#~-_d2^Ej*pfxZCPZ`ZQFFb+O;NfPhs^-V
zIiS`zRRF1yoAq@BuDb$w#rrFp^Xabl8S}=fSJr78`OSG*wLN)Gk#oAx{<0)RmR%0%
z*NJ&W7hcVy6yD+H@|h_c(8vw1Z)4UfzTxjGGy_UpSH#!rD@*fx3d4+ar!RLO;-2uw
zAdEZd@x;;|m5j%MzIWj4rb`EdY~SEVbA>?bCiS%wI)MEYxeyISm(d!8MI`HIwiUG$
z@s79OUh|4L*}r*3D-)VW_>{2NO)sqv5FLOBeK|=9im6P{xroT<ddO=~wqIG`^rlNq
z_^rqR)9E-v!jbPCbmE*dSz>5mqA=lJmH6dH`kaFLazv8xAat_km~L=_{GzHL(}RAg
zN}RWCeb{Pw6Yeg2QkX4Ajh(yKN3M|oGy^*>w7Bml=-Cq7okTig5Ninmt^ZUh%ilc?
z&OwmR)pTs<(ba5Eq=EEHHM@>VR}D5aRGq`0Z;Gh;4VTTBE1e(eJI)ru>c_X_r8~h-
zo;k{u*32-TnIGva;ML=o?818lm&FsOk1n((Xh4(DPgX32waci3XR<=ab<RR8XvuuK
zWBIM85SGGC{$+3hL-zVG#H%pJV~#UH6TW=Ht7|RHPB%G<5E-1NxsfRNxDNkO;cWh{
z`6cNfvOwtLI<7}X`MY^D*^3Tt`><&Y{%syh5u3vNTA(ZAAjSuA;1%F@bc{YB>Y~Cx
z*!r{CBwC-$r5C25l|5*4Dd;43nurs$z4lRn>gYl7Sh~{&V`Q_3%Q>XT<%YHQ8QWFk
zveO}wjSj6NrLu_o+RY&oVeeeUS!rg33m)-l3V+_TGlF@)8K&#Ms_>QO&2Oul8Wx|w
zEEZRu{!uzN*Qx4#<cF#6FjNZS2mLC^h4W>@JksC?K9j;?$*ygbN5I6R9Oi6lRX;LK
z9Y(ER3&m)Ts2@hr96ZX!uR7K8zqNGVWyN<h(X<$-QI0<y;!>O7JLqlunQ~5%<#+at
zSnyZCmNtkq*Q4XXCtE}V@UuS>8qa<@6kq56_AeT9WnfVMyt0&L8_adQ_M65`Eq#1x
zNmuuzhs4)qNl}btua5p~9NCu*N(^^)5ao`oqrBWKTIV0n*J0lJm;ps^BYY*4@_i?y
z31-T#HpNn7zehLiLU769WUN_Ubv$0~F#sOeLU3<Y6|B}TZ&t5PCv4@=yLoE+zlmkJ
zuv%BpAUmz340IT%6z*J;dAR4~gB)+afY5c!HkE3=@0WdjVrjs_%SL=C%kP0GOtocA
z_6katF_{#f>S(+tK{qVcY}l}FKd94zx)#qQE9C1iO$(0fITo+hs|NIJeDv1oJAumt
zEV5vb-Y(dvZ-oq&LGEo+l|h|fNX~YwA=L-E&IgQ1#Cq(?G2UdIjs`|`QI_nq-J=-3
zNX`AqQggGbsi~<5G)=~XK#DvaCrxg>YEAi>#d24iHDGac^BXuCPXxjL{c|qW?u}3z
z51gVCEy5dF6-ElTOg1vz``v{L$qePs#Xk2C-X(DmaCEbk7s_cGe)`riNV(3SiC-h$
zZ!Ui}kA!EFI#5dY756N?tGW0v(%)JQws)u9L06YHA?U5*gKLp(x#@lv5G!nW7FOY~
zQEF()cA%HreGEz<$Hw2x%ndXY?!j?<(~F)X^EDOu{^EYRHg<0El*pT4lhi37_?VbQ
zul_5jEOERy;iEOsXoKiZ;aiuCo)bkC1CDnAvO_^Go8w0qz85Ce>fW{)Z}3d2>V`;N
z67&v-18-qsDB9QZ9X2p;gO`ujn$N!3`ra;A?nN->_DqpQ^b+Q{?>7ZrW(Zp5$){gF
zm>pzFXO!#Qt!sg2Wq1t94jC!F2d8YOOp{fdSd(99Fq}l?nJ|PgyXjHDYvmqljxJWU
zYj}mi8j~(Rl;%^96H>ySt=9uZ;+DA?_R2e0O5HvK!+TEAJ%{=qF@TBZOam{yXMdUX
z#JRLeOyU;2)KD{qZwG_sAi2YS9G@s;lI=-M$Ki4u8Pg$hY~<$s_$&*lQ`ImpOhRyX
z6-LGS9=X)}Y&TU^wlkgOJ0z5t&PpHSd6=+2t=wJ&Zx^rQZ~>j^xr)bu$;btmqqurt
zTFUQ7GU3?35Q#An$_JJf*6qt<MbmK8Z5-)F<jBzaJu-Y}F9cJrI^ZRF)W&MfHONjA
zzHURieE#k3fHWd#uY)BqP4oEyqtL-?HP7h}m;sc*O%{qC%aZP*W%b84wS)p)dp{VP
z<JAyC4i>5;gA5)Xor*+$W7~0U)Xy@Sd#I%l{#2#k?Jq0H+V@mAB;uPdx)w+3Zb-CI
zml=8Crj<Gky8Lk#8|!`*LNc{4n4MyHo7-8yF(1bU>cg9#gs&?~emem6#c&0lOD?FV
z-fYJ=e`76S-H5yOlz08=6APJym_)(sYaQx}MS;=N>56|rWA=9?J$2TBQGp6`Ln7~&
zcDoDH6<P#e_-P#6^v8GZHFl8(ig5x-gMQ1IoRf7SbA0+6@LnLjbzd=>c9YuGP<g%0
z2HbclLZ2QlQKQ-1X|dfkhV6Gz)&3Skt`q!ez!fxEj$HhLUDu<KeaVJ3cvk`t{11YZ
zD&P@vCEvE0W+p7cp)smJe?aBoh7Bj5XNs=kyg+Lb=1fi1vd0DmnqI1KK0nycZKI7@
zi5e%>G@bo%XY93Yep>`LT(Gl5*kbossSD=Mx>3J5<663s)H7=ij4T&=rJ)OxA)4ap
zKEG`$=3<kfh+(_Mvff$v-f4{m5H{U)dDVOAB=OaRoiY!bPhXE8JUcKn+OUUIz}9C^
zig?SqkA)`Tug-w6P-J{rmlyG=KCK?%bo+`dzLBDL*3iqoZ7Ju{ffHa;#S;0%T&$So
z(Jn9!774{4JCi`j@zj6NJZ=uzEHo+jXxU57&&3tyOA9NluQ)lMOFSa$j32WmE|x2#
zjpIIhYR2I{mOnnM5lx0l5;6Q{9<lHS#j`Jyrt-4tH%E&<*_f>T(P%TiM6-<t=b~s~
zXz)UvAGGYeTi1pk;s#3cuOT~1y0`FEikL>l5AGofs5@p;xa_T0^FF2D=8_tPcspq>
zW=v!Oz8J#i1mOgrP<i8+T^(C^P(?!5h0Q6e)S@@tYAoOTXKr7v7~@dB)JC&afw6-)
zX{aUt%F7&<v;&e0-ds5flRi&e%bc|1I={{P59*j1VfTZq)Kv~<drR`gX&u}4+=Y49
z&OzX6k<^ZUK^Zif8m;Kf!?cHa9(wR^e2=DG*^j+K_m^>(g3sQiaj}n?=U5<b@zX&5
zC<(2u8sRZ#Y|r2Kms59wY;Y?R7$BmTo81U<d*+EVgl$Wu?Ren_;V=?w*Z#Wvz~c_n
zWA@@%p5V>WsA0S2vL{=RQ)dkxwo=td{rYcP&Mbu+{Dnzi?}+zhipO!34SZu7`==%O
zJ)GOm&yn|e<5;K*{%RRciIM^vh6pmkT<*C5I`wz@*3tF5lihdQqm(>-$A@XTEUP!o
z+YVSBbzD8~=<d6FnVy2F@oa8My?vQMZClJY<0me{$vU&fAZ?q904c3$c1n1^X9%Q?
zdzm#*qG#lmCM$0I%%ONF0gIPv_`ta!A?;98`8PCf#k{a;ql}(zw&9a{Vk@@=?1pg=
zY#rbg>OdaSYv2T;7TbSKNPIkM^Ehlp%9Z)C1E@t)t_;zTbO~e&t~qMCglQS+nU+i+
z^6a=13&OS#`vkM{IwBxbKwH{#!U}AkY$rc30M))RMy(#U|EUcPuL7v!F9`>rb~yUh
z>9~@VJ1{ptFkH9XHQDns>wUeAWwrSB9N5lvYbIjJh2UWwGoQt!e>F3p0%<0cNA1?4
zHViJ{(tTcH^PQDbVFQr_hyM7vz@7du<yw4<5dwdVUrX?etu(5^Mdvjj5#=#1AZo&;
z`mfV;fW0}Ne4Db{0Q96udx~)6l;Y;P$&f3*_2s>ItW!Y-^@FL7ZI$(x=hQX`XpP)N
zY5uhNla8Ya>2?0o>8s~h*SPrUri=RYxrX+*-Mr!YeTGug^eukOm*9fh<$dE>0U<Bb
zfSxiEXWL`1oMQ=;=L0|BjsT}ZygylgRj~!a%7{HHSHQmXS;cs-Ur6@CQd(%NMUj^7
z)u&S3#zrk#6FHzgx4?6n|Btk{0E)7W|GiNWL>d%SN=ZRNYH1J!1W7@pLrRctkPhjP
z7M5<LyLUzDMq=qDq*)pkSQgHW%JV$`^PZXW&KzfS7+jWn@3^k-uRdQA#x3(0g_&b>
zpHJhrRawR~&w9$wd_PZ~S2W{$8U?g%3C(|L`yOH;5#dHR$PPL8&7FRhvEOD1;C4bk
zEsopjr3&G(F`AQp+rt-K<|7~TY(R*2tth3^fC8(uf3e%J)9M-EqpYnEg|9?GCk$ce
zK_@-jgK_{MHN~NqKDf4_n+B1v7%DCgS3x#OT=bXpT_0Z8DH!VH_JmX)+IrTpl{Em7
zwAdawNv`Ct(9lZHMUHbS6Bv!fX=$5dRL(a4xRTeA08!(P^Z(v3{<<A#7{AqY*)ZOy
zSgQe77eHR_3&)&|X{1~eB{p!q9Z~M~D+Yd^CYA}${x{77s=|%BV*B}n-NafR1lIg>
zOQy-{hT-(W?#%3jCG9HC;G!ysmHl=lH0@@a4!I@VYq12^@&X+ZvLNfM54}H?z$U8w
z#YG&}W;;VD!p1u{(&7=_ebO%gVv@2eUmPB7$%T9c3>>RVPtbgBmzzTBYW2OSf<C;X
z!dGU$^}*xu6-JZ-nB<^qk_Izo;1N5J7aaAf)O}%wZdpkEprTiJ4n7&@-?yKaMxo!{
zNeeWCljc;p#Y(1CtrNb5+S2!!bLU(Ngr1FkS}m}^&!eduaM{gM1@-4tv~oR8TM$xX
z6}^|!ylE<^)RmRb^1>|5MQ!MoBOs<7tC)oi8$5iAsy;!$U;IIMQpI@31ETQ`&75;5
zVFacRlw90se@*oc8I37`0oWDA8_i=0DQz9@7G$mLg5K-&wsA-Ga464nRALznSpMYP
zznKuvm$Wtf#CD^zj46EDKTTn7+E}J62ufUKNf*JE`cpSx<-6ge*FGSr4KFJ=yvAP-
zvHHq5IzLdc9&LC$YsH}f8-v$?S#OXlYc!u#V(f-Uu1EzEyYbxI9Z-Q~8$x@!-&V<8
zd0ZR1-3I*aiFKaoaMgj>L^+Dqfd&mv6JTy>P37bl4(iMVCN+9ep_=6W%5Xu14CJ+g
z1dp32q_I9h;s6VUJ(X5UsQ)RuC3D9TiyI-OgxMCf<ly%OZ%DhJV%V~LZoV(Vl0h(Z
zZaxIDb=3FZwv(qv@V&tXpH7an<a$u$B7y4zxH$;RAZ=ULDP;lKXQo#uD6G5{`QSr}
z#SvHuWY^4)jN6D(I35?Y!Dlw|Dnq$*mzK^TxdDok0J&g=rGY23D^QCht-EFo3i9~*
zRc!pMJ!-NJbXV7_b8QwD1t}3BCvz+tx3&hhfmI;uyQ-wYJnYf&wU6f){AEdS@+4pB
zz@wCL!(nMQa~m(Bu&o%3K}Ikt!iea`+ThA3x3o8xIg8)TP%T!L_ucFJp1WhL2?c*1
zzs7&D-L~^M=3SW}aD7E$^ldviczl?efaZ(!N;0;fljHF=J7-yqO(4Vh%W$s4RTp5r
z99yjmQ2y!HXZo7F_{5qLASi!SslZi|)m7r!Vm@i*P>{=S4zU=5%=xC{X4toA9Kf3{
zGgpAy!PrU565~7MAHr#ap!$~cWl1XI$l#g*vd406d~-@pb!PYw{P1K%`80^e3HmGa
zeZ~1ARB`P5bH+uIHu-&2|D0bu>IP`5z_XC%ZIb=#0zR`$<<L#FxQ~{a?z-xuFezPh
zZy68*NXwWLo8k4-1e{tP?0+|&aF#1wJo1dKYyhn`xN9;Q0J9FhhxsEEZ(g}v=PWL5
zl-u?oSVScv9N%eTgSTI2^`p&33%1$Q7W^z_x1a$Ciu@n<#Y(Vc6!WE!8?#KZ8G_>@
zuo?ogm2u;P=GIKH@nk;5J&$Q)kBIwz=*cSUG0&Qwfk0u;1WhLJWr{P({+$<m=m`F$
zVh#2=b(8QY&3gHe_rkON9yJ)z%pq43zJ&y<!&nasEPzxf*JjKIjjZ#zr%eeqHLzUW
zbg^F!+dX42`TkEovI)pgh?H%H0FRDb0?WAYZU=a|RVC30a`KBj0irM-@oVL?!b16{
z?0DWfv8TDfuCa;%MiC@DxM7>P4LklDCIWq8MP1VF{Kvz_OL@QyIYyhn5gY>ev)<q4
zJHx)=wqC$jc`|6&Fg@Us$>_R%_ixqATQ?9hG?;)%{dlnb3$utI9$+n4wZ92KqrUc8
zZ1Z)y*N-BSrFO+H)||u|_d;@!<EAfvFnL|o%svtUbzF&uEZ>ayjnpt#Ad5g2z3!>p
z4fs4t*+^X88&hO(VBGzeEb6^OVXhp7##uxF4bN)7W~pcURc-DVs4fu`Jmm!X(w<Ql
z-GGzrFI|@w`8jag-d!a9+t~p4C8xd8RVUd2z%u=(-UoaEOgUV=&*iuO@3?}0aS4CQ
zd!Hl!#m3fb=ngs5mHb5(kWT=eY<E}A-d=hH{`)&G3vuqrTE007DsE{74JVqFeSmYF
zUPu&vvfe+hTEzp#DE<((8X8xy=MsLmqe*K)>YSEu+FHv=)w(GoQ9<oBDX>~XKQE?W
z`u3D0)>7F!kDD+9$E{6Kds{etC6Jtqq0GlHtk<<`u6>|5Q@@*~_<ss}0gjKS&0mk6
zfD7!NpIo-SMGjw<-*%{Of1^OA{B47*v=^?>j<Je`;>Hg%lu*Wf6YT+5zQ~pdx1~9_
zn-L9vPH-1N@;2^uI`w38lk;)J`vRCr#NO8z99Pb77~dp$f6vsva<90Y%COy)wEv$|
z8Ml%EM<$Kr<HtHgsge2l7CfspD!CO6WRvXo5J-1YRtRa5X{{faUi~+npaJQ;2|Dm1
z-HEm0CtpUV(Di}VZpsHEZ8+xk<;w0MQy#?C7%8)_D-<tpb2GNkvUM#!LC+mQ>zsZC
zGv}R6;v$?@_57+;YD*osaiY|yKF*I{opSa}z}kZ28i|9{@cx8h8*r3r!0gC>YB-77
zu9CyF{Lm^#2|8f1xyI+`gSauUg9xR6nl=9+e^zi8i%zMKS7Ra0ImI@gp)W&8)5}`8
z42v8v9?Ugp9XltjAN?<JEJo=!ajdEj5hs~bBz5v7KE-Y~i>I4sA>tQ*R>57)kR`wl
zfz0D2)~7<CDGDe^_5fWA`9ISei~1FGnxd2%oSK#2c2`)sON2Brl6{vE)u?SY*>OCk
z1<KS{pXnOlpiwqp5X$6WSkkTc<V`_c_hKm0nY(>@f71GQZs^k8aQNbnyCK>vlD#y!
zu0;^M=081mM!kJcc)SCkTXQX=L<TTJ-|S1(W;60f-(yfez}l*whRQ1h$o;IA&kHj(
zPw`a=WT^G{;9zzJT(jksFOe<{Q20*^L)1K-eEgDhjK5P{i!oGeWHrEz9Vjr`iQpCQ
z+3f}w@MNadTh{-^p|xpU;?R<Y6o=k<10VQ&SXuzOnW$FN;QVUov|yj_l~=eWwho-E
zG#FiPP#Q%UA%7qeAH^ZMcptTH<tG&)$%oS(?XT)IaQl!i3uHg#b2yqN2(UPhax;bg
z%+O4eF(toY8y@siqH-7>ORV_nCl~R7#D7-gmGIK%wSVq}XA8T7m!m6&dIrE)sJF)-
zxS0@Tq#(8kSpWXS@r3~XmtQfA7@~6uciUl8#}X<_rcXzvopp&gk{<?l_*M3p)0Qz`
zFKjXUZgDgfjv1OXXKFa_^19ZEgw=Vpp!~bUk^Q^F0#7%f#X}J2FrjU+Ex{oRwcn(~
zmc9KECpHMi^JATDK=pYekoNXINa^Utf4?;@&!xH{w&Rs>9w_Aa1;(D(N#FV?d~~9^
zGeUI~7=?*X$zG^boW<u{eVN1_N-R_YI2~TxQyI%Rf{i<^zVXfR>tH;cP5Ia(pWDf`
z)CYAM2tTWZW3x`mRsHxL0cmo!CR&QV8AH@oO~Y%PxfG@W6Ia;o5H&?(&g+MZ&CkhO
zC;a=cJk!j=ajpRm`fXlL!yK!(s-4}dVn(5zgsB|~KJvF@noFxb$7rI@OBBqSaMP&N
zq8<ETzqdSbpwGJc#Lved{PQ!WiiuLjr$xjvWvx3sIrLUWT(qEDmd`Ln5TvZeqFkTC
z8>FQb?!eT!#POX9TXzdo^i#}z+ZHI*lB|UKW>rY>SthVsAr6R-OF|DkALiW>%GmG$
zC1yuCipozOxQ_~ejxJh=*^Wj$od7}`7iZZjTgi5~XEaVK!tg=VaGHSQgJi(-`Gt<y
zOSuy#uUZ)f$5*^jtUamv<rAqt%DW&d>Mk`>35*2IYt!RN0ykHVt;)T6vJ>RN%w3;8
zpo2oP-2HJW`c`Uk7^6%;;qj({I(Y<=UL(EULgD&9lU|yH<?Y!`c8iz@7-t+i^DdwL
zr`t%2mVlmS5K<-mLbvGkJ0dmwiWzIjZhq61)nNIr0DI-FRyLkAX)H~7QdWcRuS%|>
zU!PH}8&&Gv*H#<@uqr1GANAiM%bt|ZcsKGXKlb0ItCGHKZ~l4X`zdDV?EM+dp(?PP
z6909@^Y@Gq)|VqR8O(*<!d;J14;1}B#fVZuK|s5ZMGQie^S;7)5K;wYIExdnKn>n>
zbE^>RJb#q;370xe9^B=mJ{e8(<|nS`Ve&G6Uo;V|z-DsfT1b<%LxenSw$gl8W==Sr
zr=p@d&3}OGTC||*-xD+h+H)*gfp0Ed0|9E^y=^|Y%NyN4zJas?w(Ni#&_StHK<;RW
z+%@A$c&vWetr95`*76oZW@G3c4r&BV-}zbpo#IR-6LLZY-GCq=%xn*{Ay%w=8GTD6
zNk3{hV<JNTwJP#*Ta%{aG*5bJ$Ci*iS0ym)$(Z*_7Nlch{G#K2x+@fzl1dS-k%p6Q
z6Q9?jQbCr>or47JGC{#NSlqVP4$A)+0(MRsCEVJD8NNHbMX&Gz&?0}A(lQ1Em5^8H
zaXno1)|#&yYqq25c&I*3>Ns*aVlvrY;+4%GM)>i7c^vh(h(&o)lgp1{1N41a`#5cc
zO1(?`OH4D{fb2-lH}gt9H}dAyanH;^V}`Z~kY&*t!1U{B^EyaWL^Ns&J7`;C`zzAQ
zD{z2W&z<Bj_CA!ONQ;53#akRQFv$K_qP{KLPo<^Lxit#Shlb<V(We#57No3pA3=;}
z4*#0=V$`(CJ!aji1bjKNzfkLu$WmgCd4QrFp(m|bdzaOyn%@NCn)=S><gY8-Y<cWQ
z?0ML-3yP#BkjulT#);3YB#ff`!ctqV5NJ!QA5~J4RxcCK$f@F4;6nQZV$~d_U(VkF
zq>4(&vp#n$vMqaMkg!&;Exl2fAuD*ZVl120VtI4r%S_wcC^E)F%;AO!4C<C#?kHQI
zyT=bI1cs)-0tToTc_ge~FEcj5gX8|`kHCVRW;d)$i-JW-paFn7yxPoaj4Z%?lq9Lf
zFG~`N*?r!9B}E~$g;_ZM_9f;4D66SVWvx&bT&}0<HY=Yf)W~M3^|F2fcEyzso4!!@
z>Xk7`u+Aci?;AA00kxpbJUn<tJYyI0gthRV4&_vc3~l;OY>D-*ghV|f$*th84k<U)
zM<pTh8ua%Bqx3n&I%4`!yP#XFA9V9C@x;M5dB?4OmI#t{r1cBqJbLM*y0JMPO`7_m
zP?#SASSUitO#KE?eZe<(Ya>lvVs$oRcviJ9W@$Cb9XGUn+m6R+Ez^?q%grhQ_wMhv
z)8b3kt5US@h#8!Do76@W@Q}*s<9A}dl@oj5-Yd4Ft~X|x5Vk(|xr^7jRTG$A@0%4{
zE&vD=ChtC4T%95EzHjb+9h3ibF$!FEF~;Sz$rc@fIjI4(+wWO;h4+4XeotO`^Z8=W
zF9|-VReSy;S(w)QtDKin1-ut0o?n;g4r(;c)`8yp9>wutVi*YLtXcX^oEleeqw_4_
zia8<R<dPNO@Z4PCURB{+fKXUeS}()7?ecuYqjS#n7R7V^A7pj6o8Y)fKWip%;(oxw
zOe<YJl!*o!R?3yBS8G^!N!K~IR@tXh8KqPwi)xG&-=~jwA-$Q!==eBfNGD@R*iIzl
z!<b=D>%u>KGCQRs{&X4XgJA-br@;xAjvT^YFI`5y=VAfy*ya3C{>4%=gMEJ%gMR*8
zW$T~g!m<{f8464e&DkefKG=vCDq0L%32_2YOkAzU=&iRuT(^%8@rxXU0d$B)qg0>|
zJH5Jg!HW80!1Z`QvuZBRaOH)@flPx9(CG1Z2>vwD%_3w+^wOj6WxP<%)PU}=eufku
zEdU>G)k?BGEuUTirh0^u+C;oQOaQHSVJ+nvOzL>Z9wV7D{y#%Ti;6@pzA{0V#JR*T
zq6PZFR&QtNWx(2jNzRc;&E(W=DtA>rmAmO3cpBSTXpH}>qTA>eUcaKW_i6KewaYD;
z*e2VH1T$P-QC`V=2cL_Az4>!)I9pO#vqe)n^0Zv1(gz=0-4qxYoxtGE=g7+lEpxXk
z%U=+%%hvV#cZv|>fSO+YAx(o>q9B)uDzyZL7zR%K?B#^{)t^5v0QkWj|IwBh&Dj4A
z&RxFp|A+OSP0KE?CeW#8R&%wk0=`UuQ8`B4+`mgLhW9=&<M_+)U$1q9Aay$77&5lK
z-O)G$2{d4HUAl7=ytuUJ^Q8WmL4d(>AoKTNSDk_56zykgQT~VVW^x-RKC~DBKt^3g
z)VldE-e0<}yUKmPH0rR_;O9WYn68e;7l<Wa|90G4-Abit(wa0I>I)G~*U}c+_XXOn
zvZh1|P<O?D)IH0PHaPY%*2#P}&%n>qBHi(Sh_@u)?Xdw<T{*_D|5`Vij+YT!0pbKQ
z&RQSM5Tp2H_IA?s!Pmrfe7g!(W{QE%_j8q;q6vXzM1FE@N>*%z9!Z$~y9S_h_lz^n
zS!m+IBAO9oJl~{$NnH8V+_XN3du3ah@+a2nw{|7g)A-7_b}FZ#H6{f(-`xbBeCA%t
zq2gjnUw%^TH)KqrugLuYKnV*4egh)5ps3_i4XzZ57Z>8p(y}0w5$q&X_qt3h1JZ~u
ze^Levz1Bai?!p?l99lv-aLko3s%Miv0Ds1LRB7Kk#pyC@vMtXTi;f1a%=Lb3V+2pf
zd4gRE=6LMQPBjSZQ{w+PM((DM>kQ9)`;+0G|L<WIjP;mW2AEVJb3~<Ne{X%#P|hF)
zzsJB&a_=TeZ)n8H?-e)PwG=;AchFy&HXF+lus`lh0GjT|4HBh5<Gaa2BO&+N?D}!y
zpAC8)+zbysDh`OzVP5%C`{{9`*K^~s+V6^=D5_n>p8MvI1NT1BKj=g4TR><vch+ub
z+wnhL<T;m90Zzc>w6oKR?r=)PuXBL1ID5-y^LJ=<1H2~xH?%tb_9wKuJn2BBIx0bQ
z!9Am^F=D|tB2@<6NpHKTDT4RQ0|6Ga*zvXkRs8f(!BEaFwRG82(m;oV(I4h+muLTV
z?2PpwZ(eDmN>(2;TG$+%&}#a%F6tR=RCUu^!X|2}N!TrUFUB~tRq^~y1t6G5w`#fx
ze$ls``dquQ^lcaD*A{8&j})c}Y27z%(4(kv)JLD@_1Bv#O@NaI3w~72)i2@ogIDdl
zo_9?Gy=4_&?4o4r6|P%!J4M%BG#~-L(-K}YaZ1Bcv)c0Btm_D*xu<1!#e3Zez{NCf
zBd;I+7(wo4JycySnbx~F-rAfha!tj(J?(SB7RWSP-GQc6A6dbIiXdu$QdB3d@y0GS
ztR4VqPlCO7{sUD+_Z^==NrUo2xUS5)4+N<RMY1l<<m>K9%voG<1)3IvQZCWcv&S+A
z*K9zl_w=k6V{^NQ=ECZ?o6?{2i<3lRhc-Q{x^;k=hf_IW{VH7vQ&12vKa;qDbn#t0
zm!j~|d(AhcfAI}i%V!+f)2j@LCI&u@=A9xW--})LZAlv(zVdA=s0Cxj9&!nH%wzTh
zS=nBMS#B{D8k}ZZ5PFGZ$Z*)K9mRM9^QPE&?E5hK+Z%6)B9N$l2K8hQ%TU&jodAqF
zr%CeR#0My7bEB**OU!xxh+EM5vgc~QsGyYAlSi}LkwjmmCdQ~@4uBkR+`4qMckoJR
zHvy{PCW%sG=+;;Mlb+>BG94|ggD&9+`RR@A5v|!tT{+jvHOA>BXhO4je`cRt>7$Je
z+CWscBh@d5q?T5+emK0c$?Z*m;h#L1w{EQz=_cR1V{s?YPk#|SX3y=3j@eBl*o(4!
z|MxuD?n4pWJH4|#Ni!rqG~gJ%fn01u@yGX*iOl^m1*7^rWBj6`;W$f>L%=xensj2A
zDbq7%-=~Z%=R5VflV{>u3T+%Ub>2Drkjt(nLV%jVs5{KFLD;`htJeCtuF9U+^yLJN
z3w)j~JNBbym8(SW<2YG0z!p$dM}P_7C8t+3c*xO_avk%_`0aNzg#+GZ5^2+p-DK#z
z7VPP@)#$PQ@)r1ybH>s~lb<yv1+xu;Y2e?oSH%?4LeIi`Hiaj*VF9jGf7o~XEmG(8
zKBgRN9ph<ym21Aqf>tSmec0was|eVd5TV457UKOi<_@yU6F-952BzM+FwSXp?q=Pn
z2)|u5UfNg9TylbHlxHazeiQP`Asub7&1PnqQX>**7=!%S@x7AQYb4#UtNwyKd9ktW
zoDFy>SVQMU@`U3%ujYxTsogzhfU)Azos0#zEurweZSp^R)yL#xo%J?jGT_ntff7M*
z1{GpWZ1KV29j^a$?X+wP+2ZywMf=bUV}?@RBw`h->h_gbR{}D~A#<@>qPkE9^?PJZ
zKjR#O)4k~LKipxmYI9)RJ7i~2cVOB0K&9L-nNDTwMdZOWO+aQzog<4S$oJHlCuZ2I
z=bW@pxZd$OjJJzz7OrU8JH>dkQMKZ*a7PI-hla%9#NU($7);AmrB3UDmHVbpJqi;?
z^~u{zoVs;c8NOccy=EtIYeL-AHbt10#)0MF0TF_Xcmm2F@Y}{l@4B{es{6p$%RZ#J
zy=6CQy4A=<P`<A|sJS+~+Hvv-$llCuAb-6tav|ghzAr3Vh)NvgZB5i6Kbd8+Ea?MC
z-l3NyZ|8jf!|^L=uu8sAyVq_m`l&|9#;33h`_oEm*lUeTXtzitzro6(t*>#P7Asnt
zgmq8G%dwkLN)Bg9@ufkfUZr;dO%*Oy#6r#{Uxk7+R{6DhqZ^Cvk6g&qL#uMW`3bJa
zs@~phwKVF^0d|vVO`OPlE}@PlDoKr;ZM@=&8X`JZu9|v!7*<j<Q`!QV{iCmnHf}1@
z47l{TjI9qiO8ywnNJY%8Va{#~CF&2=yvw-{?OPWPpx=T?jS04%r7XE}6V$}7^VHZr
zziW-YE%~3b7O$%|j4EXt25JJ}%<!tT7fHxUgi!mGJ-{q%)@^zT0PNf}X9u{d{B|9q
z7CnBl6EBo@d|xMaQF?^=SAltZKJt9>1Q?t1vQ*yHYJA{gIxP`^UZ#GzDg58?^qTXK
zZMEm4*-i{FIu|ofqAYJ1h){_=fe4kDB8cnJlr;Gnsd;3p90UCg(X|Cq?G(8(mlczP
zPtwPqZoFO4_Ftimz24&d`fYKeditzRXl^mc*+0hL0B&q8$SpPYNW+1rofcn}!0f|D
ztYXfhKCU8mk5L=S0|*-GK~WHoz7kawFejoJm<{C?vzBTZ0|xWXB1qk~UkTJ;fGbjr
zK3-HCyXvV;sQ}IyL$l62+X)5RUn~aZ$BO{-l-~inlT48NR!j{JnScm|ecs#Q*!h3&
z4g(XcBO!<1EgE-LXhmID$}N*cHaz#k=$v%**zE=?wGr%g9C-vMay<p1%{^-^(Z|F7
ziQ9Y({0zP#RF{?l9zmF&1=(V&{k`oPlI{_ZHQW3SbAhbiJEOS)a6>@UHBs~(-ilI8
z-%Otd#l5_6WJH0wSze#OxhES-c=r%zw;<@cb(+%{9$@pQ3bCmWt0<&uf?qn@z+AKq
zMm9C8;ich4j?yPWihHiyI5pjCgMWY*2_~n_Jyj85@)HVN%@>ng@Y0j&dT9pd*RGz$
zg;Tc`#B5<dUc(h<2cGM>jIR{CDTWqnP1bk*J~i>kN<5g-5xvC9Rt(7TbLb>oSA*2!
z7XUt|2cA-zzeM8$GBHtJ5svpYQQfm3FC09TlyU34%WV^YR^Da7cNQ!`0X8OWn3zSK
z&%%GCkXvbM+;VXbf3+ds(C@pm&~-6>GwBCz5O|A^ji13!HwCb5{vQ!^5$u)4Lzane
zrr|)Epn+W8D2W3Rs~SKnUh4vQ)>Cr5TZrJtY0hO&axVZZB<<#{Il3P8Yw-*40H`E|
z`iVY*Fn=OcgNhr1v7U@)c%y4bb5v04jvzy@fakM6D0M6oY$N@#Irr?v-wV{YH&(xq
za`4W@QnpI-QBNqScby19lLYOtnZ5m_=|Z(O&Xba>V(Y9}WjW&K#l&-CLyEJcfoaG!
zkg-sgu!eqzJFu-XDyHulC7?7jUbSh;PTN1<=3^pjb?!&(JVsfhWQ(utk7>oS*I1dY
z(EbTvX9g@v9<>1NfFIK+-FvrbEPOm?oVxM}2okmTrGnx;)2BAkhN&c6GaP@G`2)zc
z6}}AA6ws0ZI1GU($es{&h1bEi%%JkcM{ex4t^ELvE!3IIIr;{E4VGH0w+bfV_a0_X
zt5W{}0auA<FB2o$jGzTA%*30)^y&*PDN37aYVA>EMVurQLCQf~f};(LEp#NREzE^@
zhtdF2d6!MGy+BdU4T#Z4nA?82;$k-%wI`TAX`zi#--0Ym&P6Fp(b*Ut#00<YFOT~%
zvBByB-8S---1cuDX|*)akP=V?p{k=Y2uyI|@mA!q4C&4UU1+}AjtDvwc_i-e#suom
zLF_G#%xx{(lQ}>}&H>>M@qUN<vsnuHlv_e?uTE8UVjo~fuT4I&6|=za--?3BP&uT@
zNA&bN-#wQUSS<e&b$z>;QPOpvjJ%QS#PUNq+&&n1UV#Cn@iz>$HItN%#P82|<t?zt
zjY|N-O3=Xd2D|(i&k&~I{|5E4j@y8z6>l;!BF%aRxD8)r;Ix)6rsc1$Kb#b>Dvsuf
zO4n-ACS@fMD1N2y-@l660I?;JZYZy*6U}2pERe7gM7bMyC5$g){UV(~I?AvQ(BhZ!
z6f|udQ0+oemAt5u#j+Dm6G_zHd&p}Hezk0ZZ%j@m<tkQ~XEoVAOr2I_r^rseEb9VY
zg2Qt*@%wMT2;4&BlWvw;7`@e~Utud_B4xZBu+$6$P4P+F_c>_u+rXJ?P`pqOhdlM*
zW(2GAtJ<{Y+!o(2nVi~A3v7M-Xo)TRcl(2g#w|40hHKB_co8XU0LqxP^idaZ^Vv0*
zW6P?$#stzvZz%0UNPA~-BNqxn3cinvfKDeY?EWT&Y+eV!J2L9}MwQ(+yEkhKPCWM`
zFPB4kVk{AQ5K&xI!<R-YH*81x7CWC^Od<+_VX(&hW9p{}Bi|fD^xJN8C<fL0f0GHv
zc0jrGWYy+bVyQ~B@T+_k$T((C+;-Cm;gqi5iZ<|{g988JLB0UPkb94=cVELAws#D?
z^%%=t9LQ(ZhiDAeIsczt>Vw(SHlI>RP=xdI?j_2e;NRpQP&B;?(J!B>GXMXiC;l2A
z{dwAc|NDP0>OK=)rf~4OupexH;xSR;COJIcu$A_6ysqS*bR+OaK-wh1Ubbw+?87#<
zDxaQtO(GBndO2i@{n9`Cy!Fco|DRv^j1RE#%mg55qUJ+0B4QZOs^1n!lY4E|+#J$g
zhFlFZF%9rbW=9jm<KCH&3_a)8h>`1M`rxK>o9BzE1ZPE4r$aXUa5BkyP*8*XfO$<@
zB9XXh`%@yFI#}lV`F;5~(she^cohCykejb@nN!lo%@>75kvwrQGBcMPnWdS5{@tAh
z_N&>*7YTKC5f^`lCu=OSsU`E$__yu7F8tdW)f+XAuq~<XeDJC+u*SZ37-;n$h|9<6
z-<*@jyInxybRf3c*zyuWGI#}?7aX)Q?`gdj*A^W4^XB{s=|NmL){Im4o42s%Q#Z5F
zsn))2&d*5^PCIpTo<-z5OIYZ(PbPPQ0CVtU-2L<ZC=;*|6z&AT@Z;XDm+vD@gS=q*
zRSYq9uK{lOg*9mm1=l+gR*IrPVdVS%jXg@Oqrg@{eq~Zt<$Q=G_Zp!N;0$|p)(^`5
zRBW@`i4jY{F%AL8*wVqY75sIJ(!BP)ouug6MPhtSgsG9i*c3=f^LHvZ&~IlSKacu7
z52prTU9@8D%-f>7h&GWcYmGy$9tp{O0+1Ciuuh$gcf{x{ZNwBV^8~0=AWtCqu)$rf
zM`7Z3_Miq`18nIXqAJ0P)OFkC9jLFAyUH_WgQLMRy1Bsa7Kb%s?C7+<5bht6Y~Qhp
z`+#(%_$zhVG9}*Fu_KA}WYsXuPZDg<p{3dfuT|QAI}*~><f(_Qhw32G2E*%Fyo!}_
z7FZ$~X5oGJ><7%DBH4hkS)+TY%w1$_RsM(%(;`P{^3Fq>8FwQ&9-IFElxi2JlF`$@
zlL^^4TOG{&#tGAl`5Ks+Oly<5DhgHA3N>ga_7)Op7R;cAM#R8N;TgE~r*4}BweeWB
zG&5JkO)oYpIBB5p{jAjCL4nVR@;hdJpwk!en7CPm)>;Q>AH}wKclMJLr>*;0d%T=F
zTKek{U^4Jiq@K3*R=JRu3SifI3}$&!8gUwI#^kDyi7fg(AC?l9!nGBayBy*m@N{G^
zta~HL(G(ySb{RChnBfV)1qbpJB}zRN9I1@z;jy{8&Vw4lqE1rWaXWrJ21Gs-Udopw
zLA@Ve)Ct!BdsXW*`qOk~FiS7Kk>FHbH$u4UJmnVaOH%q>mv4@)#r+p-zF)5ju{{hA
z>N@i38lXe%&*w9(3UJ~V05{Wa)mq7lE*@gqiZo!J_;9Ln3pV8BE7_FbrQ2F%t1G57
z2^S!UzjHwHn7Gs|U{ou$u+IO}_<)CN(<==mWssxV*Gc<XDKBYrHSlQ-U9OnRaCKZ<
z#kJ5k27LA$oQq7c4!(3=M!%_rg#MSDn9QQwl(S#2VoKtCad5q;^2!FKX(u34v>P-P
zcUILFgw4A-qoTxSHANoy?2G(1g%#;Bc=Hd5=pY>!gW^oz`BQvNS8sKt-EDqzgyM||
zKKu-kEvtJ40J8A$QinF1ivWvO;c5ZbeT$Dz$p+(Fl*<_<(n%*a`2-%aoTMmiX$^2B
zP%0sPTWqJ~^IJS-oD%XoQY*$A-_iQ(S-ip!dWiB&RA5kV%I=~}Yv|UakQj*^&H3pt
zWp<>iO{95rZAgy?C)6hx#?yD2xPOYc%Q>RJETls$<2P;wpIjw0LwrzygIbRQxuA^A
zN6p_!tvzXT+e=?Zt?uTUPl;w4$Rw5^#ABFVT&L3+*+7I@jT&(oMwK+E3^QkWK1Z}h
zKVfZoGnz_T4hMLG%%^nEu%?-(o~saq92v*g5Ht%C%7q~9*EbTJ`}fCiX57o9{AN9w
z%#^6L(AYT^aZIVx=n6BuxSB6wnOXy}`-lVUBP3uU3#p)-+CJkfeZf%@eAO0&;Tt6B
zTAm^K;N!ytcBOl*s(t|Rg_v~X(Pee%lK4U!<i=!`aU#Lx@*zo|Smn;A?tSKvkqiY8
zDz4?wh@7|OJM0Dyn@KEL+53_vC%<2_(ixw7QFkX5%7Pw3+&RXu<2Tq=&`Y9vJ7z7(
z2(EM*xB^h8RWKHDu*lq;u{@`)rD>HviTCoAWKxC~MW4P{`&OI=L#CY8(Xlld62ro@
zyrNVI=^5N#6gE+zjlLXDUB=k~x!*~GskS}hoLFKbX9%N@4>KhU#@0B7ZB?9QGCjz7
z0FHu)0z8DOnJ#!yX~6oY0O5QWFr-x?<FqH3475}d-@bavW|$TYOHc0XbcZqCJ$kmC
z{DqT+{hYkV2};IUPW`h3C!RNns4LhWXx)5h?a0j&u55Js@nskZk+k8B)|bH#$=$3a
z-7pMlXk#QuM(UF*+1#cJ6sPrPVu<|&SGC+!`v4fV&IpgY?bf6=BfI6l?%220Rx<zi
zb)Z#4r~cil*i`C6m%UQK5WX`F06*3N=ZBwv=4{C<ww4M;KeAzW;YzQVa9C#r93==H
zrY1^K6Xv7|B$qe(9!;68#MT;bn$lkSaTO(R{DB`1g&oN>^tZ0&f2+Bo*c2?hj(kDA
zu~@ZFpC>FApZe=cQh09>ak5M1cy5*!KUvUJ-3ZL#-h}Zmj=gv)sm^-B(>U>f6>u><
z>pJG2<#`7}ArySuwEUqd#8@&udV!SnQS6jK%NvPr41MP&kQ*E7zQ!nT_Qq@OpyD@C
zVr$zYjSp)YSLDMSf^_^OBLbq#<NOCsR&k@G@xsSncQ>p&qfb5XL4!{H-T~koNFqHF
zdwYD^={<~^LcYCn-WVj3m@l@Sx6ot>!AHF4bZ-`Da{4GsxAETHLH5Ifb5sCJAWyYp
zcitC!iQGG~b^S(r{ZpX+m-}U~s*yz^)L|}vD8nRQQz_BE+swzfFP85wxK(R^JJtHN
zzGzw_TnI3pDHWG13Orb%MGE(fZqXvoGX?)gq+E*&qzm4e=Ru5TOtD|!_~<d31J}1=
zVPlK}hYacwsQy{2rgE53HQpleP8Q>M-wQ^ikwb$!doE49FKAb$ANe8Ox>}s;^FKOd
zpn#8S>H8_T&>G7O?PEKzvPFzvzuf~LPH6~X(3w6>`5BJ$_U-J!;mUIb;^orvScaDG
zlLeb2!SEZtq03CaqwCKe5ATx;c8Izv?vd8pb(iJz)$(Nb8WxBUaPE~nNU%`e*^m?~
zesFbz9|(#?@QBc)8s}|V?k9YEIvR{_Z*BbzCIx(KKV@C|*k-5P?uy0>0U$;5?eCi{
zZ!DU>(X{^QNajyI?ANBi_GURW+09(y%NA0YlvS<UefI<37a%q^Xq9+d?gPK$^UT4;
z4KM`wR9nfKlu`#zk6=)D0`n}6B*@=fhOnP@rGy7)k<Wr+bjw7_q+Kop&K@oGnW<X#
zck`cRguHg0phX`Ej!${g%$NrWEDr*!oT>{-c`MJq#k`Skm_LJozYoYiO#V6YfMRLB
zwzNpCx`f@XG#S^L{y_||0O_JGYMn1cErl=qIxfN=KEb>xE^7UR{`u`&n8B>yIT+r4
zR!)C>lO#m>q2%qGB%g5a+@_U$9pY@y^U~fdV!<?GIo}ile0}1lTU*K9s9X9T;mQWO
zx1u3@;MpNWZSvj(HtXtYQe%<wvb6f9Mk-a$_j91gua$+Q+fO}Y*Fo2?)*qS5GKx+)
z^~YN-greLM1^K}K@(8KJXjuwuu=W)6d{2d%$&vUF78uIJ2A*cBM6Tlxd%1Bw`HNhx
zH8VQt?I#mx_*@`aOP_F`)V?J=bGLLT@~G}U?0`ybsqj82#lsh*`Z5%iHjArvQhYJu
zRy0Dl(pcM6nvx?Zg8K0K1Ab`<ep9AL__OQjB}ot6%vk6i8^#F5ebXi@v=r%RUwADZ
zdyl^VN2k;~t@naCjxqftyIpKKqhrf5)K-jPO!tN-v1A<hh001d3EAP#C(gezm%cUi
zvKht?U)JU8V-8^xCzaQle*BH3Ym-zJ^#!Mb-mE#3YlMsVBEoX0VMXE8=PkW6wTY7=
zW!aHgeT3-vvjmP5Vr8*E_hH?OzPoCZcM%)Yf@BC9?6<$JJ`$&_U-L9{z9^J-0Mc`y
zyU{*q9&Vqg3G!?2N<YaFF8(;XF)ogpR+ggL=QXq9_@;be%T|gU$b?cg7`-Mhv1rCE
ztNUn`D4uT9bBzofIMjgm)_53eQSRJHgC%cIJk=fYIj)N1%TcGrX{@RFey))~xr5(D
zCa9AYWqbpx`+RPX?GA~4a8;a|<PgoYof_j8l@tp|YVTV!uqy~cr)W>nQhGeOkb8}c
zZkUXN59}z?;7+?Az9IK4+uo~_%TjGhB)%gCLE~;nv7+3QqUN|9ga7%spRVC$^Ek8{
z(=T*SH?9NxdUmM36)M_E`|VGYq77l)Zpzsg^-*%|^4@H0ZqLyZY}EV-R%9^#=L}6+
zyc^l3HqfGV+osz7zWCVN!=V_sr<-J`AI~-p9skm{9^Sh~cJMM1+Y@2o*I()X@e@YR
zTw~w))4^M1e;<A40N|&7R64xw+|zdcQCs!k8rk#>O!Yti>n-%c`zBi4`{t_H^}oU2
zx0FYJ4-a@rx`g*XEx*r#G5(1@yh*P8qsTBAu6`8wfmfL7>0HV}G1Yor*|Ypd)GvkB
z@%|o20O2*Ve7Vi@C}(Kclf;k9ihgfLuDsasXs)w*)pzd|j6|GwPFGyG%z7k+XNc5#
zzQg+aJH5sE8k|^42`;2X&dZN!G^~?*hEJ0Cilv}?RG7d+QP6u4pMb*CHxIYcX-hdR
z@gqV%?OnNIk~QzSpjhy|s~Y0#m-bpddoBq*4q*Job-K%Cm47;4faE{>=X`zBvskVz
zY88)^b%m(cma49EK`c9_;Z$*y_S<TZ8%Fu@x7?18gmpvKW`;h!$?Iq~>o!gM?!c&+
zijI*{i3y(BI3DjlQdo90l+7jRMR5_=FI67HD?=Ce${4oI3p`Fz-Js<xk$eL=M?cQd
zb$Ed*T-9|+vPk*m-rX7NBN45YiiKGfzO@)LwOzRN@{)6W7t&8CFw8HecQ=gRx)RnY
z_PRO6xSKL%Y9(Mrq$atPdu1$hb^FfiG%&Fy^ns7#(WXRvreUsim{-lv^yyHdTv%j>
z_z%8aXx4f1Or!B$%o{Pis*#D)`A<#isKF#>IW_1E>My+9`<($W9ua?1eHIqoAZ$5q
zSQ$D2`?0!T+%+y`<-WJ0oeG<M{}V6X2fvJm4~t^~{&br4t_fCA-|obueBQmjY`<BA
z*M;Wt2kyiwJ>#QJh}Ctbu8k}$(1TjYes$v5&1>-l+S%-R43-blceN%byR+r5vodkr
zbslWOEqDv|f>wet)z3S-GoZ8aH5KCCCi-8M;5S&4PDclCs>eSdw|&YWqctUzaM%ll
zXM22)v2Oo9WswFODa|eYt7h*J2yt1qS9|?3iTRSqxQu;f6;|e`&F3!)-%ZIZ?>!@2
z#_j;#%V0FEV$PQe*V>wBMpHN(yM6;ZkBc(}@99A2>Rti<0_C}tyqG0+51gtV%lV>Y
z1$&HJ-1`tsdcSnC>7~gu%1bBPEXKHG9UD6AK*)m2RTG8xO;6g9t7{5)Z`hvQll!rH
zaZFa}=uwl6B0Mygrj<8T)|-Ldo{nWdI|Z$8Lf&@G->+9xz$!!?y(5=Oe~4Q35R|l7
z%R0)mylZFGaLQKrBBdT9l=(;ZN=n^8lhcJR9L*TFo3{7!OsKf_aVc8XwY;1zYq6_0
zI$x1-$?|MNH`H`pZR(2Kj(3eOH#y=-N%wfWhM-hph?;yiw3gHo!j2v)kjq|7B{VWR
z1q;Jzrpvp4#ippu<%J2MB^_He+tl55sb2P_g2clpAyAEFr2JBZ>X^mDcOJ-EG)kt7
z;Ys3B`7p#Nz7HAMJK77?iRB=*F_Sm<yMb#RbIlfp<E7m>B4?$MH<?KFW<c@*>aM;9
zd;ssD1JYfjk;{S-O{l3^8-vTff%Nza5FScSeMQxqZPXTdWAsF3r_qa7WXo;AB@a{f
zvQ<W`6Cr<*zd+J;F{BZWIJ&m$Qc!Keh|St?H>(><-FF{NKd`ua_9SuN>eJIN0sP-8
z=!UJBbL_6gBdHP$UX4GKe(R)>+Y5Q3VKp};v;%cY6<8L0=k)e@d{5qe9vrGL>vc8n
z91*br6qzD|nl3PpE2@{;^+1==X;WF%oZ*IH@YYIhF}&jXb=DJU%X9by@+}2dwHvcJ
z#6$JfsP)7@2k2?P7<u13e0UksYjsv}o@3-K8MA#*KJxPbraXmW8$Fqh?7+s8^}lQ9
zw=fHx`Gs3X-(!1!Rzz43Q(bCS*C?W1p$=Kern3pbR+rO3Onm&X2tLWWL#dBb0<ZXZ
zPv6zaVwjv+6>9v?vx@2~7Cz{)ps`^bO&numjK;YpOH;kLX}qr!S6)q;I0Jo!EGj|j
z7uD$n*u}U}EH#`kB5Ux@m$E|W!IRL*>8KA)LuYZqX;xMrX`bLc@9p$jxVCGan96yf
zAUN;Jo)X9Uj<+`Xm)F2`R(AC=PaEda;fZ^&y^q@x)QZdNMgD(I%tvy4nfT9Ia1^s)
zV)d<N%hG%}mj~6KMgN-zj4~!wtKT|KNj9>N-9(RYDAZkTseDTDv`?V|eyo50E5mmd
z^Y;}dSW%Je*-3;SVxR6E30pON{9$o6BciUNAzDfCiAh#L^&(Eq*42&}MDFiWpBUUP
zG8r8_P*Z7y-G!LF$2hRbxaoKw-cSZky5P55eogyLd7_HjWI=v*Cg33Dko}v%jfZMn
z&2=nsUErHP+@JT$t3i!{Z*%Q3oL;QLWv9}nZ5xM%M@n>#3N6qd{yD+(Pb=I_-iK1^
zbdF?dMfZe7HOxs3sqd`5bZhdQqj{magxXXmJ5tOMgbQXF8$mXypVk_hUS~Bkq&_lu
zCcG=2+Fq%2+eql!v0|gGixt_oAAPv6E-o`2qVL<vxfLxLyVhI@6Bzlbo{{7p-VFHH
ztGh5i3C6=r6)8yc3oOPP%mz5VmeM3}L0v0F)P4KqRiU>t9CPu?viWjob2ik5{(01+
zHX#eH^}Fk_6txa#g}0dZ<lAMMON>k6({6I#19m!tcr3SiuT2pr?!rBneeDk~y(X~>
zhs9UeD1?J_wOqh7pGzBazW3@<TZxek%(<aEsz;n6(Xxi&+nvMHxyw&biTFOsz&;<<
zeOBJfS7Ltyi6$%)cG;21CA6g8czFV~E;cgx5Ib}@9X5BWx~k&cc5f=^)Uc#zG*J=s
z%r5H`e{JJR21HU&^hQR#4OI)_ilOfPVmO3H_Cy#yv=lR~4xr{CYy6gJ_@gB6kW%l6
z7Oc7GNcH3n_Fkn9ROvk=gr`p72K%6+^RH%wR++A|I{RSA(kDK5V&EXN-&hZ~?uyZA
zY;pH^1{<SrANgB`!Z;wk7T@^Ho8hOnVDTWZZ>OEihHm}uZxgC|@{m-^WFT@QUPNq`
zQm6a-k@JNcQN-zA4de~^f0RZx+hHU9MhF;90bfLO$5jBn%6FZ$&Y98aqO(!2Jx)$!
zKd;`6#%TQ^bDq7;g3J}FzAo}s*<8PBykd^i!OP`%bjSWy2Oil!(rYVns9MF<@%ZJs
zm9TxYExY21`E1n#bhk4Rp1gMs*M*ku798P?dKfz|`e3FZgrjH>qX)g{%XXy%G4Rqm
zrqs+ZvPY|X*%Al?4xYfpIeD(d4@D-}z8ED`Z0@lBarwhn+7$!C+35ePlNv8|()QMW
zb&}xvl}<V>)-yIulXCn&NTnV&wYcEh>}hQmsXCMifo!gNvWOn!G%8%xwaJyen&xP`
z!<b(<p`dM4yTO|4eZ;nB<@SVN!d+x)gG`znI^@EWn*&-Y2MJ-QU1;h%XT6X(bc$!T
zcd$yMNJWm@LHd)mLMgDNfHgcOE&KE&GY4-YkbKVo@q%0aQ2oR5`26-5lJP^{9;4v7
zN!}r4Ba(`)Y(|;){#54@MUI=UTbk_O%b@>xLZr0R*pp2ROFbbE-_62_6ZP?)2Lx%|
zYpZ{QwJ*Fk^q7yIug8I{ezCIJ{l^;K^ZV4oU!rygsB&MAnsAvxyEmgeW-t|^VJ~Yf
zBy2oZ%(;tkZCI)5Oy4QUPVziktgWh!lL6z6q;<Y=kNTw}T|Ie*N|~Ymw)gTa0s63}
zDh@}}=(@_MQ)>O)WGz7bWE1vvdL8we!?;9Z2I$sG)%Ju%a@T}>5X(O!x9^N<)ckeP
z#lWkfN(32pi`iAm^vaR$(^>Fs!31+3JWt#9ptPrKz>qCERUy+(e>ZP=vN{G_>XwCq
zqL*HUo;EC`!(HCTVPWiPlMfX<K|c<HJ<y$wg2?TUo;G5nbLw!hsNJu9Zbt3R7cj{r
zkjD2y&mAFoWO~IzAR4lrSZG0Z(OdH=(KN;<Xvw6TsHWCEYcK2@)8IH|_}+lY%C+*=
zTEaNBs*<;Nla~9(;*7XS8YQ26v8M~8V)UuUfoMu@&B1H8#T^^~5~SB18cubeO3@@<
z!ZLZpdJ%3i>8rs)NbQku)1G>w>MJD6j4$EnaW@~E^@wvWv}+&jJHh*R<-z()u_4=W
zSIyXBQBF6{@RLH*p8R=pis^oHnqkC?Y6<nw-GE-|Nh!?_T3j3P2kCgA56;Yr+?CiQ
zhEqr5Z?@Ep`>;sHoV})?X_#6S(HhhzKBqxfA}E{lj|if>6hZFJkIw2o7bH$|y{YD=
zsOppy`XhocDGc<C$5`FsJkn&frX0*M7+FC28Bp6}2i&qdwvH2f5XHUeDl_&yp}iwo
z!ASGOV&3=cYo`&pGurn(b-J{jaQ;|!VPGJ`D9ESLR<#d;I(&<H_kWy8)@rIuV944)
zkoLW{K{r^1xDfF=4UHH*@;y6{#h<TnC~kFiv%AGjI5YIzBx@26uio$Fbbe&?UdpK6
zJ;`qylK$*H`gU2>%8DyI^D<%=jte$&>d!)Ql^Z6EX>Q2Cfz$uyO8My9ZM)UE-iOZO
zeW>Tf)iTL8MMYH1<j#0LEjKSfpySuw*`@z3n)=`BZCk>s+s*6=SKeEF%I07I<h=X#
z{1-Y|e@7dnvCWPBD>Wo|zt}CCpXxfR6G+0YSLc?S*woLqEjRDB5!I&B;6UqTMP{9D
z)aJO9dFU^A0xNb6a^`%>(i1L-`*O<p070g>p(PYGfqv3>pqV{eDfl-a81%9lrQxn|
z{;iS_RM9kx93CSPwh4Vm&u{^PB}OEu+p!ir{aYE;|ErANom?s-cB;F8GWx}*tjO-a
zuPyjQgVw@Ipz9*_rn<I`A$m?(H?Q(#$<pC$uyAVcsD@ce516&927**+=i0z^`YPLa
zmN0{AuMu}ct)40*V#tnt6?Q_6HBHDIS^r)ApN=Ojw%GS{g$M*|e*C?f1nK?~I{tG*
z8YAS!dLQlr`_=xnPg*}WxWbx5B88Ni-+@@Rsf!9&(<Z|EVi=mgUUV1M3D-YP!+qW~
z-79g6#<HvR{nWe48wjVT@L!2DzaZ?3TI2vmKOX(;;KpH`ypSQ|tOVRr8m_T@ofTAP
z&oCTXsGQiO@w7;NWlzC9I6`tsXYxE2Vgj;Gnyl`Lb%n2hR`e!MsY~VS7Bsr$V9s|C
zM_1D4fev*dM<X@epQp^bk?TL2iJ>Fx;!-nxng6Ys1{?Gop2uyAP}x%^VaQH?2|RCa
z!bdazz;1*FtVxJYDDa;TTlJ=bt9NYnzdIhtX_EWzG{zQK*qEMydyPSlKE%Nnb6Vws
z=x>cMlHOTY*Ub{Kl=jZ4c(;)&0?mhbYImQRbcv|nS(3Nrk!*5YJJJarzQd-Pw<`1%
zBUFcG2%NV@SnV$HH!J#Wu6t`OC4u6JLjA5p2UIYM7Vajv%+GU;tdbShgErhQpZW>q
zD(YVBvoD}Rn-OAV-&oKBxezti?103yGPdGxb8}Nqd$OitJo}%5iE$GUO!C416-+(>
z=h*7h!pHQ(fuP94tTTA+TP|){U$btC`%^)ad}_77g_5#TLKj|~bGU*Bd-nFd1}6#~
z(XMJ(b<@$&l?w112&Y@_@G5L8<VnQAHgPXA#R|Qf`g;Wjlyth%J^y2HZ)Ynxsd#hr
z3O9m$?MLm(k-=N6P^(a#+1}cKM}ODCboi~Zv*KH2647QBYI_EpwKeeLc;Sg^@HPzF
z(gXT(d)V>z2DY=1`Blp2{4Efd1Al4za)xhH=}3<@Lc6FsIfR@>8hY>jebc!kBK;FB
zV*PR!8Y6?A&DVA_6)(HMS+W}A=UZL_Ix;s{sgAm@RL}J7f3(Z*w9()9z`%(8Cl!N%
z@tFAE&IhG`rN{5Q)PK`V9EEx(M?D@-uH@d_7LmTse{)SLK%NMZfO<PCROwCz8=kIM
zxwoTs%xJ%?YiqM05CD}lK389+fbB7!PurZYi?%otfB1qTfZaR`m^=T5yp)@F9)9zk
zOQeWfx?rx(x*$K(Z^4UyG&+FY+l}qIcXH9$9eF{H#$UX6Wr>7#Eq3i{Ywo9~W}B(K
zsp;-jbS=NotB7BwnN4f<?V0j=v=X1~x)=fTa$vU3VUOfHPhf3xJy%novpt3OME6|G
zZ1H&|1?95&{%ojmNs&>x*wx!x-v>xR3@C*s7tf*qt4AM4GF)k)h9j}1(RT+(orc5q
zi}<|#>Z7U_$1yoDVxts!jtWP8Occc35;8uu&QH8}R4Px?*>i*=pMtN1j)(beJRfF4
z<w$5@B4y~Zw-030wcI}!_$|BM*-h13t#@n&RU0eK!cY?fcaT(w<^~uNo30Lz&xtP$
z{MEdY0)f<%7u&F(gv<;8C;+8(m~-nCYReS`1RKhU7i=8JZ6|$qqV7S(EG@O6o8LHU
z&p!g&i!e$}mwE}ynB?XA`hzX;=4Wmbvprr;rCDGX>B(nb_l{J6-<)6R?wW<Y7mTpi
zFq0wYs@y2-Z4eCZS)@7`tdcaHEZxh@XQMrAhbC4{LKYV(9xIPnhkh26bx=gPpt`m6
zwmo%TuEe7kGV?26zMg#H)c`(pv`o2q6!;U(mz=We!d5D}<q}!AZlS-9%(eQcV_LE(
z`nvYhd1a0{d<{HU^RxwDF^Lj>;wVx(XRe|ZK(In)#%7*&nz#ocbuZnTLW}Ra4&m}O
ziy0lBHcCAl{e2=-CbjASOH1$h;nMA(TGY(`5x%^)@xh|;w9?@ZmF?ZBr?-qRgO#8|
zORS!FIbL>Sr!0Y|++1iCvT46xg?*!Pnn<Nt+oug??kj@*S)MlRDm_}WJRg(xjuemU
zJz^&OxC{EBih2|Uo^fV;L2;G?f#8H)^JTxSr7`H%wdN@dM|s9fZ2SGxi9S*du5Dyh
zByjr-m!x#8jUhd|$B)I<Hu{e0KCc@^>%yGR&TdF%98TSN`@Y5Oy=P-f^M7)kV&GXS
zV^Q!OM-Z6@F_#(4iQ9~(S#jT^a#$K(npIdic&=1(@u98-PScilm?TIMd)i!2p-cXC
zqfV^KGULKxZPxiTj~xAo8ES1)gQ`DVo;-*Xx7FKL5mQM4FX+zHqWqTkGnuTA_9|YM
zScC2tqX(6uIf#o(mruMDPmh;Phuf*O)>moPg|(U22TOLR%frh(tH$@+Yx>A%`Uz8y
zV7^@mYnwNgzY5O^CdAP!yBrvwJUy+L=0bFfq$eE&#|Kh(h{!p_dvxOW!h-feE6B(&
zz%;BcZd90OR1`&h4<XR3FuZgJof2`q2h}UunwY<S{GP(Gm?ovG_1J|{-Hr4@=tH1|
zmb_>5`bud0UI~qEXJ|EHv1^CJ3g|Kwof(KcmbwEZqPt@9sVmJ6J^(qu9=>(H0rNg=
zxcaHL0DD{;sYVbsao7{rYJ?e|sz?=utJ^y`Y!9ROo?HR7EZZwBwWG*inz|^G9X_tJ
zrD`4oS!~U(dSg|Ph3m?PL|EZP_t|-{*8eK)%%h?1-#$*dX(4srZjw+cDeH~Q9ibA^
z*vgh|MrDa%Cd-gzj8atEg_s6ehOA|*gPCZNb?lR2OqMa2m<(nxhUZK5JLh@MdCqyx
zdH!{nGoSDExvtN3z2C3z$62KsGkto(@r$I(PdLuh2MR5;#coHjG|usEW#n4Wq4mP}
zG<%-afTniov6aWy2D6q6J0f~h#?~cOFjdis*A60a=P|1hQPJv9Nn|+nZIidm3~U5h
zd4$}4Vpd+)iqkGpIs7Ay$Fs1ack!UL4==z{ow(0Wai-kU?|OCh4#f19zHY0cQ0~s^
zSKem7Im5*HCtTvu;i#r#ebT~<-$q~wN=f;T+US?5oTNkFhFaV(C)c#^&?$WawQ_YS
zrgU_vMj%>i%dAftIpGd2gEqzVPKD`=>%5;AP|~3bL!4Y>xCT&8Y^Z~qS#-s2!r<yQ
z6!F$rqKGFjR#kEg=8Vf<dJ#r72fOA5T!<qfvjp*^_E9l$m_q^X5DuE2DAx8@>H~x8
z!*}fqn<EqBN#gVrKb}^XGj~Ko3;KakA2gYv+`sbU)#}w^L-&GSlQSOSG?SrAV@kVl
zhkap@%(++yhYd1-(jreSd#ly8b9c7awC4z?b3-T>p{dQ|g^8;{#qP5sw2Hw?A57DN
zApwNzA(nQM(aj&mVe?v8qJq)OtEGo#;!jr&rEqlX;e`1Vrylr-Z!6Z-j+4;TA=zu;
z9<q0<+!N0I7ANBy;KS155!YJ^^{dP;qVJ(=Wc&Cj8JwFO#c-`<hxA!$#6z*G-3vgQ
z6cP|Ur<t;c6-1u8OajV&dZ;<hN`73-uC4T}<y%_LlN?^(H+}wby!ZR5$uJ%AL$2bM
zq3>y$1$Yk1FlSwlm1T9Q$pRa81wByo2H%hmbtj?-8upZT6Y2XR;DpL;vxx=7YuBlx
z-t_H_Mk0T987dol6dD17EDZlBPyC-xDEB^c3j~4funI2iA68U4$kUl&-Yv+W37EK*
zd8PagC3oU_Ngi2f4-i5y)oAlX*xeItE}lyU-HZx~5qFO{;#dLgXqWkPw~CSN8U~4@
zI<I<#j}%AV2=`P}bF4J3(VeAOG{>58Vq|r|Qiq`nPK^6<?|GW8;|g1APc2NQY-2q~
zDBrY0tR@G}Kiypq_ulpK%)^B`CAhxxhJ0>$bg}Ez1QD1MZ}CG^5Z=45Fy<ehdTC?<
z3UYs<Ks?@MLI=;Y%iAtO_$i#L*UZg&j;-ow6GpaZ51e>@f?0WUABnUv=~}oXvN><m
z>95z%6c|v#9%n_lChuQTCi8;By$5G9!L`2A4TZ|<13G9Hph<Z`$y{XN+L()9?NRM#
z{1c3iGQq2E=^Yp=gVJsxXF-!glSF_yNaI4A$1#;s98Jo%^_A-yv64tBmKZ`_yh#df
zLcan*RrfA*zHmap6D7j;r0B3Ncn?O;$q$Dq^jeRRT#fZprWekxtlX0upQEWZPJ`+?
z8hZ%y7fRna$PWGsqu<cyr)TO1-&&pg+8yfs<9ZZWIZM+-Z{45X37tQPbmZ#K4x?F9
z%f;qA!a`c1+Lx~!MELSRpVNd+=3o5f{LL-lKt$DQD7~h<o5pJ`421;j89M5KYO*ai
zz0cmdLKb3;;IJnzwXWSei-J4TO4>t*wLePnkcXD$&!X#{e{YN}Y#%Clr<Ji>4B3jt
zG`)^B)uy8<JR8qCzn;f!hbd>Gxpvo$S?0fYi0VJM-RCf+U`1`Clc(eS?!T2DlSn<$
z5TAcBxe^h1=pqy6m9NkQ{o11fS@<-3>aw1qwOd|mSfiqh_KC#E%#enC=$tOY@}(_1
zWr5bBh|N^Y87eI)V+Z>VDio;FXGAb7TX<8p!{Y}v+~)`TkSkh@HCirkr-c;LJG**K
z_HQl)xVb&|WkA<LvnVSozW5~lgIlsHqW+mg&ffG;RVhbkwyVaMUQWWLLvPR-to$c>
zXoocP4Kj=Ix!IJ_SB_B*NDZE^Z`%6r{8XC0gwzq{kPGu}til6M0l8g`J(nwxsVp5#
zmh0tWqi$#>#^<fa0Y=qE__SxDtHAhbX!Smr_krtO?L}?Gqvah#-(2*p3ZDvhWBO7F
z#yR>tdbvjItn2+4vL(>8ofb&L$W>5^soGl`?gYNm4r1#1f|V$Y%Cn6OZ|yayS=)`u
zc!ZY7$T;~t*RV>1;*GKPx`7-^szvW>$dEC8z$f%g4Gl$_u=~>ZCO-t0PND+d^IE}>
zy*nn2EW{IpeZsRJ6~|A?9;=;>!DTES3ZS!nfLo!eRY-ZFJrjj{Gxi@a+3zu60kWvu
z6c(O_G){<gP=~{N=Z0iPo6(G<W&?>dDdV~jX|C;1sSr3aI-O@KxgC*J4m3fjH%NPR
z@TUM@MdY)K4Yk@Nd{AN3W$?2ruP@GW2GWuyyIZTN>;iT$lJ|bt`HTKEq$gD1g&|k;
za-+dWqjlA;dpO#|&WJDn=v-}!eVux6+#nsq@i$MOC3qbx!=D;58FLx4Pqzdqwk3M0
ztcnOneA<Xn$s<Gr{1BPg(n~DHuRyvqrQB=16+YVaa0%C)r9Vt+Pwu${q16QBE_xef
zD%xnAQyglCo(o}_tPvAmRUvlxeXZ|zC8g=x^_8mHnYo2EmsZgtr>7pOxnmtBo_G<v
z+_<@GW0xkXh^C9)YgoCD>f!Bjrl_(gD7zCa5oYYVlAr2J9l<M^*!hIHlip=U@k87{
zp5=SqC!aON-mNsPF_P^MYix;}XK%fdA>7VJ7lUIa{z!K%Xst;7%P|g@#ojtiPPjyI
z%?s`sF8rGp|BauD6Hklo&*PgYl3@8pB0EUyW<wM6?bm=GmX-^Mn~iAJdip-xIpS;J
z2)%QpQNNn6T`iXMnRMwsrAP-hMm-NXHQ0b`oJd)I8=n$-Me<YFK96wQ#G`ZFU<LDf
zm)O$H@VL1%*RgfHR4i&-z4YSrklS+9l(xPEE~M)CHP}&YjJv`CC@C$}uFs_BC1!hB
z@}O>cEn8cWNmIkuKh^<X@I1`Vy~Fb*b^MG(WH20V*w$E~mZ)7pXFrcT8>WY_3P0i7
zR2gt``1`~(T6<b@F7xsi!ss$_tfMxd7;O?<8c+<i!cI-|3I)ip1@=lqs=ZH8>ls2r
zhUe}#AmoHzPHX?gW*E^SGDD3{N#G;eXynyaAL+nwQ7kw;uVV)dB8X!SX0RH!u@fyY
zp~=zG+?)>TV3u2k_@n`~@<(p-<Qg`!)hA|PzM&W0G@u&CC7f@ekJN@JRf~tu(6`(3
zh~22Xo|>uIzE7wA^Y|yaxl3!EWcI!Sa!eze1mZ)wOUpk+AdLap=|ip)-EO{RLQZR=
zvi<bl2t8K7t88)EbZu$LAA^sjY)n9BQkD|d0F}wQ8>Rp$YgBl^)_p&!5JZvGU{us<
zwVQlo8P2nc=h`9>Stj90!pWX*VixN&w3D56ItR@j1mjWKRiA_DMJ?B=BrSvI*Y0J}
z>=$>t0<Lp#FD3oRUcv=#WAiCpYILN!=ENYs`9pi_?cnWyJ3_2#eY?36Lqia2(^!~A
zZNzLC^^D|390IV*Vt;?{Rz{CA%u*r)UFjD6nh3<sVD_UF+7Fs4wY(oFUm=SIW8C$B
z!1;$^<42#@-%3OFCg_8wuLyvzkV%?AQTaJI<MQLPq<5+GIS6<$zb>F@2Xd)sveC)M
z6poMnFh4^^Ub()KGhsB4pg&&$B)~vD8|Z}I#9zf4i#!*mjY`n|qO!7p{l5nix1Sug
z#mZQf9;T1Q%tG8hMrv*+W+W}wZJY85Tdx_gRJ<ABXC#spd_5dd;GEF|PmW)U?&p`A
zMuh#01c+`oxTr+GJ=ZL(Z4bWV(BCnt+01^dk)uyr`(sXTfgsM^$*wu|WDR4T^h8m#
z#51&MHA{uRJ6qPHc}Jv5i9Z#p<pt7X7lDa)C6CXr@ryTp%4POO9%3w8h*!)fxb(7l
zGSI_}Nq+SIqFcf0x6}xSA3~7FcoaGB>@J1=5BO`&UyYNdqhKFzdJf}1XJyfJ7k^}O
zEYfEO<3CUxbx<#5!~b2kn-~@27?(YF*VZ6F9J46X@jjfx_>-&;p4g__(C?{7NNp8v
z)y?%;bwJ8h1U3hQA6mk5SvRZH>z@V9?9<wo>{kofZx#NxZ(?+18Nc3Zx4-8~E$c6=
zQBeZlT07rIlrl;=zG;-ko{M8cp6ziYSzp91OXbMlFKJ&&k?&%Xd$TY`u$J${drKu&
zNUVc&j<&!dU71B<MxV(dOc?H96@@_A{b*8)|0;OOn6uJq6aGrAXD<`gO=t)#_~3Kb
zCPg0>Xueq33>;N=aW;OliMgyI*V?4%L>m7wt_~!y?p#zz^0yo=`{=lbf!wn)kagT7
zYl&Dg<&(>|P6_TADs&C0Pq{-U>w0xA;BPH$oKFun&L<wjd!m+;%>%7HoN=N`qiQzl
z+!Z@+<n&KN7SmkR=A^Tcamw$X>rK2S{ev)pwtT}~kOee(&ARp}hvJ?ee+x_Vp4ztQ
z2-#Oe5<@Htn<=>xviYU{ZeJ5R4nb^g;>Q)#UNam-FLz_ipLH|24c07P&ZE2K^)^A#
z5Q?@Fcm;B<3!WD7_Lj0u+QS7zs0w>!jpUo|Xr&w)a3BYJKa(%I^LfLb52Sa61$70X
zJvk1qT;qul#C{*n!;tTij*)yHDT#=D>e*-g#pGEVxyAebCv{PI-_o6cNwQA%%a_W#
zxM<W4R$I2c0McTwH%U*>Pay7RpwDO4OdyfCNlAI?zt3h!6O_6+$wN_fGhNZ8u<U&N
znS7`NQ9#hpMO_Nau*YI>G<zQY9FLfNymVszY>%`xCt>Mr>LF0IhBhZd{fyV-Q^X+3
zy^9>t9j2Ye#8o>L0vQ>NQ;Ef8NPJ}CW~t&O_+5<pE+2nVcL5+?^8r?CuTdKN-gc0}
z^itD<h#vxwG)lAivCXchueZN_=~V-n6^oK(XA@ss`E2bD(s3wAY^NlOS>x?5>!313
z)Sh|bKlY1t{MJF|H1S#S@jsPNl9sr%%W!3Ub%(I@HUBX;zp3ZYi_4Z<LnjhZ=j*#_
ziUAE)g`GVqtec^iGqV$LqQ<FQ7l^r+H7NZ;K2^O=0WR(FkC$O$p2gw1M}>IHll@uo
z0>60Z)yHtRkXH>^JKye23srqqIam|Evf?u4hND!v^vKE^cdD(&<9b5tJ$*Ogf!{N?
z0Y!cNLX#;+`Al4v&SLeHn+Tec)cjFN4@yg6)0DhxScVFV;udUEkgj<)cltXuonjVf
zYp^t&SAu`P*4h7ghPuMWvpHAS3zk(`lc!*{O@DxG{luQvZp_h$B_GEepoP1h@epX&
zZe)lXfxKKNXMN00)9>d~mC6-*14i_c#nVq_)oZM}l(cthmoo+##F9xL#>(Rip1~On
zc4cG*!N%-1+vk?o7o5a0M0<AzrUj;x)1e$TK2`Z9_B(D)-bhuKhmX#2;K_SofgCxQ
zyZre2hFq-nZc{Es{8$<n+8{5qio88X6Sk|kH_xYX6Pnc+QqT~r!n9bwR>rUQW`)Ex
z0drdH`IPsjN-X+QCDz$KeUP-O@b8Bx_%GR3s}69d74pG-YM*wCDeZW<c*g|vO+n><
z{6cho)YJRL{m3VeK;|iBZw0jVbpF+s<lnE7#Yz3|RN?<g9!`2jO)V0v99sPRj`wfp
zPF`_Fp!=^sL<5AOZR%ISA$GK^llEd;hA1{F@uVGE>}dKm!wq7tjZPIz7#qEZ=VhGm
z<Tfe3VWW8adur&EPgbmvSWEfCO!h(rnm&4B>(5UNIUV!Iao$wQN)p2kum_&ic&*RX
zo{2f~F<=4TGPv{2Iw5(0TFtBej7o0^nvXc}3#!x-H=|fH`bHGr)c%vc__^j3wluuQ
zZ^LcZ0^IhNRYZ6V8o2T^)>c$^gJhIZ6HeyKcubLOIHw42NCvAjC%K?0q+zwk&R@?e
zCXzA7b=W85i&PJ8WyBr?EPq?T^8X*&<LSSI0aoDdhQ(cgxm)#N&oRwn<|43!&Bc5h
zPF+@4^{!)3a!)~v1~<=$$}t-I@rB1>JXQMU!D$#3si4_2KRbE0pa%ry<~>YLxw5P_
z9pep+M|+0MSo#@T#_N{ka%;NMqc;)5FDzY}qO)<`QyG|In(xs2+WWg7qbv7({7ob(
z>3nkmF>H9xZPkQx)Y-9s)sGy$G%QU$)4Rlu<oE^x|MSc~u6^Fp*$Rto?H%>gz@}r@
zLIV(cpG7QyM%NJrnqIQU6~`?~uufd-@i|W%`R}clk6)B8#!w7V<tV~%>Z$uwZH%(~
zx5x7pGgYH47eQVchV?eJz4im4m}l$XO{dUOwD&MTw=0ZSp))IZH6gbc4p4;Ix0*Ml
zRYkWWR2+uS-HbU><5Ys0+Akzdircb+*x&^&cRiGsxe<ZyDW3hg8RFt;&?T1%3TG_h
zm%@fGjqhU0>1%-g71ev+J)}#xht~1MTRNS}-!Gq0O#ALM5NjKE*yAlAOI@e{O#N3&
zqER~C&^NWhy}6%FPrZD(zjOvWcq@$A-a#3SN)WYmU=*7oqiD%E)Ay$Xb|PcE3a??+
z0|LCd{mwDDxhGYIaE!VaUVGP-gU#2<)f=R2Vsd5OPNuBe;~SDvL$vA*>9EwocF%5Q
z;bcH4L);0Mxm2v|dji;Wa^;wWSEc@Fj)Efn*mUCZ`&Sz8|BS9b)?tOLXGrzM66%~^
z3n%;2+*|Ar=q@*t*8<8PNm-He^Qq(YR+Te?reDvX=dmT`-RIROH`)Qb;DaCNb9dmY
zW^e{#se0(!-d}vdsPH7r`KGH8;lwttFQVepkLcP)zQEn4(E@qNXa!GR)batlUkXiA
zA~XPiG2`U+2CTxMVA0`2su^M94GZ`G8wy!_gnceW*xo9fCWP1&f@+030r%?IRRD#P
z2`GGer181#Vr+hqgHv3~?T1H@-naiqoME$C1b-5vLp1uP?k+4o%ZJNY5>r=tl3r^#
z$n;*OudDxMQJ=cdN@6gWqe2PWOFqiIQw)cnPsp327bw2co7YA8EcZ@yaT3rn9{1f_
zTZ=t2Q2^EuGf>!$=~$Z@BE!4kTr+p6twGFCsXmhk3DX%pfmOQsA}vj?5NPEonCT&V
zUcBSo#ZEHc<n!cbQ>Lui1K9aKv^x<t*j00+^po2l=mF~|I_<G^_?pr>0bRODaw@!N
zS-ofQP>)_V78C{pWOCR)!%v;@;laHs-~bWdsvH@>QqO1cEzB*WlAQ7TZQq~vrk+eJ
zZx)(<;+aX4Uq_DE(3RBQ|7j2T4~RJ^Y`?)_0#3W>#h!WiiOK*v>4%=|TM3`oU?Q}U
z$#Qg#%b9sFZLA7_)_^L7aq3z-9r~7Loao^{sDQcfq$YSN@SvSeJI1{z)I0Xt^u16>
zUUX<|Xrw3KaqdrBcn>dsW}>Wh8u|BbkL`rv7feU)TIU;um2kA}Hd0)<07THl5SrH`
zlnTH-VXbcLpF4_|A_p4~z*C@4ns7xG0DB7!DLj!R@pdcIHPmeuE)f{xWsYk<4edQ2
z29(G%vO_|Ckxw%{4&bRz-_XF$IwJRQvED~~bwt=)l!XsjP}?R{fU8TGm6_h4JXbHR
zcSWXteFqkluQ$5o0gK1_>@#f-WM)Z9R|hXzySGw`%fy>W#-eQxE>M?$w3x-AbeyiT
z-NOwqg-50g7N0<-ddV$Sn3=!alVj3C`5_-maw~bR1@korJTSW$4(<rX3ZSzzOkI`~
z)ERH>aWkQ=VYU{O!J?GG@z^}ypNm}E1ur^stvHZ6MCpI@cVF%I?7@<A#N5w}S!n=}
z&^iyk^M^#HkwK+Mw(0Y7!T!?@5=;H#7N_#3F76fICIiSZ{Hc@DQ(Ct4R0Air^7Jyn
zA-O?daD978ITnqJbY2ocHG72mh{$+ctioVxY&Ca4$}j)IL(9Mcrgh-LT>TeL$oeFf
zB#@&i5tB^8!I!N^C-(*Uw!`E-Mbw(Tc~j!)B|4a2X`bY|odFPRimxn?%AncmPzA-E
zD-B@@Wy^j|jM;u$cn~7nEv&W~Roap%d%6NVvij`j(3)2ga5~5qKuHgjZ=$5KpkFAd
zH^@U1c>G6dz6eMzHp$o0T>F!~`eqWliEVa-<|r4*7VWmxN`AhS#CxBL@?iC8awU36
zFQy7bBdIymS_`tB^cV3#GkP4Z+q8H8Su1+_6B_aK+9TSd%^<V{tM2iba*-6-;{g;e
z>Kd$CR#hWL2Ow-Bj4^|okVp*aVp&-e0KPF(tV?s7(o}EVnKmUJ{FPvZMFUoKbingQ
zx)|J}${A|LUKK))CQQ9Rd;0S2R?-HxUiYJD)Fk2Lp9+D?EV~v)3+$B<d+%}b{nO6a
z{dF(^zJ~ST!S8jylxXac_{7pCBCm)A){`>X8nW9!4JAtz{^@2og%Gc$<9K*|BLy1W
zGIkgh@*$%7W_MsBUre&nU+!EfBUUZ6llG#)|IguKb^XOBSF;so+B^Nv%^D0B1RM8M
zxw@yP3G<WNR}BC2a?&B@p?O+=+IL$BV-cqhLdG1)RDd6Qe^qakUEW+jv9Eg(+7u9K
zT=m_kb6`keGrX(0X36C5{?D>^rn-cHgHqq^l^Igm$C}a`&TFY!s2-y2+XOfMz}%$A
z5nr772iW&tl`LFlgbGwm-nB`on}(~($#Bnd^4F_hEVs+eV(kP@kZ=zC82e^ONvcG$
zOX{Rw^Lv{HdI_k*bJ)NCMSZwRt4DN$qS#zi87QqX^sy)Rj{;PisI$@6TReY=^8hDw
zzlY}N@K`zT$vs?#?_j$Gj-ppvH1M)F)#(L7efxFOv#o{o!@9r-Gko>nW<hO9(lnZ-
zzk(no-mD-R;jg&fJU57$@7LdJtWPVXm)E{H#y+<NWHx@YiLblt_n9I3rE}}$wlX)s
zy5`D!K^=<?5<+gwB6EM5^gp=2(jHY7`^178y(_bzg>M0r9p4wxf=Q~}ewXI9&&RU0
z(Q5OLqK0dP>2vbo1ON}Oi2#G#j>9&li1~zAxN(W+6;`4B15sjy*G;ZF<U8~Px?Pdx
z=~vrS3@)swz$fnd25oI^_kZf;pnxzF9M$iaUh+GJC_!)47Czsz>#ok_a*;*A+$Zd}
z2W7V&tuYE#FGVfKB|Amd%f4pf2*X=gBaU0tE*c^>CZ0cmYSSR7ni`qNHB|Tt`jxsV
za<0u!%pDS(e5dt?Wp)DYj2k>g9Dr@BZM;VxQD%r@Gcxy_jZ}zFcZizZ@_?JnV?<fA
zhSR^8jmcxcEO2Pn=?AG~{L@vqd;Oexco~cb&1d>Rg)W5JB;NMQUsLM+&0V=}y?n!b
z`qE6f&+WjFgMB_k-tW8FZ&U%(!mPc~mygK<IwJ%x@}1xFz|UO>d|gqY|IaKCeqddZ
Y`mM}rm&@*l8*8^PwLVvJ*5men0hwKe8UO$Q

literal 0
HcmV?d00001

diff --git a/docker/docker_build.sh b/docker/docker_build.sh
new file mode 100644
index 000000000..6e82a157f
--- /dev/null
+++ b/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
diff --git a/docker/docker_run.sh b/docker/docker_run.sh
new file mode 100644
index 000000000..5fa4b407d
--- /dev/null
+++ b/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
diff --git a/docker/dockerfile.base b/docker/dockerfile.base
new file mode 100644
index 000000000..b83488d48
--- /dev/null
+++ b/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
diff --git a/docker/packages.txt b/docker/packages.txt
new file mode 100644
index 000000000..1ec90dd5b
--- /dev/null
+++ b/docker/packages.txt
@@ -0,0 +1,3 @@
+wget
+curl
+git
\ No newline at end of file
diff --git a/docker/ports.txt b/docker/ports.txt
new file mode 100644
index 000000000..2fd9c4567
--- /dev/null
+++ b/docker/ports.txt
@@ -0,0 +1 @@
+-p 6006:6006
\ No newline at end of file
diff --git a/docker/postinstallscript.sh b/docker/postinstallscript.sh
new file mode 100644
index 000000000..457f47c17
--- /dev/null
+++ b/docker/postinstallscript.sh
@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+# this script will run when build docker image.
+
diff --git a/docker/requirements.txt b/docker/requirements.txt
new file mode 100644
index 000000000..d2c490c3b
--- /dev/null
+++ b/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
diff --git a/docker/run_docker.png b/docker/run_docker.png
new file mode 100755
index 0000000000000000000000000000000000000000..dbce9b9c42ffbc77d908daf4329b11c4864ba5cd
GIT binary patch
literal 88368
zcmeEucTiJrw=OCO7C@9HRf!al-fIw%fFK>|O{#S1f*~lV2!!6dN|P=vbOO@55a}gI
zPiPWq2q8D<@B7X@cjnBU^VgZVbN5V!oz31mdDpw%wVw4nYbQ)wQ-$&>(^V1@63Qp4
zk9A2%$N(fH7l$Y=6OZWjUv?w@x!|R%@`$8thy_J_L+YTYp-4hf5q<6SIT`W&m6xi<
zUL+*+oae6#Pjv77CLtl@J$bCC|H@($kBr9k1x^xz2+$+x?!$`G$@x=}Mf*8>nw<16
zIv2lc=D@UZct`@or0pY6;}g@BY3o>-N7u=s-jf<F=+UXZqoqkhe;<4AZ)1^r{YKlH
zE2I{8dNj&BlN(QO?XAR!F2^v%iJFBK&mHY8zV(`4<P&G13T;oe|3Jw0+nK*Y{CekG
zomg<}-#wCm-Nm+^mAg0F0$k<}i>~Ke-%_@ggx9*0laR=phPNZ)rQ8aVT+dIo_#SZ?
zz<#K*n;>`ZaeAN9cT5ZfV#Jhf-iRZwPyljbqdI(Lz7?SQC4<btC*_2P;KTADb@0AB
z|4e!B!ro%L!fC_tacIz;%TKb%gYd4qs2zOx@oxq`$H3-r&*OqcJ(Nds;Hd&+H11GW
zX7!i&Se5g>7m|VSdv~XE7<g8CJj=j8W7(rZI6`vyH?)AwDTih@MCL#8ptRm?nBsg5
zCw>T`<f)DC9HHvyZfzWjrnz+LKqj6$BBRclx`0w7B(Hp!LSnOVvuhoAeDm=d;WRrU
z2rEdiQ#*s7K6R_YM;7CJQha?mW?sCXn5JpalG*kcO^u?q_IIC}kB@FKBuxO~<{q#W
zN*~7^uM75F`1ajBEcz^R#a*r`%JWroTaQ=LPAHN;$ZvIK5fQZOvm53<zPxm(zSzEV
z%CL)KhnRv5@PiIP$EJ&VLB~}q-j+Dd-7p1Q{xP&3sens6{*46dYfcoGhiefWZcClK
z_BhHW6xiYV?|vQY@SJa93>VB>d3GoYC)6%RL}-IgJ0a6q8?x>MHSme#fQazkL3YQQ
z!O0$6_-r-v_-8p?&|V$^r%*UdFgmT9A9g?WxZgZB_?WNu?Z;)j(dD}9-CG@JlCBgW
z6IeD)ThE4~&`lDO{ZA5%SUrUUy~Sh*Dj9RfHgT2<F?j|XTi1D4iMZ|B>1{Q27B2fJ
zq|>6~$^26CrZzl|#+e;*<X5%ugPIU=+%ch$9e6gDO_!~Gp>G)C?v0VOu$bv-4%&OL
zxFlvY?k6<55FaX<R4&hGd_3SsaKyT|7%U@J^^gOjfxi`6QptI>@NT+0(7;1fe~<vg
zgv`(lvdun;Xu1&`OX;#!%0;|Rc6UjEP_B98V?W_L4?M=gX)Og3LsA6+bt%VTL5I%{
zu(h3QLFLCydP2Gc9v&PcJt-jY1{5qv10^N8_P5IbcQ>~VL-@&<*!pol{Iq}Zfh9Kl
zc!V9Vv-@5Ud>0>tO(l-Iry3S0=0Dl}L}iV~cP|Y!P~HTnmaArA?`U*~m`Cd4G<HXZ
zpTY36s*yi&{JSI4`Oa|qa(+gCkg^bFqjOC}2=BmaMVyOT;1G~dZ+RoTuizj^4~#Lj
z<UL!pJcS;w%fdN<pMzsP_qVvu7Hd1X2@4J#<hAIk@}N=*+cV?t)(+8?9C*vF`Nm=L
zW1-5dgQnYUJshtNpY4tqKH81<w<?u8*|}MUZVEi@UhMu^5rpz{!RoekA_+L}K3pc1
z5Cl8i>6#|7M=e7NaA<VFCQBJbpynb|+*aV!){aj$-dp9`;cnAFo7hMJ7f2pbm$dZ+
zVYJSkFu5e?Rl4S|uQ&vtVNXrx=3jgTDPMko^S!@ALD})Lf=;o4K8H49{oBo@z88IM
zFS*8X+ANVHH+F?-af$LjgiLpoQ+D+7DLb~?&TeM*7R5=(d>qT=B$1u5%S7kfcpu~_
z+y|Z_khJWE`ClLyYTgaY+lu!PPadMuEe*4g5z#H_kxaPFbC+62LM*^v3FfRQuDuw|
zQ2s;Q?SAh`>a%WlhZ(y!YnG>vWW){Du%bb`8;k8hI=N$rDuuIlcK$5@^2}lCPz(KO
zl$k&DX7|Gda^y+c@iDowG<cPqpbTC#J=_O6>bd&A<hRrsx61V{zJ%z-YU{Nr93$LU
zk+?#IV;AXWfI0kf5PoBPE9cqEdYU1HL-So_I-~obcX;mIdSox~!9|m2%4w%Nz?0}j
z1pGvDk^Qvg>d$*;BRlxg<J};<(b<fg71yOl#F)4r5p*A-h98{U;YI~v;j4~YilsrE
zz_UoWUbZ3Vq*rd>OzrgN4ku2n)<G}99fF-vPJ7&fbO#%P_d5yq-CB1279I3YvLVM_
zySr=sfX51_7{87+!X7ewrmrG!*U&k?Cg^NtQ8cmg7$bS63qG+;^Y$DhEFN_-7%D>u
zBWU*y2FR|0!RRxiVbkwSP<?1@f`a^mzDaD8BUSfm;JjIXTh8{2>=W^WO!%#Et$Qe~
zy}aNE_#Qr&t;aPe6_UlOR5^^)?v5j8n+u8U4{0*&ww9DU3yiXw7T;HQ+Iva+Yx{Es
zETjYu2zttoi7iS(TJg}QW5iyiBDbuKrY*PaISdjjr7oLRFfu7jj*|j*;%$;#clB7p
zyZZuW-SDDx+KISC?EK}aBUO>%OiP*d3gq;{{9_lpvq6jof~`|RI;t0KmEfT}d`NBm
zB3)zo7Zdt63CSDpo6ytKHHGZx@Y8O>F?%vBby-4QM)xd|KXUN5u<3(J4`Kf?8W~V5
znk2Ck_deFbXO$0?`_g>-g}|szDy+Hyz)g=PlPy@xjUTUOY#Gw49A5KgAy;?+d`{X@
zw6Fg&!QsZXsKV)-+(SE@AJBSp?X2+FR*4OWV;~^GNdFrGQ-&45>|f*r0rm7X@oRF#
zXa*6@+)giB;SPyZIHubj$+ib9j#{ml@`o0|9$G!pbQN3~*3RtPjF*^5BA>t<=rcpb
zCaRw9flYtDXFpR&tlq|0+3HSc)~36v-IK^cAFg3Lo{~8k1X-M%+BCOVxPDuS?u^d6
zkP5$*P4L9*C<>nptU9#N&<8n_4nQ!HMF<B`g|j{5H-q#)Pab1Zsf-dYxkWEtqx3)C
z&7MF(@n>g9V>rPfNbs7f9dseS)7<AVDmi>fc2^c8-Lo*exP!4NJn>ZeRJ`3Jei<<6
z*e$~jT+{HBlYsB7>QJ%ya33V@6pK9pw(L`PG-lV566btNw3kqDc7VJ$W1D$)v{)q~
zbU$Q_uvuOSN8Nx9sDM5&3peUaV5A|8O&=gn)Y-mwZ>y)<ipppnrHX{ug%nPLlEnIY
zWfn9LtAnjMoQ%e#9G5ECk}t(>6J0~>I?UQ;(2s{rwq}<NHyHAI6{%j7l|c%K_X1b5
zTdvt%zB!>iOoS`d@m9-PXUkgoh;Q9X6~~xlxY@LQ+>9+#s=>KKXh=BS^>1aLsV&e~
zQ4Mb|HG)<}t!%lZx6tYNV#Wjd>_$^(-o;LquUUf)PdR|*4uJHm?<&Ij*(;gduLRkk
z)qaS)$@h)A4o}PR$Fy6Wo+jQd1^t0Zp-Y+r(EVq7nVb7`IifPy)$+BrTrE5t)LJV;
zS5tKVQ*Q9fDVSiGAF}8a2R~EbeUn7?9(cN4J9dbq%RaeOQ_hZWI{q!&`9dMHR*oY1
zWDCsNV|{OCCtT|^ba{t>M5+<+#<$H*#@(|n55x^0w2Yp@$2?<jLN#O;<K;FMR=WvG
zV^2~;6~>T98oOcTKB~@6{UN>Ft`XB#%ZA)oKk$iG{1W5oBbtoUswnEmSP_z-JHWln
z?27zUxetQ?0DRwk_ZFIrMWO2C!EU|mg!lzQW1UV;9|UbVtUhDAz(WZ;UE4_j3Z81w
z#l-4C;3lMk>;%oz4#DCP7xux<Zny#ICkQ5)5O`QmqaD|B`YG7I>jgX;AL-KZT>PtD
z8Q$+^39cT_*YZFi>-BCMLscx;QYgua=P-Q>w`Upa-^jp@w+xEkS&AT-tQ^g;M+kzK
zaEqdTT|CXvg!to<0DtHyjvXtEufIKuOSgr~uv^lx&jqT?o`PuV61wA>ub5t5``I&?
zI{-W=mT~9(Q9R0k9yaw7!kF*gOh(2lJ{{g$^>6{M4HmP=p*5Mn&nA>D=wVj5^J_cr
z76AtY%QSrbao9545sN-|*@U)Ar<Y$@csAFso^aZ;Wqrvsf>!PWnvpL6UG@PgA>?Do
zSVcI>{0zH=HzCD%0gNP`6f_s_7~1{bUM-W7xc!b99T?qV!=n&ym&6jWt=_xepGo0&
zK+NJd{H6TwezGJyO?2m_|K)oZB-y(b^4X}6#O2sX^h;H2H)PCz?+Pe?=v7qt=d24P
zBq;$@pr|pW?<k&*o5WDItyK_Ws|YkLC{y!k8ql0?Or6Zjp^IA77Ja7H4^f0ghxcKA
zuxO6~9%k39M$x<@W>bhcc%OY!FHNhC^w<8^N;G^Ajt_7Q-*jlvOYIY@69io?`lL-f
z(|otkoeAL(*_zGL`Sp^P$kzguue~%eCp&gO3I!OkNqu(-X3kf($D&q6FM4Muz11JC
z=H&P1*?lte#domh)O7>S54?*ka|1g_=vGdK#(PsfIXF5m0lqa@6h`R2QPs<79Tk&#
zX&trS4^D}?7(a7{ww)ml5KCd(I!xSgxR8S+Kn{20coH^`mUeRAGN+!gpG^k|ZgpSJ
zSb3`f_~h*X%@-|zZ0zup)j%q?cW42OSW)<1GSO2KA7*~8MmCk?c5}e}wTnS)lSC-J
zm$+HgG<!3TCV;zz7nl>0%H93tQuC_J`}g^K<MdFz;OMipLHs0mDBVH&)z3^V^hl0<
z1JNtk<kMJi3O^n9k;}MU3J@=%kJduk2j(VC&u#caus7uKNInJwe?Qqj5}Jo^=!e_`
zhR>BWl%^r;iF3iCR8dAskyp=@m^KU^xzM$|Q!4mCj0oUxq0J|a&9A<JqlvRm{_9ai
zDxG3o0<a(x#2m0`h@atRyd4_tAm9Brm_3!(LX7Jr&j!cBqVQI987|CF_6~`UKg%(6
ziJAzj2Av$nT%*Trnk8#vGzKy+rb~9)a(*#pP_`tLml_?zj`y?i4aXgl%)ts8iX5yr
z5Q;_r%)0_H<Q!mbv~xpZToFO0q`!D`W_T6rHdB_@TVL<??(c0Xi?c(xnFCYExT<vC
zO^I#Q8dK%D>gquTQN0>Pda3D?!Tyo{Utg-4CMZ9vcL}%3o$C(zda(em4Ncf*yDfdJ
zvpaIMf6b%iJ`ii!SSoq{5~m}{HnT_!ro4mOA>i>0GJH4?Ytht*gB;$B^1hdWttE)F
ziyN3IX!l#r$C3%}CFk86ZB3`r>?4V!xx-s(nIc1ZvxaQo=_&<?^+ie0h5n$M3UOsv
z!$Smh(0$<1-bxL51l}N|I_Xr+^3;Zi&v?pjJC5qLMtD||URCdNcPjI_x#=YB9%*NK
z<<vh*^2laS+C6J!!Nd=PT28d8;%9UmUs*<CJ<xkWAFOPhTGVl5tKNF1LL<EqN9s3g
z{J)XbWQZU17qzJ#?2PtkhV)sAQ!u!)1h@JxU{dfCA{-)1GYZ9Q?e(uVC)S$GmsuPv
zVCQQ$;o$UQdcJ1)XBc2#bGGaOfQs9Gt?&D|Bh&Rxmm0dzA4c&7KX=ElF82M|zH!42
z#|!Cj4j|S6DQ6t?4bc!}9`X_?a`@Ss^1EFSw@L%*fTU`l9a65ELr@5MtFc&`sN|Ub
zGXd!*^Kypob#H{0$)4=ZKU0HivdX<K6{G>)7vJlqTABopL*)$bh<pwK){5eT>GwGA
zzIR$EN@2oDcGB7))()yC;I!<;Z${=M&p@a8m??$Rnr@DVfDmEOoAvIXiqpm9Ju(2n
zVD{8R_<aLJzVkJ)1`1v5W;;8yl~>$zC)U`!lLdPZ`ukuKpP!He6fD+<r!V>jm5k52
z_}22UJ8bwn^O@TGEMVDM_n)7Q9$ahkPYSV_6At-2csBmE-|3zU#yu5znl?I<eH%Kk
z-Hx%lv(sT!r|Zlx4I+_wwF`N&$f3w_P`hivrwJP)156}(hZ84(5hX?`!+yJdXrdhI
z3yU3Z!HWjRata!pqk%qRQInpY4))H?_<&C=2P=brU9(i}B_wB~bYewHUD<ukjj~12
ze5?53SzrZyP8J&-xwEkNc$>#I4;O1yrvS6~^RfqtD+seoeOgK7%XeQmA?K$if|S9$
z><3AaJaSJEwD&MqzTeLIp;^GqAV~jY0`e62T2lzqKC{Rvo_&&(bz<fP9%beU9hE5L
ztV@oj!bcq5w*xw=n_#pT<6+8MX|eP&+WaAj#B-IjD`-DHMbAwDD+^vH1qlMr#^7`W
zRkEio)n7Fwv(^9<gy*h8%H&~Z2H@>%yPSJLi`jv)*R=6xN1YCYa=yv<%kMC)d9Dra
zG<Ed|<d`rcwq7vlQmMjOgsZd7d;~T8W_b31$~+xl&nYBU;c!lDfwwkKal((FbNZ30
z4}WGEHVo-K)@039Ten(PDKe`ZPRkcz#+9~^6eP9#JM+v?Lw3vc@90Io+n00S?q2np
zm+SxarL8ATWGv_1!A&yOkXVFwVEbbC@MeqKj%js6Fr?Qm__0z=*^*-rw`Hv`vOFyx
zNBc1`&po%mvA<TRAiFTOu%=C2&u*|s?v8QbUCoBUEX&4?D<A`vkMrCs@s~>n1I?CM
zCe~~ItR>@+*f<vau{wNT-SNi=$gL_e&1Gp%V(YgD-?n>|=GKC5hF{u1R1Bm}YTF)*
z>?f<QeS0!Cu_>k@!9j6kEzceyTmPd@GuxYDQ1hF2alX0TwOmtBj5rgWvkh>~b(-jY
z$H%g_c7MUeFACJwcB(eP&op*pOGP1Ncg!$5)&bSn#iVSpU|IMBAojoodDhmmit5P5
zXQcu6LH(brHJq+(<DVZcW_#nZ&q|ITvhPJh48XtHAy2@If-U3xrG@;|T-nw9;NaMu
zmNH9$dx7VqK-Y9KCJ%(fZ&^Qf?S8n-Q-8Li<h^yx1sq7M;k63d5*iqCRh47bBaks~
zGx=hYn>(>P9a@2aYXWuvbod=(tTGQl2fV7Vapxy8EWP~eh+k$%Y_CIz0=+^OIYa#L
z!2LRfqqW5d_z5w+3*rgkL|zSc+(hMj8*mbiKTA9w=yqbyChUMrj3qE$Ht!a80xiM3
z8m5$<^lhr5DfjzJP@ICD)6qeK5T`?Ar-}?aJHA+Ki=jm))nvzW-o-9`bLSTPI2rAJ
z3(H(5Kb+oOfz6OC_x$};*1(^{ctss@7%^6BXSj?kyfjhP))O{K`gT&Fz!7<Oi~U~U
z8`NMXqoDqRN@PMcAu!~@da@a>Y6ASNoKZyBK>MI7A8bm!sR$eR!`)V(xwT7G&@$!=
zPW}>W)6>4iZ}eW7Kw%3PW=kuZmYPPUnEj1xb44Lt_{+7Fgejq7S;W{~*?BZjCTj@=
znoz>Yr2msyYuZ#%UfH;eIyii|z*ZeYSaI-7-*PvJ{JEd^0sw+-#i1#tl&xi-!5g^U
z;;a2br@A)9P(x+UGQ67xNE|WRh-v{HJ%^{PHnTpiyn_@LQ{d?ods78gSj^>CxsaDT
zXRb6aE3wtWx<AAz3M|XwE<;<e8;b}O3`w7TJgy4XC1X_Bb0BDfKUoT3h^~_XvtX;M
zF3WTjX&DBTYOyt8HrgIArXDX)2R#TXjMd0O?yNtwZALsS3RqAp_ra;(JOw`ko0KxY
znh?2J`fBRQqk_SSxCC46CwapY>s2^MujM(f;8={;UL$0faq0zBzq(#LU#>PXLv$~Z
z%xVj@zrQ5McppKq3<5y9RQ2Mt4iYz=&FBYo(xcV}c}&wAk+=uDsq7CfzO`86{}te3
zu{K!1B0wzV_NIx&8K9VYwZ_?j*>P=|>4Y=+>di<t<yV-wu=P8^Xvj{oQ$u6qfV;F*
zN_1JZ@FlA;BE_P>Aj=vvn?k%xkS|%ADuBn7KJ2SWOXg5TUxepLZJ6^{2V%tj&}Frp
zXNz!;vX&H-+2PEoe>F?S5k4mPyQqqmEdA8b4&lO&KznN~@&w^C2RlHT@<rZhT>YgK
zUoa56ld%qwIWx`YQwc|z&g;;X?g=|{UIdYZ4pBHhr+?tiP*NNb7FnL=A9vLDT}X%^
z`;}{XY)r%|5Op$2&eYKI2(8lzjA<z`>7FQyV@jv;ynW_hNGBl15srFKjKRmD@T`Q+
z+6YM~;kRFBHxfs;``po0`KTJ`|9hw$fV8qjBTs{mAEUshZ;zjvaEHgn9oF#lyX44l
zM^Eu&S3#j-pxU|944Q;Wp-aBX3*AP3-x_=u*4R6v^(`=-xJWo{Hm)AiZ<ECeeG2Bl
znvq}J**9%K;)$rt@~q`}5xs-Ck2YpIQv;tdY*<+B9WAmy!{7jx*rU8^&xG}YpaXYI
z?0Y{r!EdnuvedY!mn}Gd=yDm5M-4w^-F@HJkf?zF9k`1UObP*r3}>HA1f{MTlVc%V
zb3r^nyd8kO?3`%zZ#-e%eQ!Cw5pGbICE2+$c{-H7Y(yo<1pRq1C(ek~oPD}X(SeSz
zsl=TItbX%-PBeG2nUJct=#1^#?tNS|$D}~VOYJHou6|O#uMWnZSeIJT0hD7x@vqol
z1F8A{7N!4{Llfh@eZE>ngj34mz08((8C?)d6FGi+f!=c6r2xRKe@c3WE1*^QcW%}+
zlvlp6<>$9)QC=`qru#@OMD$t|{`bbBlKctHs~ljV3U)=$2dSm}ZGh#oU$c~Z`K-`}
zgH*HJxjyr-+=7cSHRY?N2U7osCLvf5`E<Sd+sSLE-&JTbabhh@WOFpZ$Ioh_vdxBn
zpF^;N5K9F-^Md%vaZq@+HlP-=m|fF+y5jo!S(_?Rf+(=avd*kL7mOB^n~j8YodQf3
zi{r~(^(|ztieqXxT;=l77hlm4DR5Piod`@?glp7Hm>oOu{NKA@+Eg9=KXyDnxSX$(
z=MI5(_}W{FUZM9Kge&0eTd0;D2qt9?$2%P=V>Q$lngWOftr_<p?I!T`uGKx_Er{Ej
zqf2e7e|Jfm?=b!&@_b@AXLJ8-`@#P{)#AoE{e8auxb`0!&>!Ca{NO`D3ed&#9TJjt
z;?n%P_y2GDA2w|!10;uX@9zIx=7HTBjWc*K`qx|wu6G!XA59PV<ej@12}$crtJoja
zjJW^5&&gASz+KnXhX3YO|3fpmle3~ICodn?K3YlnXH+D0F-H#=3xi)iJLmYv&b1>h
zQ;+J+y^iOO1YGjNbNT3e6MSQpdw2V$(7k^dPRd^}B-Z|h+soI2`Nc|X+H!p4EAZXC
zkoXB&a0$mhXT4$gJx>m(K^>;J$+~eb&b<I|+}4)c?kI`VWRlxd&Txr7qUN5-2#zH(
zBrX?eI*m#6!=s7t;X`7zy2lS(Tu>EdxK4@){<I|y=pjl|9LN*Z)e_xi!1sZO>jXpJ
zG77*rq|ntWG~0aL|15$GMR%h5`I#bWinY^sF?(EHNyxJ$0^vGhpK<n@u4gB#(+!d)
z(T~*Gv@e)335h)?Cr_j_m1hcLEyw^NrISVkF4|~-4f5zo0D!{SdQ~N{T{v75!TG&g
zCw9F3S%%GDgP4^b0U4$1SR=3|T8O9)5ZRHFio$%^<2#5C(bl$*vKmqVw-f2AUiGbr
zUDreGU}j2L$UbhAWJ$Zgu&qa*@B7vxs9W4Hxl)CUZ5!gZPINv?ShtWmbEs@Z1{2O=
z`*EEtC1@4PIG39p?alTjl5wM&?O!YT(-pyQy`3CjF%{0C4oywVsH4T_+M7$htoe>d
z)Z)(S5|;sM*2z(3*E>&I8Z@GY1}-N4G9kO^BP2rAX`e)P-7CO5C&NI1Pzpb)qtgSx
z4H{U!XmD;l+O#{13w%1T|96C~w~69fxH7Fvr>E|fLrze-3zOHk(Y7`n3G)Vf{$70>
zAzd<tb`@c{sid!$8GQnh7e1Ys$`3Zhs=bym59EM}kf{I5dl*{NVm0;?Hmwrzrp=RF
z(_$rs>$FfaQyAya4Jco}t7M(FP@LoR-_L#VNT{;2-J)zTrF2P3y+DWD+Q(Td-RU~0
z$nl3b?gyfJ#k8zySS~nr&m1;yWBv*8_+K9*xX_%DzjX@=mYm6J>s<0pEW7JJJGh0a
z?y}0Ia^sp&7yNcRtW#njI4n<`S3WrQ66^X?%qyda`guR=rb!b^E3fT;h1DPLtZX0D
zXy#tO+5jDQz5;4#v<wfq%9Irz-!#XJB3>JYTRyIo^fXPu-u<_;KNK2m+v4rH7%&2Y
z^3AAE2gi=$&0K=N(n`sFcwHgm9~NaPq}%W}TjG)g*QY*q{&<Cx{8v<NdP7`_QzPCF
zCOi{z=pzcxtK9SL^aj*Li0EmZ0(~Gx5Hsl7Kk<B+_<0}=o1sqfil*A!OMg#ZfLi5h
z5XtP4T8h61KHd7C$O_8mtMJzoIeBsE@U3^%xntqgf4|NvW(rVJjD(lD<KI;FTS7>K
z`R~8}Q>sE#tE;Bhm;E=d`j;vHQ-uB>inV_SMH2A;YAQD3r;eo<x7mQHH<+_MKHa~U
z^(8WeYSZyI8{^YXuJ8^wh&z(%+z77{DL^gBT-L9_H?>3-a#(Avtu-Q!M%vF(#v@*c
z^1GX|jSf?*FI9<}vUxYlg>Q8vt+E{W-KWYtqiU`Xx`Er5Q?Zt1okaXGs6}R0a4hKe
zLfx$@v@f8wf~v|p)DGXi24*O53?6emtF6IUg>b&P+ZjvY>nzM{&@$`~x1+}Kxr@Y9
zrf0exSr01lMg6iU;gURrdpr^n@%x)rKv+;4h_C42u5d<aCvYIx+tsmo4%z_@yrI*8
zY`A0g9RN3<D{uEkj=*kyObDO=!J?sHl$EW9sqXU!0Q1n%C1=$F%KVu{Gxj9znPzlu
znQVn-`p0OF+47D_-J$p~i+5cfT{?Fp#B*pf5do{e<^nylY^RU!BpoGfoe%7dVFAI)
zmV=TbZ6=bWyY%Y24qT{b+N{~;DEZX#FMVK>3_0b@MT$h!n<1R|b+!&*uPIpSqtKIE
zJF6>OElhha%SIj*nay`wEidc5v)Asp7s6g^A3)bsDg9*Qmpw;PV=IM#ns}RP`Khm%
zT8q-Q;PmvRw_P;>m0QZ*rEqz20J%*C)jl;HCF*I2@^sd+qgHS%n@?p!dC_kjC^T@`
z&;AX9mK)1Ye*66ED>CyPI{ub`k=p!O#pPPVZ9S@wbuh>mdXROoODwe;e2h1DIz#vw
zMUZwm;V#tqx5UGHPJTF1o6Kq@TQd~hQ8a&8-~oTH7{F37IjCX{h!&j>tLMxVeJ0So
zb&s%VUwy+g6Imsf0q<_>Y14PfQDqtT5^m@PHfEM7f4a!^Yw3q7oEEo!XME3^vU6tl
z(dM$h=BO{_P^5&8uRwgjEjRYD>Yr&=xrUI&UeF<1yo}kcbuPxKM=J%RNqqYsQ8ssA
z{e^XKSHj)r{xLzfBJeimxFrj2^5Q#qKB{r+07Ga|OaF!z4_hd&>u8sK+pg*+eGcCi
zk+SeDD&un9dsH)M;V6&MPE1r*YW{;gnwW@Ij3-~)e=FwSYQadIm`Kvi8{-%+MqNOi
zHAU-+g9#85Khtt7()iX7kL#d(=egd{on->IskDjeWv(s9Orzk~`|>^rPrd>KCpEub
z!(6a*8IJY3gSC0;huiE*6_(GU6jMtd;N&k`A6W%Ch7&4$a#YqDNM9c$tPZu+kK|B%
zFHJKjE|34<I86`4x{B%NaD*fVHYQjWUlKNRQI5%uHhLV|r#~~@npZl6vXPV`yBdft
zze*7h6e?@Zt~RmxB3=Vt&g|;sk)4ovI-;zZ$Sws_NGsJu%iJu}D$sOLw|Qq>8eA}y
zmRUhzFjkFLp`SVvI7OT+{~~JG^66-<b5~w?wl2?DudkP?PPD<N&Qdl-FOg@a8Y5-Q
z;#1;`sXdouV4>#8R;_oK9TV4#@q<<QP7`E+gY_Ro^{s2|+QDe7o@axtofEiBK4q!J
zQvG4CHszz`oKDSq2OX-(d0U<?bL$N+t4mW~^r(3IlcuJjFdOtxmwxy3UpZ9t;t2(1
z@v&?p#`HT3+>VwqjH!`D#ha#HCsZmSxq)&1h1FWd_4JDM21icbuRH?_{P!q8<Kw+E
zxq-S0H(02hh1utFb4nx?ea-1MLg-yfPIWj5ucfCwO!>HDYMuZc<A*c%ta3I&n?fgH
ze{!xx#2Jy1gIDQ<fLimIs_V?d9s+v|LZWyn46nXQ$)bS{5XND1cNt?eMv}TeskjVb
zS|6hn=R!Fd{Uy!AluVL5X$#WmQ-5p}$t%CE$e?e_MR{gN_Ek1Z@r;=XKyaBKH5kM2
z0RmlJVgo3do<ka7`133Sjrhd1Ix#yv&2(g9aO_-kX}U{75}xdFXL^>X$8Vl|r@s9?
z)JBIO<5Yqr`&5OpwJ6M8t57DDt&wF!dT6M7%ca<&#RlrOcNKKvB3bhjN&d`l@nK1U
zA9k(wXrm<3<Ah2QKsOL6(dl#FZiVS%xBr2iD2$=@-b4(GPRg=hp5^4|=vm!(Z<{te
zbK09Wbo8s8k2%`eM~JrS;4U@a>S)t(?m=%zCL0~0WMuVEh+{b7Xj5f%8Bn1y38se?
zPGf1e8fZX*jL-p$DO(L94JiRyC@4A+33W+$eoq3hw3Eok?xUE$>WONxjK6A<(4s`<
z%zBY#tv4hcCBeLFD5+M=bo@kFR#>4w9dh8a$VPO6^2UlZgV!>=HIK=A_ytq?hTeOy
zS`v8(SxeZ%(o#)e2-_}y5cn%}wrc!#{O)4XCfLB@bn!Oy=@dE_9Dnf3`(1@xa=YAc
zkr`F7x6|ufkcgV&Z~xEv>IMG0Pd4JoW#v7IjL@H`XAz%#1;kY;`R&7^&gB$7#R%&n
z7#mq(0jWe2?W%NJ&j%ky(~#JJdk#L~elw~=65?$=zbhK@zfj#1uTtocOpG8v$d||~
zD8IOfnLJyOy{uQ0WWh_X^X=A#zrX&}Uew^QQRISmp!3e$<g3&l#^aTUi7->4QAzlq
zX1PxBG>MHTa7;Yiok$zlB7`#~iH?3K5N$B-{Y@z9&{~{77}^<OT4{xPSfj|j^}ap&
zu8C`=*m$SLrT00q!sc{2j1e-qeVs9@rsabseMD_3TkhwJudn}1)WlL}aO_DQRrlP$
zSxn)$)j)ytMymya2|7#xT6GtP;a$~*9+GQ1`(zfC_9<&2k~WI%-1|zQ8GZe=srb9|
zoGA!C<W9$R0h7rVCCx1}*Uq+#7GC7zjA(SgP(^)uA<c)~=VLODil378&IG^G+^>L<
z0~*(Aqt3Rs+f<u<?OL`>f@8-wcA&&}^k3}sa}&q_mnte<MEjqTi_O%G)gWY~)q_r-
zpw&5F02#~lAL>UhSl9yOR*qPxt>quUf2fmZ<V2n-c5oBv)=f<`!{4|RP)|&;ffbK2
zB!q93<ouPP$D0?#(=Sd}3Rp~7EjM4~>e<0S)N8Z8%;hfWZ>b51<n3`>^6g4@sH${&
z#q3(T+8R9!IS`v8w%Zj}Q000pI0(^3WL~YjozvdHo*uF+YltS}h81>i<&QVaMYd=%
z1;-8#xi^<&nC3T0k^&A!uVT|6yroni$D_SWlZ~n}px8ujWq!5NuF+nSJs##%yYnI`
z5(5lWjF%WhmIO`1*4Y(@n4xl|@%E-B8s(Pv=z#XviM5eK>#r_>qt*&Q$3To#&d!CI
z!a&(SL-ctyVzu0849hpVF{#qDp8UO|vD+nAYbHWXj#%r?ycRa!e*9JCXuWpAKzz$k
zYeb+qBkz;tCyx{Rm*g-j+eZp@1vzSC4%qd9eNMpx>ivSWOA0{TU@+whxR+tnWo*V%
z9VtBZ@y=|-0hn(an(R6|s5Ozg0de7Xv$xL4uLjL<43}+xK!EX$4TZ;g@{5Oc4Ai9?
zYUqtzE=KIYDQwq29cIS6a3*oBjMA;!|3<xc)ImgO+=j34z-Su6bK}z6L7&PJX;)z;
zcYdzQtUUJssw%MfYFjL1&&LzG*FaNLlo_8xZrF7#rRy49LShbA`o%4Ng%2rcv@OJy
z%^;8cYB+;6o#J;>)~}ID73N<Z*D7cE)+Wy#K{GO3Tw_?|CgZuOe(PFD5R(yT*j5u{
ziYlB|6X!|RZHqO*wtpK^igfzc;{gvqTOAu^eMu3`Pfwn|zWr$B*K(Z?$>XSFk$h3f
zveFlq&8MSEy|ZSI(#g5dDgHk$&=;dsA4J>iEr{IYFOS!-qj&B28Q|7`Fd>olt)4~X
zWTy!;^WH>8ku4hXxi74%;a1qzPe#0?Q2wsLsPY)e7p=_wC$a(h6KcxdHx$LGS(gDP
zP9S9`6whxbHr+8!20rvTDAO%dnW}#G$fvpM<kVoOLd66rm+-Y=3Z=q&KCaeh{=38$
z2WcL5kn3Ao`LulBY9muC4r+<X;MfYO59Af*Ci9Os>yUgYPC<yOBrUCKjSKEyFjm*T
zqg(YIA7^E(HeFG3uNm~AJU!l&6gEvpu&%6kr?QOrrTJFQhTjpnbQ=R^6l@jpUv3-o
z3W@cg{Ww>I<sT=Ky5;wY2EC}RvXQX(^MwGLA4V==*p$`<GZLNu7W^cFMiqE@M$9}w
z!@D#->)^;9)_XA;1<Rp;T@z*zi~j2UC8Sm(heKQ)9|@=*`qHjy8Es{&ms#Cx#tZyl
z4bRFflXdEK4{(j*KllzPHDtT)-|S+N|0f0i=G~`-lK$phkA2teI`~5CljSAxIf9&f
zB(dk6grw=bH(rg)NsPOVy-5o1S<!c`@%%VRIRmj?<ldKgGXnExl6fh7<W-?waGYJH
z_lIzO)lB4DzL~Ibco>Z}rCh12{sTi@9FoLs|AJnkfcED$|CtZ`zgL^DJZk$xK4xci
z=vC7o-pc<&#91f=#|E0Rbs&fh`h5H|^LqX5dL~?2S}GNp|CBx?3kMs5|20?l@~zdf
z%s|xMHx{E>^*!^9{3`~v<6DBLF(|Yt_2U>&3~k2u#R@>*e&XrpS%~N2DsHBT4uJO_
z(34q<!fW#nJYw-^&yNv!U?-7KGqiRoW{?C3Te3e+$lnr;k*6_Tn%u}(Vq$+9DY2b4
zE&ONbG66(iqxbXwLFC^X%tKMYC47hvuF;QZ$V+nM6}C*?lJP1uX^J9x83`MuEAD1y
zPVg}7_6y{WUdB)O_S^P|O}}q9a1EX@#7?|XM|Upk`fI!sXUx|@2L12;axn{(L^SCL
z{6y`x)9LNg6LM7pfXi54SQb~Pw!FiL-S*{w<UhM>AvYNuxo0kciLnCtopXAzQYCBF
z-tmFHvgQ#F$s)dY|3V?_>V81Td!>iViqFCP=3aWN(dHL#cx4%%3~9dXPVLGxOG)tl
z&XN^>zMmv^&ClQX-J(ZIpv3gooq%xVx<Fc-n9wVJz0+r(FsW~C9sOf+Qf(@Nz+_uY
ze~b}jA;ws!Q9GqtyyVL<A1N_a-A`;C<VEGSwF4<jj8!f)kOGK&-y|uA2#BS#`=JZN
z*r#D#mi?Axq)TM1kw~+vT9I8EN_aA<U}kki=tg^|0V$TEIETW|2f#d2da*v*Xp>#@
zQo|JO7Q(ea&z0I*{QgD7+XT-PVojhiy2aL&=$$lT*Ik;DFv@5)Y--XXBIFa#f^1V|
zkKid3Z5NHcK9-|9%*qq64XN^xvY{y{T64H3p#yV;^8ne;%Z$QC7$+_GDBZbU;$Zi~
z2I3@vU+G484b?VL0lb0UP@BkEwaPU$_Y(}2IYLW~!Um3K{oqGW<de3{sx!meoR4Gw
za^)(9G|-<Lw?XNAR(Y448?lble2kv7jZ-MtxRtHURsL(FlRe}Fj^?Z1$xP{N3$J!A
z^-LZG#?4cRC8S_=R$fB)*UKt@>qJWIc^9iBWs*9M(Qfs9?d@!FEGn%p%e9o|JDjSQ
z^SRLv&W=+N8eMBbVdSTy%bD&42(UhuO+pXt$-Bclr7csk`)Qycp?aQ(RmX1SH$VK~
zA7IXVusTRoXiZAATgTvwyAq1NZ2c4Yh4PFjiwX1o&s}%!vty4Eod;jTkDPM7s(Z`G
z|MF*VB$wY8I7TBFv~3lksxJE4#lHJXTWX{trS*`E!j;c~UqDtsW!VbG;yWFG*j^W=
zDq%V?>hF2W-anKjB7KEui@i^Bud|?a`~kaPEjVudpH#11P8CRon#djcn6Njv6urG6
zt}Ht*h#9^dQ3B5X5*Rzt;ay(~RXey(o@lW8S!~SrzHg-&CeWbJN<%|vbfd|<2<&W|
zbLbSjH2k{>U$hrv#;t_iSvDKQYz)L$a^I@B0;sSY?wpiz_s62~uXZLDZfG-KXaAYt
zjL>0WQsI;!2jI~S&eD^SY*i$l)rDuD7OIf<O0PF{zcrdGCPKeIyyJY#1OK-geBV{P
zmJtIo-+7q%?X$x(oWL(81U;^FnkXzY@2g9Huc0zpZ#ChOsN#4NGn%urPa%YF?smn^
zeDaabXp%N*F8ShE<aupIe+wKA6ZVl{ZZ;wf(l9=;oiPd`HX)xh)m)A#N(+w7TGLp$
zPsKiL(ee!8D$!5_`Dr0Jj4YRN2z@3aUoTc;uIX;^)rZ|Q#MDW_Kh=NJ(DIh+fLxCq
z6;(kFV{{~9BRwP*;?t~I!@_r&+87^~1DctBGmC;$d0G0?-1r;6TbVK6+j@w-b_RU>
z&b6pyebNW4Ec(u)V<nua#9<jypZVmZYhHAwx<H|UFU_)_SIv{pESr~;`JpWQ%rrXq
zC&J0kI;((r;%TkG1I(HZRal#TDK<53&5wJ04_3IQ(Y@_IFJRP?U%@msj^h7KIaFo&
z%>IMaG=uX!ld|}>V?~z`9_B*B#N^a&NMJ~8m;R~PH|d54&9g-<azq=n_{`O?Xre|M
zefs(8HvNPK`vzD#GyaiSUMZDX52`+p-;y%2BNzz6q%$u%zPt)5<Da=X_9;p0h9<rg
zksd90n~(XcT>pqM1{n~3=;oG-XiK*@4~e}m9q4!7rpCC!g%ohPJ(c2d1tgPItDaOY
zq6c7DYDbp{c1R;3V?5f%Zrve9!+i(3@Ol~QEUf{ON@WZCESz`^=}&v7cg`i{es)4n
zr{C_&I$J;yjL`j(jiD9o&l8O$%yg(aBv?AwVC=x~PuhGXUq&odUa1xs*>7C1Zwr4|
zvd`?}%5;fMCdL{Wnzj7^>6sb0O7l)n5kCDU74^MM1P`;hUDuocIjc-IB~N%>bM~u}
z(5M(=c&8@WR1u)X^RI!KHu=YU6ZAq$9UJOeAqNR~%rz5{selCKms;?LyhKXJsl@Sx
z>=6*h?|222b^vS8fYaUwS2sV*zt4y=oCxe|)90hJ3TlhofL!^<)5#54iCRMc_$OI6
zLzKz8?okSbh@>|v^s`T4$rkT8kfbd3Z5^XbiEN5>io(`>fw#W>TSzFhT7I~tFfVVf
zSIv)@KJeb&%wywsr#QSH+6y8j(vNF49^96mIY4-NZK=*Ac++BSOLyh9=&f|2ewFIX
zDTL9}G0#g{-S?tlv7Z^HcU*O;jHljNH)`jE*`*e?elw;rZq>A~>^$mErg9xiFBz;R
zj?3dJsBgpJ+11H84a3El%gl{X=zj4Aa3u!e-lADn<v;jhT$4Cs34dTBNqMgiIp9Dh
z7Ngk2HK{#%_yYnT$k)rq6&E#$edne~uA5L~Vir|MzMAA1Pz2L`Dl1tidB}uPBMSUi
zKt-6^B51zeiuN!ge#L8EK)*#GXIv9r1r~7eA?V@~lboto3RPVSI@O~vBsQ69srB2n
zxhvBN8=Mjvb^@ihr;LV0lvC>_e2|_o7@c(S9F9Bz9REK`3el(K7#*Wrra#8oy$<VF
z!VyIjUcor24_)d~o@|OzEB2K)alc<|nzpps0qB(g54}EHr9O%7aJs(GheD)(+6OJ0
z-KU1E2pPL6-$-pMnzU#sVnIL4|I(0Edo9tu_S175MVI=Pq@8rPV``+<r4LmHjtvJ-
zqoxy(I`>7R@h9J!mpT7H&culy`&U6+85<1_(qDaC(ur8y0$2L)rFU`-<SV8TB7P{2
z{s#y)u^BQ$S!Sa;7q!l`B;t6cy_p8jpE4zG3jt=9LX|RslE9C4b7J~`50aF3(UAlA
zH~9aNbY9hzK?eVYD&7Bo6Tm-K`rq~T{{mmjRn4p@c4%oW-k<)6YRy`$3Fz<sFo5%n
zN)A;e-Kg+MuOh<tFRGiXM~6hdePJwzaE+uDO9PdwTDglV>7nXJ=(unnij|L<0wrvR
z0;1xF8fS?Qnob=CvmJJeS&rB{&24X<DQ-qrmT**L5@imOg}m9P9q1c&inRBmi0vK&
z?9bl}aapI&i&7F1n4Ij^N;7pJVE_4jDnYa43kzxc?9hbFC8B~|(NduhcFY)#ce2(B
zy3oF2@woTM@(ZA~r$=)8Dtd=A^MJ)~{iP*lI$?X_d0N2>+`z+J-)wI_JM)i!pId!j
zW>W-_m+$1bDq$^LO)kex1oDxgFsZ8@T@oe;8+iMa$!0s>i;cL@JrTA?(XztDZj$IR
znSoK(qyuc)&Z3ycw}wFi%;4;emH&;TMOD(qnv`e#ZU5*bm<xWAvUBrc&dlP(5X+@n
zUdv&A$Co<MD)C^_+T;3>wRpXU>dX(21qbzqE=;CNr1R&t=c#!|^4*%<s}>XL*|N98
zRD1#xI<4VS!#=YXLC$3WjpR~ylVPt|vN)Dt)s{6QX=iyvTewV~i2Cxu)3v*I%WDsw
zUVpt^;8R2`Zih}j#LNdEDFV;KLB0i%wm^X%H|TkMfsgpcyB@|2YOn&!^3GbASY>7z
znk7`CITF@*r?bRUfJU2$(1f~^n90t2y*V^ZJh2ysXFN}^L)RJ87IUvhzR8_{FtdH&
z_m_fWlwL4B7OQ1GKpT)HFX5G>HQWSzQ6Ajbq_tV;$sacv{fnLs`^}d>|9$81$^2sb
z-O}V;H;eiCw_eQRDdF8eC@l07f)66Q4Q|uT>Qd!7FYEgMnC?4@Yg1jb@hacQ8#aeN
zGnd-Nu+!Mj@Zie4s-I~d)UL?zuCxP6hwD|%#2Edz16FTAo6NtCV12t?qZQ%;Y8}wv
znZEsZTGNJ~+I<GS`c2Q3!>4#MfP4I9Db&}VbVru$4$pvO_eR#?&XXG-Dce-diXRlg
z%exYN>Z){@RVvc69rR-y{iP5s<a!m6#KLvBbDayf>Dxg6w&?Y*B_S=@l%1bSZufj6
zbcihpL@9%J;l@tqBXV>3!w0hcf7KMyiLF<6eRYx!F<bdmYL7YMkhfhsdS>0mi)VaJ
z79K<&w9#8>_VKU|T3>5fzw_4Bbg+E~?|t7_rbnM4#dX&Gc?_FM#?RE)$(|*spS9bQ
zwSk<-(TOR<?i~X?I+DivTZl2mvHLc*@x1P$-#kQ3K9&lQ?@94boGu*BPnORnnFAR5
zLl(_jqSpAw;>Pila_9zZ{k6ny-G^#{L`r|9nzb&`o%#ZI=-w&|R7`q2bZ<iR$m)HW
zgGU`*ga=JeOMQCyiolJwe7IF^(CcxfxjF;ZvJLga$RFozR#s?!M(OFV<jXGsT>?3A
zPWe)<4690B9L<IOnyZm1_gL{9fqY-^O2Qh-bEoN?)X7~f<n5_OvKqZ()U+QOm?NDb
zlm!T(c=C++LPzvvZrQrZO_<{NuSDB=A4u<qrHVo-7)xZ{#BzUr8C8dVHcpB*x~u)o
zYb@8k@NGwRUVqOAqJGJx@Tuy0kNd+<0;NjAcbMpM%rx^coR_Jcye)u_1`2?$ep>b(
zm@Ttz|LJh)fiLlV&lHX;ZLdMO{`0_Yyj|LcuqxYOT0YOPzSR#q0_!@>o7GD$hi_OT
zte!Q)#@ek+$_IwErVdS2v^?Oa?J;H1<P{cVXensr@W!*looAn-4u0*Q+Uiwr&7_4~
zq@;#BS{Xtd#6K2GT^1hQ<G$3@OK_l0!zj*YjlcI|&I1X_%fGjFH8!!@|EiK+^d^ya
zCtk-u&nLtF_Oqc#$0D|T1l~-9I@}LS1d#fACgnDJ#TyoGh8)^<IV!=~TPpGIOtBJX
zD%{{vCYGTY&XCx(a)<4<@XflD7x6Rx$x7@#?dYweNc#;nf~0g=e4b#_v4Z9Cep?{s
z_4wM|9f(iY+GHhBaS}xsqeV!XAo65@Zftu6qCIK%n}bbyi;tW=wahbYd&H5(+~^hI
zcN<Zx^b@DN*H|ZI?-o2y(`)!SC^=bFSw<{2zfP9RcH8K&ZrU)~17Z_+DAWi7Alu`F
z^38P<w^c5H*0$YhQRoSC;loGorGu5p&0YL|P@8tjXaEn|%G<ZV0wz)XVEr6yQ@KWv
z5*-P}1pk_4`2~LrKCCNL*Zjj1pn33EmiLd}DHE*Xcq94K6~~nx)tsfkn@xpMbLDfH
zY|8$t_qToRjf+~NZZAJOP*!+@%va1nkYBEXmgI0Thg{^P7N1aL>!=~iHV<YD**qE8
zD>=)@>AOZts72C?&m_eu(sb~ecaJsXst_M95+2Vc&cT!Nx!bOBp#$oxifG*yfi(?h
zwG9s!V)c-0(RG)Kns<RF!U~`9>D{7P%ix4gs+LJeyMLpR5QgJ`Z>B(9+Ex*M&sG0v
zXAn*1bQK%=Ln1Fya$3@4dBaztU~fs!9uoL6`?O(aLLP1mpWc*Zz)XpCNo3nNmry$t
z_LKZRIiQ8=OJ43ye_j}pSCG|nU(=a&QqAq2QsR7*asPhK=0stvrfHz+lXbkL-*0Z~
zTOcB(7qM%;8)29g`djC^I;ti~GhPL9VB)PcV!?NyU*4fFV9T!{h)*l%HzKkx=0uT$
z$e!p*KHlJA6Qu0Umz@l=8NJXv0x-DgnfVR(T0B!k0T_)L%Nc^+cMqBI0J4gU<a1yX
zQ;Q+!a_KR&FKiH#o@-uQ_{j%8x(cw06M2w->{Ck3e=xCRX0|=OpDYgroy%myw1VKC
z{CcJ}xB6SZ%A{34tNfBCE0MPobgCdtKk$jmn3!rvxti`A5Tlu~UHeI8?!_O$E&Rw1
zMc%p@(2sscqpE$tTMlSAGYpAxRs7ub>0!UmCjzXH--QR~qLw3fne|~xSmtSq)61#U
z*UaMDZ8ImW;QCV^|Ma)5jVLviQL(O&s@U-+h#gL7c*EY4l8gE+!=Ta^J*-Sdt(lTZ
zw#{DIM%HqmRAcnLsNRa~`&KI!EhopDecdbZ;|nk61?ZqnPvu2-<_gp<VB4x`UBZH0
z3z3Lqx2;qeq>rmhnAWs}d%b&Z6#xAr@*K;+WC6jreMdFNz`A7Pf;IZkQwy{|&yJzt
zr8VAGf#LVthPY+Loum6y9n)qaoc{E@_#c-^cIqEo2ca3LCJUrcmM42rEe(a;y-NQ6
z9YY+<c2<tRvl)JCck8|eIHJD2T9(9h3cY^yVAQU%HT1UGh~$B`i+!htd|u0*t$tYT
z2O9+ICYV(`aOl9_<$>g2oaTx`k=)CcMjb_p9K9Vw`?CECv$8X`;tO43uH%lpXvdK@
z$q<1lo>H{_&$5HI(zAw)YncYVt0(mm8|ykY{A*ehzwTL@g9xaR<h8-qqa)7*GkGGv
zj;$SiQ!BWTm~<&daqE46kkhLe*iNlBB`80Ryr4eMB0IS#<OE!3WyAK)@Su;+VMgDT
zm8C4RXr7VwMLiX?!LzbieDc#?I!m_Bu*n|r_b(!y;9a~c<>fXJ@yQb6xf<+O*{d3S
zN969wB$l`P_lxHqNW%nt=EedswNvdcbMqtnVr8AW-cGwu0N=jGqL->3uqXdrhx$So
zk%Om}{k||SgG>qDf-gXKld|mpbG}YKJ~3CAOp-+pG!LAfENo+&C--95VlA#@Y>+f!
zbK<nvywkl-l*;>N{>Df{MyS5yV5~<zI`T0dw}06j>G6{C)%7lFYv#wsp)Ey+pvCd$
zj>F|tpTr&<BDyKdeobo1O!6v>1u+M0n6;ct^ekEMjrAyEBldJPl6fE_`~0GYf;vI*
zHu?aU0oMTYy2a9Br@+>Dnr!n8C3U<HkRj}vMLOi?QB9^oBGQ|w(Dilt2d}9yqsRm)
z%DQzK<BvIk6>d8Hm)!A?t?jO^{zn4U0`|Z5!FGyZ@tQu}p-z2Q`K!8GBW|z49@*Tk
zF(apRrh{E(-d7!Rl7>YtyDZOxQ1@3`a&8uXSLye$)?ju3z@Xwa0%@UTD>F>@M8y6v
zd6>YZnv(B#*zcw``=1(bl^9}SW-6ROM@p)-EsnOHh_R}&qwU;njpQfUyt1{rbJQ9R
zWA*YIpWTd_w^~N5ipCY6_QK8y9JMz;)16bXM{RSM*Hdesji*<WF0C2Ak-t%V%IO32
z*IpY5P}qh9UTQV{N&H=sXnL>I?_ZhFFI}gSv?X^V$N}7|l%rn`L8GJ3J%-9#QpGm{
z_O|l>ZRQR50Uxr&#wxbuo=U;qU*krEPLINzx-I!zSzU`VYs(Hk=bdrQ4k<6{PUK6P
z3%al>h4i^ktXwZy+xTDXop)4|+q&+t5tk^jgGfL{L{N$#Ei_pwNC#<B0*G`7NC}Y8
zRF)#rq)Qj+HPQlvCX3!{2qg3#Vn71{A%gb{y4K!j@3YT6_y6l(j^P+c<~P5Y`TpML
zdEN^)4)d12kBrQmvXr#15*Nxp(h8GpA?=tW21#4Qd*&d)g>O!A>sc$?u99W=8**8y
z5zhVKKXG}6*}$2!j(FQ}vvsplu#p)eRiu9*eWo=d?BRmp#jE`;*YUbX^idNw9hHx4
zT?{0Z2=*0t-<Bim{nfrEf%jfzmvgm^<Q2>G;{!elC><lSX?H|M>FlJWHcGdkeUs1`
zqG{#KLF=Bu7V`@$MtC2unQP{YSw!j0fx4fI4D#ctcItAt!TNYpazxhKRZuJzFnx$?
zhNp`+)YECQZTUSt%9#cJo=BCeX|lQHBiJ9`$rqU>qX6|6>hg4^D|!!w`dz^)so}x7
zjoWx%cp}ThaM|T8(FTUaPt?94ayY=Hlc&K5gncCT(TH2J>eUk&l10IHy&pw;BL%+t
z)1vY#B%6m$eoZCj3Dx9=y_Il^{nN``YM524&5&YB4R$0bZinnDBs4j*`4-3G&JpqQ
zm<5hG%ESr*MSe>nRLY2rh?i<i;XQUg*2VCfu#}44?{#gUWsf0!f`K+O?6^E>T%qtY
zO+cn?>=xPmu?dS99ymy<SgG-xyi*IiajUODBc?;ct*Bmj4C|{Fm72#$sh%m6VBlk_
z7y`1S@czLL`@*4&vCZ4%w!CcJcbBMeX%}7%`?~NV6PSybK_hh~LNAL=;9iTdo&5LC
z$_4+2vp<i2Ao;~qj*MCp8rqiGsw>a;Vcdg1qWX4}I{eJhCPoUX6=6OJay4MHX!C~M
zowT`ft}Y3$FPgg_eF*_i`<k?i6abO5xqMU2-hofVVG|M2wPzg{)ptIGTPjY~d^Ifu
zKOfBv|KOi1(7<%Bp4Xr*&uBtquuo=um_~43{BuBmOry=4UA6sw>^i?lg=AD5YlyP{
zx_R<YOVkY>ZO!DIL&U<r-+_V#8vj(`nK+N*=kMO)l!y;BPe4{Z-HjbOvF1|7N}l!<
zi1QJrb{dR0q;D`Zy_G!X?0SUM6=Uv1GP<w2YPe$T<q)wP)MD2(_gHXoXPmb+1m4Bs
zH0*)Ap!bCX{>WAA%tdhslyDyZs;6jmxHWxLNUvXUL<|=x9E)!Ahp)cU1w`0*n~6g`
zdJ*SkjX+)T`P{hoIr&?PZ&SJ)>INNUTYxTlj0LePpCj;K4#MN9Pg-630|mQlZT5AT
z#LWU_3H;4DZF-ad8y8%UD@K%UR+4yq`()N3`+QTi=b!Fj{0*v4pc@voU{}F?{8Ozv
z+5?Qz28*K^pgQe#sDFH>FZkUzZ^{x~pMll)Y(@Wked1`6z3id2c5?1@URur9jy8Kr
zpSwE;nkY<pC_MaN#b~co_U@zscXa)0UW?VwMie73aECsz6?X07sXu&5QX^VYzLr+k
zAn6=~J>y;71>n`lDK!N!2`ivX$Wd4{_AmU}aaTh7QXak0Sbe@EmhfYzHibes{LJt}
zgN2o4;<)<yt8ae{=yWTmj_PGwBdYQmF3Q>PuxXwJAKni1hqyNgfN7zi6X@^PQhKt*
zRaNb4JU<#&aS>HOCJQ+`F0l$})E=3~60TQYgk1<-=LcY7VT0$`jv11W%<8Lx06*fQ
zHD>IgQdrQ~Vky(4lw5l$c*b+8c@aDjO2ygBLr4Q)jc)4AFb@HKJOsTffeRws$>93!
zEeVQ~%UB6oNVesX7vrO;YT*p>pBLzv!1~CUK=Xbj@Wq_}I_)sIv+bR;8A_}wLk*+D
zw}3<|F&V#jm4mNsIYRwXNQJYUmd2<2v0j2GBo`U42Ii?=koh7!^$N51hq7^+dt|?Y
zKj8xY<($8;ao<K=jpm-0AYAXydK)X`k3B@eG>jF#;AjStx39~UP02{TfUJ5`A6aZ-
zKt&1rThQioY~b*SYB2h{*$(o$Pl6`rVS8yda%qw*RokonL*P*dV)Z&KiRp{+21jU&
z^R*i#1-0c@!{^a8xS&wLSv3f`J9p%SaeNYL*1!IgqMqMdVJ~MJ^81ZA<LzB?0_WqH
z@5FJJAl5i{|Am;B8q(a=4+%n&=Ah!X#S5V^Y0~j)cN$!U^KO*LeS65UFk@=7C!`<H
zHGgF?X(&(b)hrtUJDsTP_a(1Ee(BH3;u}JE8#bfqRUi>-$1QeVdNN3PZoDEb!goqz
z=mq2(GY=dgS7bZqW|np_GzQ^Y>EM6ig|CXzo`tKornZ5CRa&aM97@q2cpXFoS;ZB1
zFxXVQN(u<*`cf`Kf`<_eM$bC-Bs2wOUX98h83!l|f+4%LlRe%!SKyN2)Chy2sRO0J
zS4^1L(u~o9Pp?Jhyob{S8CB3&Mcp+beoURXh<BI(n|ban5XVR2k(o4Ewk1|2v+`8^
zgQxjI8)_16wrsvlo_lD{S73`n{GqYL;kWRQCsYf3sasVSW#-xI)_KPy%nyLJI@eyN
z1yL*BC;1=WkisMxxz+2XyrrmrguR}5zs)HodioPRd&}bLbivV$WKT_{m7-7eDd7b=
zM1FW&k<4bYIkpCgTl)RS!<6|o<k0IbW(U<wa#G%uP+fNBZi_UaI=JtGO*_W)>E+4L
z=GdW(!%>=(A8L!cDRm{xBaSByBbj+Xm2;LFpRnR&p+#C+NbnAj6Y%`59HeaLO2n4B
zv2Yme3)e-ZPC%*RS^L#rfQgCicahVE&U(fny#lN)8HkBy!%yPuMkzv91I_ESX1p&h
zUAfq(U6~pe1UU1u$e?Wvgs)?#?x!cG<mJmu4={O)PkJkyV|vDNnLU^c$ez23Ofy<c
z4V64subld9FqX996tKnN4dNl<g==#q-Gt(T9&j4a937-DsJBSGX&w<3B>P1j;pr<8
zItNH6`b}1I4l+Hx{2y@a|F@(u(Di?G<fV>LsF!AQ_`aTnH+xuUt)IX7Q%f&l3uqVa
zMUM8YQHcB9GNv!tCmD<yuu^nb%>QK6y>USv=0RE3L>AO2_mJ=Xq&!4T86&Z=v-Hh=
zVB>+}0ddo$=pw03%;Lv#M(`s*ypPeif9gQawmkw-Qt$a5Z}Ar|$qBITg^HHm{q&@2
zJmfk*QSFGzKT+p~vmhh3*}zAWCK#`D^1FO<Ev=H*FTx%jryS_(a1j3U_VvWzj5|&%
zN{*91wepE0g443^^4~{pjDH%)3h(&~Mqi>tfh5nAh*OcXgoDm5((HNRMI){^u3;MX
z^*f&w(SXX^BjbgXge|+uw-Jw{-s%{mx64-)qX1F!LiHPrZORJ<-eq$34S(Ospu5XL
zf79Lk1P0xGKLNxbdObf?7|prz+2Zo1jen|pXgpK`y(`qx7^hh}-w0`Z;PPv`qO}7W
z+f9HR0OY%{rj$pTYgBmUT=gUFgC&>QjXbRz+k^a)KhuHJa9u0UjH|bN4QL+xR|)!x
z**B@~A=X8ibEcuq&bK<{8vtvOkm`_pi%*HRB7Q*0Z}V)(Q+jT8d?kpdZ0=3V-q+*L
zGgu%=eg4)WZrqS&lvyDgO2+IpEGCO$wg%LH6^F5tZX$!tdQTb8>r(>ld}w!|^o&v{
zObK`5m1W6qgZyozDc)H?xy~Xc67?YPV9BU3v!!Wb6bF#lCs6O0sbtONWofflC+e-8
zvAzjj5feYCF8@bYty7qet0;F?UU8fbOhfTpox<E`VL%c{mf5mu^v$Io;ef&fnsy&+
z11cFN)v)L5L^gKTPgo_q2{?}F8S=Ij@Lnwa!Rn0La1+aoDq6b$0LB!I?CGT2RoBR|
zJ5C$dUubg&%1`wG5jgX;?NHvFa<CBlFmkg5_mNT=o?1U;DB>oP99+@o2IfLgJsjwY
zfcs*>qXKV}H8EgH6iX`?N2ua!asX7?2Mehvn#^9fgbqx~9A{NZjdpuL7)&oJS{syS
ztP0a)C{O^*$ugBD5iy3DJQZ(1xcouXOM&ZvW~9BM)Vra5sQBP^0+1Q{W%0VF*lq&Q
zz5{m{O%_Cfb|W>(Y@yvsUHXGb0*&q1CoJ5!*OIZtW9u&>o7!%Uo-KAMBD$kav`4RL
z-Da-vxNvRT<2M7J!RHlE%)AyY`PTWBk0#nLEGzPS@f|0;FkuBgZ$;u@eMhe~xUqem
z(`3e@!P5(McA+eLjeeL%5TA6})bQX(E`(dUS+9qxCR}o*Ub$utm%5!<F`Reh<b~I?
zZ^s+ua)2~AR%M^+l;M(I%KqChU9`;4o_EkE{J<?a+&6wgaLc13OwOy#9uO-qLG+}Q
z9x3dO=}qP0yp77TBI7jCvyK(v6QpJ31yA0C2@<|%Grqg0*s5suz^rs1F^E}XGj&71
zh&T6xG7}e|8S0@zs~}55wi#hxzc@2rur>-U9Wdg7vF-~`79I|jms!BCN<5)&c0Z7P
zW#nI$_pP<s1n?s|vPyJ&jm6mHJgdk-Bnqp2l$Ks%=u2vQM}lOv<I9*I5MD$^-+TrE
zS!S7cx;I+PBO{8QVyAxpI`U1OQ}#_>;_QTm`?C_$v}o#yeI?4bV)l}EmEfH!-BrXx
zG;HTACB!2*&zC>u`#7gQWnlGNxvd%8JWKQ3^Y+XXIegc9&cH}y>P2GiQU#c?8J<<n
z91Dze(2MtdS1GGaj(3~q*%)s6Qo>=OK4hcVbmi@>Gcp&xyxy8^WN(ZOV+v{6YvfeP
zq;dvXij`>%IB|*_vW4meTqNm)5yFnWP6Y4nDE^r58DGh|9=BGd4pj3_aOnW()5!UP
zuA|a&JJj*#!z6<vC^eq%!`pI&gZ9-~{3W7Q<=To*98YFTvnTkO(|B*4P=tuRf8xf@
z>%I4+Yjn#bmzhU*K(kv0)CdHFAckF&78Z5**hq=Fk)!iGsVqh&mWSt3h>*0+?lOwl
z%J6QJ8N_xFlXBxEzce$?+?SHNFq93Q@oAa7FBXfi=MBawNnY?cc?Wb5*8p#P(<mS{
zmW?kyZm{xvFU9!onw6l6nnja}8t~iKLswfKhWSSE@WeZy>}r$P>>VSD_GE<4gbqLd
z5VxK~y$F93pX6;X6%Px;9W5G?M38?18VQRhC$j-}>pPnO!^G?x;uI~&R&X5FE8&SF
z0X;F{eHruk{4;X>qHPyK_ZGXtYxqfP{KMF4@Deod&}t7AJEm>eMG4pPSsdG3Xsmm(
zIprSydf0eH)ld3SZ_ax4{)efx1^_<oHFrFJ)y{KY0GH&-zs1WRXLFp4#gO>Npmrb3
zoX)@nc+Mf@IRklKZoXH=cxK%fe>c7U9&HNfKji)DQfR(%P3mE*A>bQbA-9Lyi>zvr
zMH^CM*9#A+Zhi_DiWsu6lJCL9s#(-eW`!3a8#Fe=Cuvm#bkg{}y3qV^UZxQ<KAJQ<
z4zID~*_I=>?nI~FnSq)INr(n?O9inU=BdH_+GpuG7V|P+-)cVhJ?%<v=Z8G_ty#6G
z+GbGq+=7YIP#?#Og9zcm4r)|UorKz!ExnlK2#@|HUU{WfqSX^R(aMi&Bev@>4+I~Y
zyjB|NfM|?4KejQFF{()+rd%MS2-yP;SpOI^^$%8L-jO<2V^Fc#C3a-q6ZdhKx;%mC
zE<mVpdwE}UZj(t%q~%5%L4-$8PV?i8S0oi~&a3?nbk#9?#LKPMVqT26|AqgF*bmh0
z0l6Ah!oF2ql@no0mVQS8pR%XNbm{lgpzUkZB0lfL@xI;4#L$*qpsy-lmGWjm{^E@F
z3FqAScFxKIj}sjkh}#3!FCvXTu3YyFNoQxh;j&w_auU$?2_*_gOM!AK?ZIzOA={da
zyJ^*Gp=`uBHoaUDRrDP#nKr*K2tTe&Oppt6ONwzZNc$a7V@H;hrsl=}hD)RGSXdQd
z-iL>%L6ZKi!Jn5#lSsYefN7$xbFR3ZJmJBQd4BA068%eVbfrVH-R8lIm5ycOt7nBj
zh^|i$r6Mc*a8T>oV7k+R4tHiA84OsrqP%qs-}KXv^?DER>5?)~?=hkWWwz8@wZTO>
z0>aEX;3m(=QOz3}Cv>4-cyUK!055SE^6i@eOXSD?YvK1kA}13g*OVfz6s}6*5Y34n
zvdg;SNoQR+6=E9(@Y(z)b2bC0+3s3GJZE~S!x<f#`YYk5)Y0y!)Z?~CmCRN~Yd48{
zcrCN?`6fB`1duUVZ*#_s4`<R3CtY=K&{)tel98tX{TGuIY+6VP>jP&f^VN$Hf{*z@
zc!|yiKW)m#LGQ4i4wIKa`+Km*iq|uiO{-bw5=O0pH8D<B&qaU~p8oYTYtPO!XUO%@
zWKAE5vl;JKUN{1Vl=_Rgk|@5kwi3BVF%jRWzWw;srHiNT;rtZD0Cn~Y6QE0sk1x2|
z_Z5ijSSsH2ou{XcRn|~XjCPPZLNv5da~{Tht+BQBe_CXcOXMWk`O#AoE7d$-fu}O)
z<cBf^YR?(_sUY{1@))Du@im6dTQmikB0lVn-|Xgt!%{jAC6*lJmPwAsY<AxaVui=J
z&_4k=8H@wW1O-+~g;AQGmf3SWD?#UZsod|U+&+5iq024L6YU$C5m<F9-tX{u-g8t_
z6`7dy>k165^J;7##5xaTFTGDob-Rua<qWAwvLsUcN>(S`mHJ?WdtncuCezvxvr#Um
zy^yX)_m3DRv!X3Npg=-hXB0@=7lYe2G+Iy0YRh8l&n~=Nt~i>u8x`v7SA(sP`}H5i
z)dn@TZ_-$VR{+N(amCnO^yr0Pi^|hyo|Mmvn=q8%CyI#4SuQ`qfO5$Cr;durrH>5$
zD`La@++QC3|4rJ)e8Pkf_)45hTeyTJkh6|tnJ1P9(nhg8V-^AuI#)f{edwOkp_B2B
zXWkd5LA&HwmJj*y0QDdja5&g>-K8>J=RTMgZYfie;~FQc-RTF3i-QX<KMqh|KL^hq
zq4chDH}75V5;#o+1ToA{F=s(Y3d4Lm=0D<i20y1t?1;Yy6P6*$v7|C|D|08rDP`uy
zV04M!dl%pZnmn`%;sF%d72uXzLeS)&`4FjlvC`z%F{~-eGj75kGM$8RU+c|->{CBb
z61Y}Y<7Y3CfA>Av<<qq(<N?v8kxBf3pXpRcdY!I<5#_#}?BN!`0*l}|?W|;3FaP?w
z29f%HEqVl3!I>ob&N-x(Gj8OzrXCEvu%S)h4D~&i(Nef@GUQ1WTR}(IuFIqsgrLG_
zn-{$_LoBWtp{Y#>3;A(V2fX`cQ8#3KegMw*4l=gt-W2for{4frg;QNpp#O2x(eN`6
z=s+}F;^6!h$l6~8#fin^=jj`b(*BwvYDrg-x4AM^t#bA-$^P|Son*b_nDkTSYFvX1
z-|RtEal{B;QG}zE#L#|W0m##~1#ZV3%CIUFI_D&%mLeOWRzXI?8zu>>My4*+U0Z#6
z9|uTSAUSjPA>}Tp9^Ya_R|58=As|YaFWQMU^ep`IU#pIh^fP`PL7)CEJQUswf^7*~
z76KTZtE}W|_6e)B2~GSJZG^Bd=b}AmJ>}S^`<9jL)@f6r&6aFg_Gz<qk8(ky|0q1P
z$GH1_2T|28o`_j5-8@fM7Kc(z!Rpb)bRU_<_BjlDum_;NfM4(03AO$nw{gm8*_5v@
z(85ZxXzONirMfU9YY$c&Qbr2{e)D2=z|b!{eag9c)LL(Xk^X!sWiV<l<?~hULnW}Y
z2jcYz1EDScw@%It#tZ`{Q~j?n9$LW<S3aA?sT+YH8q3C>qF?sGi`2*W7OfpaIit6)
zYanRJW|>>IoJlm7b<K@C$JtCy%}=?WozYw>dw~CZDGut#c31q*tB=MilWywPg*Re6
z1za7nU)T!K*QO0LaAMp-WeA)lpkg12Q?@)~P26OmjV<dOPM2r)t9Lx;DR1BR9{)sl
zH5$r?W|OA5K3%;NHn&&Gzbs&9AprC7$1hq)opFq}e^X2e7e)vsHAcpyvbuE8NqXP!
zQqB&IVWTAMyF9_0@u;d*e06s%L9ZoC&=t;?mQkdcjeLvm)_CX)nyZf;4w_FDEgudV
zaZe8~NEj=Y89_|#E-pBJ=E<r@yEQm?qPVGf4Q8gRMgo1k_$0%I<~sRr6*B&`?v|Dv
zTi)26!S>5V`09@pU>dx_B*nw24{F9%&9cQ<b~9D2#&%aOrm-G4@y6gO3mfI4!t%{%
z|97t$3i*Gpc2?}Jf$M>4=l8EtgZ8@zxyFPiVmGoon_v?qFC!e9iS9Brtgj-BKHV^#
zGOk?OfGYD%rP&=9B8jR~$G$$#4`v;OtYgL4WD>NbfA3enp$z5#pW_4P@uwKZGi7KI
zoKXrGpyt`B>R2HG*Q1N*^udZHcewE&$Db9aZm!#T)h_QE^!510qYp?)$eZV~xhF^I
ze(l1qT!HY~sQC~is@Cusz)BAP;b<~T(3B6-6%p-d8ZdLlt&=>2YKe$8>}rj9?B40d
z2WdOxE>db%Us4xsB+GnOj~iMp6He>yIYQrGvRq!u%bnLzKKsjM%4UuAziKbkFA2yf
zz4?KQrG8FPNt1rRuF-z9ul~~7qNX}tar4<F86+S`SYB>%^N3dskRj&NB7QG1a3rOK
z`|AN(mACdlo=K+y9vL^kfx9mu4rumooI83er{%mcbYRDDbxyG`+sF*ycD_4K<+r%5
zb!!&h1f*15-yg)^<rTRMb3mD%FAQt8ahgiK0HnEexk%u;q`Vl>z&(xm+nR9&bB7}T
z``U?nqp1;bn1?QhQQ4e62UJjOxHW*1pr0fJRY<3N@lS!ZRe);*@!7H%k<)aVY3SM;
z9T8+rMwc#?vFadOp4_JlH!(`H^SO9k>ZoqHaB)SmUj0=?e>09F#Vhw}*}$zO6B_ZB
zd}y3(%9_7iXLo@H%}xzk9z6+^RnEw#3T%l1HIoK-F-^LD+XzSK?W)gsbn3P{))r~%
z0;BU%Xz3<3?)BJWbQjA;pX&IERkk!7?><0K{ydnSJ-+i=6Q9(=yMN9JFZ`^_K3_7P
zyp(I+85!Z(4?crqdXgwXK6xI}kTBZ8TI_rDw*TW^rMMQ$9~PQCBjC{V(e5{(KVlXR
zWcLrw7zA9@S0<0=KFi(ySRgadaYzaI)sHs&HXQtXZ=*dJC||foqkR#`#xSl9x5-*Q
z@M)8C%YtE}+WS3ztJrK*2Pp34F%GX|qY0dY9AGtGRxsR;+ho%(mTh{2GgKDRQkXYa
zELa)mSQcm9U`7cK!eZa%);C|e2xAKllTOs`pjt3n>Jj5(i0KPGRsQs%Co*g+V(}8V
z96?gQCd&P|QtM2{U9R!Imy1)yH8@1hZvGP)4c#9QA-bHFT)~`opIiMTmZmnoayMni
za|uZ=dN}1Cyk``wMJL8=y7}E}fRYAM`C)Y1_?3&m1*A3;5LD!sXPpv%<-n)4OB_+T
zSok6uj)G@eWXW_noT!V|VUGhH;1h&6?{?8$h2vJ(P|~7qdD+Y@HONTJ=4k5H$z?6S
zw7Ec7c?nt(pFd;jDQ>{St2w2VN4B~gzS)A+4r>I#xP&i;ex{EyJ3uO;`P<>oIm|V;
z%Wy3iM}uAX-D=i-ae#w|#>;K5R7!YRBKD-b0h%A;WN)*W{v9*^E0dVN8?s3@wvmVD
zYneRs5hC<einUo}I&uMh>KV!r9-u?P9a|aR*gizoz)je0sX44Pl@TG8>+qa4Hn!!o
z*l+VC{FLh6`xQl-MbTW8gu&{7k*CLj6A!Jg)%ajn4bye|3;>6KgdkaHB1*7sP(qH^
zM3}jlQ55EbU+h78)O~_vOlk(NXmT-koA2ybw#Ap(?bGTfdr|eNkF%9dv<AgmKMoL&
zx;iG5(~Y)WkR=D5^%1!j-<I$q5>ne-Lc-#D7t%4$)kI3MTsHZXMEwAhN<T~H*EhQx
z*5O{J`UJ^sn<?1NGt&!UiD@}De8{)Dn}#N?N5kBW9K2+=$htp2|6qE`hUcP%RX)x1
z>~|uf@`)G0CFqIc6xvp9J$q}-v(KSnvKJ>@5;JB*U@U=GBV!YAB@L#MP4$QlFx$j)
z^sjb(F~z%O?)#^Hu_AEld(A2ej&fJlpV9_BJXlg%Se>0S-7d+KohI&an{4(ddV$wO
z*Dh|gZxhB?47DH=Iq&ASc5inZrvRcy5GbNJWkN!e>_WQ#gUof>*lOrKa8ayq@v<o#
zvr0<LBJDi((fJ+2n(?)jbb7a+da$vbO;!R{Nq~v6OY1rQcAo5OzjI=WhId-8$jkZM
z0(1z&Du{(}#URhNV1f(ju~~f3alH#|$*}00h3T7lYN#S(_ws^1|GLiZ91WWN{<$J_
zD)crknWd)z=0!JSydDL6-mYoM8Mj~CY_rDi%Gc~+o&w5xbC;!+_~g|1^>UfM0x}!<
z*X~%@dF_uFY={`G2=K(zp8}BRIg%iIg4dTb3g0HVgIlV+zcx>J3yu7s*UUELO$V*A
zE9|Bz+&a$({2RBQu+JT!xYdJbNEfZ_NsxV6opUU=8}_s?U=wAXx_gUNuq;O}NKf8?
z^p9|x?CPczgi-`>oq99#WvLm&?fV`GlA=zh<IkhsM3w4>5YqyOl2kw~uWK8=O*!7e
zFA3c;`yTmiGHqD4J%`a~ZiikVEzQ_m+aqrIzg@gc=>5EqKYnc?374}LtYExf_U@|O
zolwc6uF0@(&$eZg9bTTw9Y~A)(W|RA%|}0Q*~dGrosBscGXtW!n1&bG_-tMhAO;%k
zCHUsGbNfls>j$S0_FksBy9UbVNtMD$$Qpef)3m+5AE&$Wq~CN|Qm2{$FJ414O8&f8
zY4`Aw*ZL6cuHPnE|DQG~&{k`!?*D6#)|a;n{b+Q;?o7=J&)brga;9cV&Cxr==R<N`
zb13qluxlO|Mo1~c)6{5vs=n0G<jrlr<`SCOvLW9t{q-Kc#3OtExx|u$f%wI1kdpEY
zJL|&=tI?XNsFQN~3G!LcL`J2JL(lPI-T81wf`ONrQE5k7B5jBNrP9`)Pd_7`0;|9J
z!F%pop07uu#zbIk=FdtSHLrN4Og7#jW+qQ{;uhD)$~>J0#M%gi_hbud&wZe9h>}q-
zNCfMZTr)N<ybaf<qPf~BtY;gbv=W!1uKE{+OP_5~yEvTGqc#N>Stojj7x36kYyoSN
z3SR4l2Md%9L3g^DH0Go=-_e243VgM;DS?LhPFOu4L>E8UWF7fB6~Rr-ZGcscm+2))
z!(1Rcyh5~ewHo(LZ0d4Dz%@i+0~?uVXIQ0S^Cm7U!1L$ARFK5k`6tv3EDeov<-!k^
zzzn#V7cF&FypD&(KeZbP&-T6iY0-mFcKOGV+D*llj11CHFxRqV-nBLeG4THO(cdb}
zr$7<(ZKcCHh^Z4PFlD2sWS}Xrcqc@*4RF9D*KDpwuF>aDbol!o@+57v|D$~0Is7^G
z;&_4Uxi*`OAN6Ppt67`4VOIpC19irU%~I;4DVsdr#u6hhh^9GRlW=qJCi@e~$Qu<?
z(bErpPx~==+pJ5vSHIJ#bliaOa&$6Vm5YuLDq{t&X~I?7JWFE%r+PPInT(0RGL4KM
zGF@d0Btcblo$y_6GBb7@I`Jg=d7VD2U7hW=nM>;KY5dMKaW%QF_{Q!=pL$<zLot%Q
zr)5qv&Se`_053=(TLoEI{aG2{S-iiaDoAqdNY9Ftb(^6L#SELXokz|=x`u*!BT+I`
zckb&#N<Z?`=3P*&q~8yTpZKy-)c+}VWptHkdEn)$(H-t(;Oq2R<>xvUAY<o99@~+2
zoycA0-C<#pD0}sFG@GZ7!nuM#j&08tNMDCiL(8^0a(JjnKzNC%n4|H*w7wMm?@GJ1
zPNDHZtPr1!_cp=N??RsB0m;{QzuLi?s3@xpnFTey{W2Q}+tb|jY31<wqy~g_)^ga^
z*dF++hF*ItLOo8#%C#|ubm3?rr>P6#HTs*L7Q`u70x+d52-m+~?_+cZ$1v(fpr5~j
z^gk8*-n)PM%+Uw`{_FqR5<{f`HK96a^R~-#UorcSTqc#$7ypsZReNhqjs9#sV^y}#
zXw~SGqf#0R6O-ww7m>CY*{k3X>hrH8AQsby!?vR-=YRIY0RGrQ%H5(YsrEG~=ay#d
z?zyTr7?&#;d%Bf(qbMb3wN?XIml>!!X7)PzWz%5&u6~V0&E5U~FA_-50osFvpjmb{
zNYO8E%#Y4w#D;1Ri292(&`t(4unx)ICyL1pPYAatj=cPto&{TmX0PCfbob5ecYT(8
zN(mwT%ir7fNdxF_(s1(Zm+rtapoVX<$mcqi)eC+V7bMIN&_0upG54;|Z#KlRlncPs
z<L%Vqpd!<YNMKyUn8uu3L2)XHK1e@Zj)X7V5h<-v+hFi(=(0F~$DQ@wH@zC*;)BY*
z)xwM&eH$X%9Km(}dY^u|9J%R>MTr)7>72XlBlRfRY_Qh((se%*O@%4z1%!ZRXa7;c
zgnF|Zo8Ccj04|7udjaXCdf>~^q-|+05z)4CQH8@ieOkByw&O}40x>>u6{#A{*P>{(
zgPQf$=pxQZ50OES1ka^HKOvwE5)K&ACRfZyqgl19&UBGr^%PH8HuBv|5{G%Zif^)S
ze~gIs;|6js_fDPeYD(CYd-n!@Hn<5Tq$D|JTRBv5RBwWXI|xZ=^vyXa962heG`1wk
zQ!;wro{-}$e17GCXSq%qFW6>gQ%3`l|Jx{#`_LuMnRHpkwq#U6eyg#6?-N1+_6Byn
z;qKzr1;z2LUH}~E3)A|wL{8iGqxwpR%SaNA;690+AL8DOMWl}Pjw}I{+H!8Cg%#JA
zqmq`NMf073Bo&Rfn$ns-gibM#g5E)Cdp_uSddr<qHQk3#PC*P>yuyZay8%aicY`4_
zcqM*zH*iN?uD`3dxVvkFUHaO(4o6-VG!*5Nn4P|!`->}p{`e;gMXzmS@tERG#E#G_
zhd)$4mbyri@|s(}LIUCIXBQHVrB;7#2s9Db@ALKU7_kVs|1m!t5a99r{?A-tS81I<
zeab5A>a?Yj5{(b>)rQE@8$ZrGzj5MI=l5i0SFN1)+|3Sdx~g-rBNc)=Z8wxzNEOSa
z{4R%X;V7oSK&;yDvV@@9;L^`7QL@~jo5CG;2PIAj!fidB1nkIEh&*Q>;1GK1z->yb
zmn~9GsGc-Y<ygsGEHlUZJky%c@$GGKf~h3Udu@GDGNaDA^>vbsHeDu)apCOl@@?|F
zL76{pB?luG5~YmlTWi2`jb)rdQB$g1$Ko8(*|5QFP2$u_xf|NE?OTmr4=o_ladLQY
z_1hxeV@4hb$+?sCPwiFq14okvg=zU>p8J0Mf2R$ocMqTijNWSR>Vq7bk43kCWO1e>
zsR1sH0nGc}x}-!oVRmWX?t^ACJ0n=Rc=Rpn0rviW>v!`*ypJd<pwHpRiG+}p?i|j`
zcG<BRpfFI|+z2Kvuh$XS7?DBi%e$kic!!m_YXv#+$Ca*8EegK&m-j1vI^(ALxH=q~
zgMh0If8}Zt@%f_v;nMZp8ap3Z0w^UYL#_|!J~A42TkbX1O{x+QI92;Sh10vsg{&9d
zUHI*bGwSJxdeFDwls~2rI`KiadruL($*V^BtC@|QFiCQ@hfm6!?I!o}8SGKh9ol2r
zZ!v*h{uE@Va?uLhhD+665TBX_6F3<8{87Y>=MVZk+xWq3HH#W?YvSAUTYtz$#FN~C
zdaj%VbaA1J;b3W0{#Qm0YzgXZ6<~j=@E_HJQt!r^AzCvZg|9v&x!9Ui$Xj0@aX{Hp
z+Mg3+k6ICc%spx!B9&L#4-u&m207lMiVy%oA55Uncks*qLLrjFo~U(M#)LUvAi2Gk
zuq+%Bh)>K;kz2mE?_~&We)4WXGU<&=a?OI~A2o{>IA6Xl+${?rZf>#bw3BA<w)UQk
zG#be4l99aTcv8u0cK1bO4pJ(@KSHC3Cr<WD(Z=i+8>`YH-H=E_0f|BNxle!|EmAkJ
zm_2wddZb!#U|WUL%iZ-!ZL@OtRjlNT3LLqN_DHT6Igy$Lo?0wr>)6Vn7|jLdX2iLZ
z_3Py+fm*~RAlzy4;h^ti?IzhBeBHh`aH?7x&VUqNEx8>v9`=9CkI$3Lgn>~R`7%MV
z$8KivBny&G`y<+Ne!MNudwg{hAfkZaX$tM!cE;b29n7`R&)=XTWAxq6@_i>?wZ!`O
z6<Lp?O6Q+HFsa`qLgHmjo5a`jZO+Gk-hVXJKzD-!$IJa)Fy-;?!*Jj7z?D#lpl)<&
zc-)vyQgF)W_Z=XSx%g?PdHtK4;6*f0E(xAZG8&lCB&|7d7>x`fL%Sp82E$PAM7VeK
zjmI7fM`&vcvzGX_0o^i`m1Ebo*b60?+0zW=(qzfCj=NC^$_<1+xWL2siG5=1@q2>9
zO@ZP^Gl1|Oky6)PLBG@jW3H>2^$~*f`c9S;#s>{3#>E#h%sMDzuZtfzQ6ula{G2F!
z@T9zb3h{JGnlD6cqMjt7wo63GI#cc@%l22o)@eNHiS<*uUDY%J^R&&z9E!V;Sf7uY
z$`F$5+2@7@#oPuYl<w(kFI9G>-|0F|9dGlAt6#+B5_V|5e&ld0a~&Ogjfem&lwYS!
zSv78MLu9S+7S=I)cl^vr#=2?JM@B+DNqdjcgxE~=n@ibhHAAUE<f=uW;G2EO_!?qj
z+qWBpnd8;;mTs)I<2_rWSl9k)Ph-y%j`YhI;oCPaO@GKbldFur3R%k95r^&lUavPq
zR^>Sg^|&p3CRA@~>6--`99k?r@>HmXqd@0aj-xs0XI~m){N$uke2|GqsXX`%boGoN
z0=@R+VlH5O8wBOG?pE7X+ibKqS^1o@vWvrAPAmTnkC~tM)<+utfXk)EJwenj(vYRF
zXc$o_q}82R<eNeOExA3A%F$n#MFHmm)FqWps|FTop>f^xDiFt(ZT>5cjrS4NK^(@-
z4D*_Q{pq&fRSd3zS+^e1w#LD^Keee8rK3L<<ip(>=q_#LOO;VO<lf5B6700&dC6c=
zt;2)M6nGniAQ58$NlI?{yQ!5E)vp}oXLM~ODguBG^HU_F7mjWh8dCJ&b7AztJ*gZ}
zT~tvufIU)F&<Z){jUoW?t4Kpo@b5l+BIMU5dNRWhXd?b(=x2%i4~zSKaE?$$A6#KL
zHxKB8Ln|14@Ib9i5Ic;E0!%WI_1HSy7N@1=h<m7|dD~oUy9+(|>|L*USt09orbZQp
z7k@Uiu$U!rh2Uv)o0+F37OhL49@>=7q19H#JgS#mYA|2&xuc@XCENG%TKv;T@RrNB
ziz%(Lwdb+2<Euucrt=Yk_D@Gp-|Lq~=QxtzDLxKakx!<Ja3yzLo=Gx~ZSZxd?)=fA
z=gtGp>VwndUE)W!jYD-{ceE&5wUl0Q)AirdM9ZYS827<X$^ZkZdxRW30QYlGW6jj_
z7dPD4VB3zvZIlsu->rU6VZKvVYen9+Y9J8nyXGE$$OLv2l=}6S=wv=SaY|w$Jp-9w
zkvX?xvJTrM2f#pe_a$P#T*v6rJ#AaDrWpRp#A0qf;jsr+WRW;!_p^H2l1F$T>(i7Y
z5tREcBDX_8Z}?vu!H%z2SR{SRL^!JzUje8H`F#qPNxi@Egx5E_<4UMFwAMdX=yr-^
z-|o)%A=N~3IKc8%2<PiwO>`_T8Mc#B_48!5E8MBn^HbQl$$);fJIxH5BF>()IWw~V
zvQ1EvFli1B$n#8PH-*!w+PjgT<?SU0r*o>ORBK{DSK(dnJ?y*)>rWy0TEhgR<f1+x
zW~ZMKbY8oDUB-&kb@>gO(D+{seg&YxUn^4Od*&m`r>NR2Y9;h##1q$mB*Y4HrEbZ1
z!GIp7_noo&cYuJ9g*)Q4V=y`6+}Ao2Wg?93@R42%zp|9NJu*?I)3&&Ar+8WX*<h)z
zup7|pmyBlh0Lk$W{Zq4lG&j#i^esYku&=2{-RbvT3(>M19p`9j&x-DK42yXFy%90m
z9kYN{{cRxNQ;s^lzyVb3ycW3Ca(N)msVtsl8C3b8u3E;mK*`fac*U;%mG2<5l9!Pl
zr1<*gb&*KnHW6YwVQp*J&AF}<pX2w%o-R9xnlMpK>;r^!b*l$>dyrqCY|&vy|Ezn|
z*EXv;0-<eB)Y{99;ali%!ircOs|#YArGc^V?Ju%RxgVC3sP}&-UaF}tqCZ_}^8Ru?
zNN5`;TpH<kB1hp>f^hKwHLqJ`JD+L7!@AbHXWXR%COYZlclzyyy_YQY=?TUINRiTI
z?9%rc40(Y)_32yB<igNtWqFQHU~JrRTrj4<Wb1*~pPAVAA;9r|3;gib!<*pqgyzVj
z60LV%w%zq~-gsBnvGvEAUBxkRT<6Zo$L-2`qqpUz4RYKN^P;G}n_6-oX7#dG*&8>Y
zsD^svQ^bij8{FxS4mm6C3o`C}!q>E2HZWPSs`!`6#`|ZS36Kj%3B-$=dle~zPTiLK
zv|}g7xcFrFec=bCmD|=-K4?X*kGjPB+>tATZAr0(;cNZw(F*3=`ttK+#+JXTT6E>C
zdT(%3xx-f;ZOse<`%;G$*q3$Qe5^0Zdjv*mZnG;kl|HzKBae;<P4u;B#;TO_gRU}^
zOZ(rg(Np<iG)8HJUHO}Ne3RX09wQ`GXk3yi_K98XA9u?<$LB4-C$*A<Tdqf8LFcxD
zerxE>@DH(f8{bzfNNLj%sCJ3*KFKSu^XU@Y#x))Pon}EEDfW?j?auq;#3~@IE?JX2
zRuzSB4cZEwT&p<D16<DyiVhS=0Evdl^R&0+EOGW;Q46DPIP#z3MUEzJZngPExqBvS
zL|{kNF1pxsRHVjF=o}`#WM@qzf`o;ZMwU#kCKEtuZeM_YAHwC9K0yz^>sbXFE|_ne
zGgXfxcKZEp<{tOG-m++`5Q~ilB?>4i;oUzE1w7~;$)RkZt||YCU(2Z5zkg%rCC%xT
zo`8oD?UpjE@}$9m!yU>oH0}bcv6`3ky*VXgO6ToGzAc_k-0DQnnC{nf#rskT7$3If
z_KD9M3I{WnnNsDygbfoTlb~;fRU?{JsfIA+6=`ONd>slg`?<V_lNu8@dWY)^!E`wl
zal$Mgd?*0#*)znbwD8ja3t@F`Yr!IQN?2hE!*;)o@xx;9E7-cUH#o&G-<nz0*?VEm
zQYL%`5ms+qVtjiKe*fftDzv$!ZqMhFqx+`M2#e2{{PFLwF>$T^l06uZR^$yXSSc5l
zv%<KQR!&qgloj$`yE!af{tJ~i&#sZp&|mc`F?{_zFCwd|DqX?e@Ltb8l#1a%(DQ~H
z*hR)A`PmIB7xPcWc~bBr=LMA7Sd-4cg)pEhw{~f%*qw+q`iJV`lRC5|E>G>w+EJws
zV3&E|lO`{u(xQs97@ySb=O;Csr-A+&l2Uu@|3#cm?RJg_o$9K*3izpku<_a>)Sk>n
zl+(nblgPQG>!)kB>wt(&Z*OzLR6I4YYwPCL&1vxr#z9<sw6PtCE{_6<Z-3nMI7<wc
z3oDy%)LJEfEO+uz1k;}+qV2>goV`4MT-sDNG5~le0AoJU=T*_$_}M~=ux9NDRSpng
z;2)QsQ#JyF)##gsaDbaMv_Bnx6D(fDw`01mXZTybUA1i2(x&~{@Qlxn;bI|=GW77Z
zS6y|lqe|=aMCSY<*+@Vxjml*VEGT68s(5ovkXX}?P{*Nja)9QI(#=<!y<1ilZ0Q5r
z5&PP0*70B3ZOp>+alWDg0Fm^`mWy%Khcl8x6slRHpwR)9FDVCT2PvB8-$}b`FUB(S
z2v4k5q2N4QT$?5N5rC#>Y{)bkFC|*~HC5lM?tdgEi-@rD<@-UOnyr$;n5rLWTWnna
z9a;Q=V#Iw96ymq@>|)wnqh-|mz)`8#_TDYH^saAwEIv<)|1OUm!?ggjI6R%12qj04
zv;h>tmTNqQyB-LBN76eG(`=!&dNp@NN}mpRE){XezV>qyYzewcC~mE(mV;m&Yiy0s
zmBFL;sD$G3XqPd^8$ptg(A#-FbJBs4!DXBKU}r^`9a#|YSPOb*M)M|+hKwQpw7Vs*
zr|$_?m0DOyykX3ZUypeE=>MV%e-`p@x^VEvjIf=XND-FG2#5S=xs!lfyXzchk`~uE
z5G8;0%J!FwjHnGWwonh`9R?f|DFsG6OkI-0yACdfmA==iJ>#wh^}qv6#nfd8KEjm>
z-}vG9h0qBkKbEEZycOvDyG~;#eh+p<p4qY{_pr?Hp^5)0(h!GsN}RZT%UiqQ>>Bs5
zt-$nUnHo`HlPB}<x|Qwo3zo_`e#ylZE*;;W5Wfcj(jh(MLp*v|E%S!FZt3x;zW_4E
zJ4AUm*`1PdzV_;U=?3B`#`E(~i%n|h)!*!M6z#_6<z{((fT->$?bh--?96MbiVt%g
zm&pke;KQ%UzNQIiBgIxaI|566L<VanzD}PB<pd%LDWV#+n~CUh+eAkh5-wGGu@YeJ
zk>636hTx$E7FHmPrX)M1=!Sjt+LWB^kNel{cOw;!UHy@E#NFCr@Il5wAve85vvrXq
zW9gH+F!Y?Ztieheu>^+kg{>Xv#ouQW+A35fx}|-uacv$|Gjo{)O3BoPr;*<csM^}J
zzCw2(F+XK(1DmUP!z(XZ9{Dzz9H>ejT)UCCZPp{%kj=8i*fazneUx+Qz}vwlgA|XA
zSalb}oh+cb506Ai6};XGk);T{O&eQ(e~6_83tM@7g;Z#RM5{t;cv42z9t+a%W1}WB
ze_$3y|5O&A)yQ((thJc`a*IXcCAsq~A>TKQKLiV^jS*_Pnr6$&p#pv6*-h>>{cI(#
zFLFt1U!f7+RzyjEw1ndY_7f&4UIz20J(ushCHW+*G3pK!IsFiCGZbjhSvK)?o-1^?
z&6ce=?9(60K8rX_`ARvl`E459Yp2`g?gWt7H~E`1q++~Ya4-^qDerUCnmSQB*LCQY
zK-I&<{33aoSC7wI=rB+VvGB}ln`bo@(E#!`WrIjflyzh7FV#4!ZmF5O{MH~J%J}N+
zMayYz=jzu{GE?xwTM3BX)1f#2f+lz<DH?>{nI2fXQR0j4;A`O9K-Q}7U#m2L%Rlax
zm+!M7-MxMJ1}?Cf5vWMJ`)DwdJAWfz?AqV*4a6)U-yo0T!Z{hS>)9u`!$|4SzNqAS
zY<qSV>nqDFXzV6`o?yzIUB?X|fM)0>nh^{EBzm$RxH307SjujlVWI>~zhR8m>CDPj
z10qBsM?)z@F`lpNHrqYf&A&n*_}TgKQZWmtgiiy|o_OBTWZ*dX?&W4?OTNpM+huep
z7kefE(%AlM?b*TifjoGqlkbm?L-9y>s=$jve?ViDW-)&)wAo*dZLY@h597&H()v{-
zpm3*T8u~0onzv+4RGa#1^ia@dHz@oBHj1vLfcYmL1>vyShe(Eo{d!Oh>EUN9dR)$W
zJP;+nkQI1MtR#9{BH&K{f%s-uNjAx$*(tYX!tSndQP8^y5ENpv61pXc_IVZ{w?5<1
zH@NgT>!8{RVOrHIat|7~nPhN67b>w_`Nvx0<L0I8m&Dal3J&vx+{qN2;zA7afX!%)
zIa7v%Zgtx+34pKz_!eS<GZg4bh>jp<HQPrGF+ALq_s}QGUBj7$8P=Y|Z$R}R1~kz#
z`~^)EZ~qH4*~nF*D9Ncm^3PK6b0=K^?4&>=`28~lxjzdRzLP+I0HF>45t<wXpvhhC
z#s|M(NxU4UdmVp%Xew_d9`&&}tv-i%riRR-G@`RD)>r0iVoVh4^L!W(@8x4tI&Wz2
z;<eN1;!aNr^9XIe#vN<W-ibeMmPfzu77ku|C*jnt6v|5*wBtiqTr!tetjyUCN%!y@
zv3?}a?ce5^K=oUdyQFL0H!eL{wo$^Rzgxy^k>yxD>AS<IZ_|T$8&H`()-m<0t?^1l
z`F)zilWHjd1m6Pt=r2WymLy+9LrKX&&bxX!{^lxOv_)M2M6EDDR8<uKqEt;J1%5}=
zQ^S4Xwl_&%vdMk1l_c@`H<FwA;v1TP$z1Mel^yBu%D@ekHe{Z^Mts9tVcZ6G@aPLx
zNDm+@$8AZ=kKRuTC+CSLcNMof#QaGj7-BWg&)S(67TL&tj`^WExbU?U*q#cQY}s#8
zMdo8ortdVYW%JmI@!3M4c+a*pEx@AVf0B6VA&-qf_nci7b0-iLB7yfDtnL==-PR`%
zQ@R!_->*c`);yl&Cm@e6NS3r+Ip%38hdzAPm!#u_s~{}ad^R2&D%L92)xU-d;A>k6
zaEiY^^#?Sfc#CNYwYXZYQ8I{6;ym>(JE9qb9kEy`rQvfp$A2P%x#%e)Lfl2g&3orV
z!Of~+R85j>a{1*rXgrAhCpH0kt4p%pBcL>Y^I#%vQSbl4P9Qadn3cfy##V9!J#G=<
zjgN8d(Ao1iRzy(vqBx-J56jP$Ht!SB<=LEVQa!GvbT4-xtHV<o;UITo$q3~C87zX8
z?L8a%d72}<NE$4ElcolX2jcMEAXR2b@|19eC5!py<90^x+GfU6k(g+Ppxgi4mc{%L
zs>z4_Ws-CBbUW^r5QsD2{PQ!rZrISHq=|eoI*#l9;<?ZbIgDD2%b!bJ{pqRIW##-5
ztu%>u`B^6mmo{I}@RgkZ1)OC1Am(dq=F>N-i+UIEtx)@*h1{w6Z{&k0dBmVJxE?T3
zPZ`d7)=9@Ev{diErt+v`N-Hg$(X4jlB)Z3?ENVgZhn$9aM>x*2)njSx_Q^;<J<`WV
z$4zd+*iesjA=}^BFSzK3Qe2W&`v+>(MFs&VDPEd`37ou=r0RVZh?lPsdjb<8dKhZp
zBUaRIWICN_rrL$S5cm-fDBhb&q!$}g_Ill>FZs2{44pE>kw5JcJY2lKS2Od_HdgkX
z3$2MXsbt1GIS42KsRiYgGVxt5Mt;QWfH<`qKd})v>{r)oLbCpy%OUWzGy)DxMqDV_
zfe@`b@xU@=rEUXK5&hyyUrJ+ie5EO#w)i(Y89|Cy$g(^$8eua)=I%Z4OdHO-{mp*0
z(hWX;K!y15vi$Dpzjn^Amk$UJ-fuj<51f1viCoh0Y#(G38($!<wz`+dAGXZkT0t<S
z|6(Wwx&O{kY{8);|IAPh0SrZG{bQV5IUnR}i}vv#K^GSNHqV`B=06~)6YtAE53{)R
zHj;WW4^1d)b#4HYZT(hH&$VC{jz~D<CmequXB#u|p8(1^r(A`t{O^#1I{x-1rcEV4
z`;D>GA`i`T6ty9BGp=w{xlFxjV5AXrtoQ5!ulbA86Q!V$LDxsbvf29}8)NjZ6|X9n
ziSOKB@r8`jL-Doc<qiMfCiTAe{QTo(my@2K6h`Tp^lbDFvM@#ln6-JPQvoRIk!LFp
zv`FH73e65xVmF$_VJ^-;u{&=A){HxDYb@gNVu)k8uYmi*dVxEVunWRwVp}G45}}eN
zX@{@+?eKCTzvO|<cI%_!$UQmTn#p<l5K9N}VK(1JN<EO#G-$oM>o*wL7a&(0l~EVk
zo4j(5h%55RlzIFRCek?fxCkIAr>r`-32-2?Flnji!Zw4WJf_{_?VnTAs3?`c$4)(U
zP9ioaF=|DxIDF%1(KQ?7%BgcX0*JPw4!Djjpk!zQN(L9M@!`WfXU4AguK43BjB+0K
zB@-c9JL%)aqb^+R#ZS!I7%#I7pv^-r7`oeg3B>?B-EgF$wc=dQnq#fd=ZtXTvVCNQ
zt`g1U)!eTAr~C0Lv?*k&TsXyYf$R@n*6Z^HWFo^Q`ox<Pt1$;n@E&#2Ko`k7=5FP8
z8{x^|z)NtxH#DIWO8p;ip4O>oaCd`WU0t%w>AXsaITY2qdpb?ldpJmFw;t6>Tl~*l
z<-8*~Gb(%`J8G4kLwKgY5F)w#;km}#M!e-&m5IR4ms}&vJ12@dT-!D+bmV96P_}MY
z4Sjdq-u2~qk&9iU9(uy|Q%&+OhJRe1NosANwHsmK+VwhDDB${6HS;?dDsL|bY(5$!
zbePAXldg}qfvnFwbanQUtpWznj_qv77TUVS*i&8Bp<E1IN47Y6>J_=IB4xwA?{x3B
zxKxEqd^-S4Zx|Jam@d$-b-@?j%_e@wcD#_;lr(cRs(T{?Jd@rw&qk_7)#M@UuoyRk
ztoS%_+2tR;fL=y9t^u&_=02>O<d3){cu^hSL9_X6sPRM-{hNCO#v*f{rGT!&$eHe3
zSMfXj=~ISi-#X+1{%P@kYtyg}#M4(?d}hKjky#%b=8O#gaL9W{L`D_*%{~W{erWr{
z#&XWp030GS^~w8e4!O0kZYCynA&)*RvoH=o7LMetfxlr&+=zLIgdAx1FOqV(<|j#k
z{&ysWiAgVcH<k9M83rH~h9he~u|YHaOlYB|UnHc&M-3ao5!~jHXcoBdf&D@Hv($WB
zkfnu=#XHL4iYUJ^eOgbZFj&V2#!ndj^%L5C*h8Jn8@_yS26FRb$%58k<@>g`Pd7UA
zUTZ}!;++^{X!X}mGaHgFibQLcSg!!XJeS<!d`8r;aSWED))2v1nW`&jg}=snPU`~-
z7T}@>{Q2LEo3-_<k$wHBp7++z&{e^=F4s-~BjNt5ja@3RY`pB5{JuJS?ZCdYL^hc3
zLZ~kknDv(0I!1l_KeG}hCfO$Bs?nJejWxqpa|Y^#_Fn~HYH1-1rI0rm$W`4RZ*aD)
zYPrAdpSTE*j7MnDFD_!c^OKA89lz%Q7{T5mF3Z1OsI&HA!-nVW&^FEjCMsYS<78f1
z;pUawT%vQ=PbOhc{r1=R1Nd)J^h^ZeNSnm|a5%127`4CGyIlboWQv`rVl_}q-V3NI
zzIf}y7UJe!gOJ73tVVyPr8$p2x0U$XU5oi`Iyg>L3~=W(AQjo_`MhJ-cz5>gAeq3_
z84Ik=bJirWBaV|%tdO}536|5aOffms*>qXMfz3kQet?N^d<yIPY9XpZwUgD3YuDUm
zFcDj3oVx47S4y95ARob25AqznjvF>gE0+g!B$K@02U#A>Y78{;C=fs+HAu)Z^wByD
z`sL&rTSo&BZ@nfS<YQ3jucAhzs-`sj04T$}NxdzFAwxDz0I+uDilNz9$uTcW={mqV
z7`|^Ed}U!Z7JGz;d(%yip!0cgDNUkWHrama--a#V>b|(X%Yk%z2{5q*PgutI`HP(n
zfyHF|z1gx%=dw_~NE_1y#-C;+{%#$h4aoQwj)iM5^K1<$7B00A94qHuj@~!o!s29h
zKXC*<Smgs!psVZLAnJPl5rE8DX%Mc?`C1<u1x$zdBJWsdi&%|68S1g#*-7dI#%+%M
zT-WS>7nazo{1TSzdsUGC>QxcnY5avj?l3Tj!{K+|?;HQasZ#NMXKHwRgf14|vr~y|
z^S<;dr06JhP3q0loOPp2M^E6sUmudr@qb8r@1Ul(hwC3j<=Ws{>7XJ-q$9mqLApq<
zK@cJ(ASEC*2v@OCq)YF;34s8iiG>y*^Z=oE5&}{Z2qEyEaLc#+p7)(+=EwMlGtL~$
z*=O&Zz4lt4_3UEFDpouGmTQ|eRx#ZB@2pDrXGt}i=Xk%?8zzSN`w>6*&X;xg<|?x*
zeYiH*mVx%x{@TM&VBPk0mfC*z!-dIhPl3s%@r=P_81iYr-Pn7kscN=L!b%sbK7T7E
z0iCL^4rRyT4HVwGc*G+)OTM{lMls1Zx}PUc?#w@iHkR2*3wvM7lXrC!;JOWQ!w2S8
z1;h|`wWANLw{xSL^!m}PxVgp`hb^TkRV^?1hILT8yW40*Qu5TLqJT0DSbwjtNSb*Q
z^V*RoJvQy%<Mv;>BFs!32>e5R3aY-I#-`GilA85r=D{Sz=Pk+EBi-BJA~?IVAuVnH
zN0d%+$gQ_;d%Q`S-n@;47sE;vj$9%2%ikZ`?|wd+{0(Y%1`=nrv^-C_6s{lsA!x84
zpV=G!KSWTZ+e#tA`alLHC>6+{Sgnmnc>8+VFRf4(fqF>4I?=6`+pxPr2dcUtKLEg|
z7fbf+3U`|*q+Sd=F{wb<5kd#515bLi$@28<={l{qoiS%r0oRBV`iT4~O%d0j#Caah
zywFXrIb<fNR3+%&n_^h)#Gg=s!|VkJ#tDfJU{BrUnb9Jk)I-$v-?^N}LvEg{y|_Xl
z-98<m)}jaA9{fP1-#uOVr2mYVKr?6&6VyKu6R0&$?v_x;O7TZ$$5m?qBf-xvbUq4K
zc{UNt#{tN`HoO5)>dD)qFKkb|`zv7LL;(UORTpeijp_#P5QI$F-*=BRDM?<8d=@|D
z+?njQ7O!w&4X6rmf1WIa5n*~u7<On9YFV65VOnU8*^)Qs-n?ts!*+Lwi2-vV%>TXZ
z?Mb%u_aSQHhWc$(Rwna5#foMmV-ZyTEDjag6t#Ca@2cU-1*yC0L^?Z7bZQ2ZiJF~B
z*!can=og@IhD-Y7<Cn;FX_=Y9QI(w0D4;a{BMb#SzP^qjM+$9Ptv7&H1oxPbkeb^t
z^h4)wk5@4s;!U?zuP#!R15`B0v<D79kiIR)tby5Ve14SJ%e8nW_B0$9ovMstKk>2f
zbJ145d}Nzrkhpd-RO4gBl-BC=HrdIK40Ftj&d*H&g~9&Xg9SOD(b!(AE^6CbI$X%<
zbkIXzgKV$t<5~7WkG;cULf>w_gFlk??b_r7lCii8HvuxFiU}_nCvjgV^0}PrWK3vC
z>6z{L$-JVT^=EN9a6w^)Yuxv*3ceTzYDHer2^JPVc$*-Pa1@h>Q~cGV&`8vW<rjf)
zZLjy$=na+miT70sZV&0*F2l^$;n{-{7TwXCK!Jo*jqpd$(oVp*(dYA~1e6t=vzM0G
z{7G|3eFrw!JhUUbZXB35RwL{0-*OP*f+dVM<y)M}MK@5ue5gR;l>9b-W6zuZYJbt8
zqd=4Tk*`9tIhehOv-@q(w`-QLvIgz)3x*99#-*9wo3|0A^s8CCbXhQM6(I3X#fhk?
ziVS(CkNDqGBUB8*Xr7e4*I5huh5M?>vkQDyY&~|ywo-dD+mP75WCCIi_r-Eq7V>6W
zd>y;!59uzhL1d{D)U3IZn{&{h{;jY0Qzi#ejZ-moIBxDFjy=ff<yc<I(Pym|mj8WB
z2;!|XW)qi>YmQmp7%Sh$j*&{0*O||%BF0fwe>Pm;c?{!7da)?e9qdb8fgOG&^qc-V
z|2m))B1onn-#Ukjuq@I25GwJ#+Rpqww9^9c(fk$rV9*}?2+7j6%gF+RYkDWiJ;P;H
zkD1G4`mKI`|0AafE&gAr55o!*m`*sOIs$cT3vVii-fQ{JWiqtjmS_E(0psS(7Z9#8
z*D|?p4z=X+hA%RRCwz5;610AjY)=N&eC^Lm`2`GNA=P|+Yp*^zmk?nr)^FRq=~|w3
z*lVD3<q|QvMLsv+)i&Cngdb1eR-`?_O5R6LFdLu)`}1@4-xAin;le4GOMmONNots9
zKTqECd@b!lmu1D_#xHy{#NfQpL^`2f1VH9;&e68sUbL-u(}&RaVFO7(ueJ1MhRCP!
z8_7Vg;-6dZ2?4C<ZD8x&pLMA5cWC1pEwoWAFC7#3od_ga-g|OBuFLs`OKfRRD)dBZ
z&mLoX<p+cg*Qn8gLx{ogCHT6?FmuVmg|OJEIbeDMU*&&81rXVdqKlI%TVv&FuztLS
zfF)|jnHDq-y4_Dh|9PTxo*~6{0b8ldpSDuL-AwM^J4L@lxgIxiO#$eO2B~eo+=SL#
z;ntHhTpz`G{&O|F$m)j2cv-*H&4q6pxrAFgJuapmOA7cfuAeUAm%nF{y&ivLWV+-F
zx$e=j=kd}vnUuj2<*iuq(;~h8$Eym6&+G6c5mEZ5oWsdo6J0-4wOJz&l74Ip(f}+_
z`3+Eo!DDNuq%V4@h_*YX04T7D3cEi^eoUtMcxhzsl<Dc`Mff|vv8lcR3@t>YRfpG>
zuR!BD)Il~w5T<)2ZJ|%}R&g_fW{u+@l;j@smeK$g>-TByH{9k<LUqKR+<6WwsjZ=9
z=GdR_#BVB(Q@~b~Y7+lP=5t97v>96Jj~=AW&{_%VNuEmsIXv+v>72u++XDLgy421a
zb@W-Ow&_B-2w9z2sMJJV3@%uNtEIhF@$sWiyz5~{h5)p=<mSY3#rD(4`$yZ@<T%lR
zTa&VZIOPrxf?U~YmfXPu-AT_i^+?4RHt}H6ReN(JEdpVN|2qOvT3D9e@S?*|TTCJ<
ze74I3wL1T->07?=v0q-YlkgM@FV~LGXr5Hrt%CK5G=NV})iCaZy<}2dRLK~+Mflzb
zTdbaH^{^o5`lg&M)g-4C(z4$#A-?BTHDqGqMhvfBq#Y9`#K~|^qJRFg*-UK%&Jt<S
zh0I;U`n7ut5<p;yo~XJ$za@)c__zuBH5Ay$#dmRMEBv=CTvNTW*5%6SVc@NixFFc!
z(ZCehTIZy~O>P2{2xUNY!F_DY=2GmWZ0a!Fjkmm9jKwF)7Cs2|p{^M%VjBX6fOcXg
z+#B7~PKnjPjSiWa%@L0BnG`~nLe5M>f7OAjNt*s^ThY1H?iIfpZ$P~|8pOYd1<hvY
z`LH0Ivf>{2^b$2<an_B0AgOJt7v`qKd}?Y_$@pmQj%{;=Db|lMdIM;Jkw4J{ZJVAS
zdLHX?LRYPKbZCKta_n!i^#Gb+cf~yO7~S&n{T-nUCffVJS98+jNghr7;lt##kU$*9
zpEb-P?B|4b1%cFL;$7|#i_uJWuB=orUoz3nj&r>TCMeT=Wlg1~o!owXw2PUWOt*dx
z{Sdc0Jz?H&IrvzS&1u^7XXYzUf46qc4pr~YgaSpdqswd{&{B}333c*2PWMakzhVZJ
z0n~3Fk%{jX5J*0~`%Mk@|IRH?^5|})F|2anki<$xy03M>(fQxMoT=f!`sSZ|z98e*
z|NPbJ_b}i|p#A*6F=jSouhkbkdGvlsHFuj5@0dgah=lWbHkTgR$TZgg@G&MJVL*9m
z<v~sZekGl(!rH`vfMkP)A|*T6dOlH@sU!T}BuCm`7=!T_^G%GJC<OEN7nMBn2Wfr0
z7kG!s1{P~~<vjr)rO`L&mTf1&gnPygW`8&dPr)bP(+)QfD@v2VVR;99z&AttWZ5LP
z-Wyb)B-uOFQv3t10=$T*PHO#d<o;k=!rsS!f!3s~y}ANQgdaQ|%DkjxvGyo7mXIqy
zxq}z$)~`|;crZbI%JKYvuG^6DrFD6$=3fV_ErsNtt^B2TXq!Ssd+2?i;kvr%ipF}@
z_vF26S$7~FORY5=Nvx>|z_VV2)h!Fj=_y#RmKQ?7L@hHWI^=`FYbToe8J09_6w0iy
zk6(udl;Gy{R_#f5F!NFoKVCF|3-P5wSKD5S0x$Rj+pZI%SoPqjyaX41QFlO;7{>3C
zcV2ZbHCUn=f0a0P%6PxUZ`-X2KT7r(@9xB&VzMb8Ynel6aw7f3HlgE@?0a~yNTie0
zsJ+Tt>FU#4sjFvC*$jR6Y4{+GyNEt`RSise#ys{b%U@;^Qfb5lk#H+;w%UbGk1G(g
z?;Ozhy>ea(aWcwZJ6f@b@yF@W<%#$Kw702G>W3H!AqI?pe-nGyPx;>0Z!3G27N54o
zD+Y-GJhtgrT)A4g`^%?mRVRzFo1_^0SKx@f&<n`>%w!RI%yYTCfPM{#@!Z6C8skLu
z1bmQazY|Dxgpm&t9hXwM?gsQDYO|z?x!8^!VV1%-kvx<sxlK7&*G;RDr%mND1<$vY
zZnuevMbrV^xSp+~)uHP&&w1eQ5XVcE2sxh!xU)iylS6jT-fE<ADUHAIAsF?6`9GR8
zy3O&|l#D}lwt3kToiC$lyIa^@^Z&*9ex`i!e?%n!Ohi>LPg=z55t01a1JXG=uS{EZ
z$RJp^;y2*Pf4719`ezNt_%h?On^^*6*QyIV9)N8jLtaLNGx00}P{7W*ZXR3xRvDTq
z>#(J2!-Eg_=fwXhCYk)jhX!mY0Q_*dG2@y5?_yrFCdO5X#Z(ZJtRQa5Fxp-Lko3E^
z)BJm|$<53<2vhD~+EwAEpiP!`?l~0ZO@q-a)#`*hfPawmLw7!?)@>;|kQ0L>FKorw
zVe2gQ*`Ww9tiiGIB+pimTI~hf1nZwis0U7s6u?3d|2p;2g_r7u3`nZYwbtCi?231c
z#k2LQIAa+0PVbNDdykfP0O;J6U!P<w5)(?r-L)38@9-Kj=%3RLFPThta~^K)RxtUq
zi5%9Z_lH6%`i3POd@w8VkB49g&*JQqe=*-)uzlf(IcJ~*r;men&CQ?E<Lu54rgv!E
z)J?`zOGaKR-5R%z$qo&ElFn_JKwuKGqCoSSx2_Er-z9u6xXsVR1~dt;d$m0<y6VE@
z6*8k$VwtBhY_%0f((-`tH;zfR?Rx>*ZZrr@l5s<Jea-YTTcZUl{14Y=$fEbv@sT)x
z`_u=lpYQY+geuF8nUwo{)^=chS6GIjZLEoC!ib1@6L7S(gwHyg0NB*O^D`yu6Tez2
z>0tJf*NhS8Wje^REGsd1*%jM!w=oqU%5)RRn3irW@SX2<QOG6A1GHer*%Hi?1h38e
z{%fYu9tu90_;03x5yjX+xuR1%mYaBwNo0!c`ks@2{AOKRPy2P2mB}){adVNU-j#c9
z<7GcYI@AKWfwS*x<g$sb+Sc&qtreB$Y&~&1@D<I3QMti{P}lUYbudyi@T&hdId}{&
zMvl<%NvTdR+U#AGgBY|))c=C3ScM*K^PlMXjugd2&<(_q4S-#KK$KdMq~(QmW}h*`
zYbx-id*1yYp!I{=0zN^Vn^tj(Cn7w}B9kB7pzvn=<zUz><hNvF)%E@CQ(LB$S=x`k
zo8k}4%G`Py`z#uLcKjOj>)K~bYTZbOzY*+xboPE&=!ezCnHZx4!-NbA8MK15t)E6Z
zgO6>Yi(gM*^5-S1(gEAa&jwTyld3O5A#I#{RfrIB-Gco!cO<E+IMv_9&%S*Nk6qhN
zoO+9_=zZOI$vyG4J0(8yrSe7tX?x<DLeu`S2TId%d5*>Ch-aU}*YvXPhLiNV<Yxk%
z6K;cc{DRt`eL&t7a4if;I|1Yu*|B#O2PK-#^Dk7m7uNB;%CZE^5REy3lQqXe3XiB0
z{iEgAr@S_U0iIjW!pjQ$w~(=N)-Uz6-Wg?urCf)xGBQ!lwP~sLL1JT?b&hVmi@LeD
zA~WE#loC3WrrvuJuv5CPr+A*R$Z1^GDE;6;CfULcowv9JzRo?7?&`N9p%9AP{8gyH
zMTIB0WP5^i4oj`F_}AA{<J(3lg-TU=6P;S=Tl4iGA=!w8VE{WXwLEkx>n+#){mX<O
zHYFnR9-VJi=g%AZ&K>`_uq+aH;Ds>6bHA}E9qw14zEM|)oz(1q8`--3goX(O0hmDB
zBTV2!4>RY?vM9sta=L%%5O**yY%Qwn=dJ@06Y33#{sn7gfm;b|fvHd%o(fiM4On$n
z?*JL&5Bmpop$&MqgVybitF4sl<gz2>oA;Qvjv#@n<DdI=-szfI8ohp%0v0X5)z9&$
zwkGn&tB55|^NRGaCh-w2lQj4gUu(Aj+1n^IeBBt+qApgc>NYp9GCsx$Qf_J2W_eHb
z6>0VUS)Md}>p1`pFjg=Xgaf20p5_mLGpKGqDG3W|R0Y1qMaS}i29af;K_pz~$$uLt
z@-f-b+0~V3`<3>vrx_rXZQcm=e4=XPux7iI*V;P9QR&C3a_`E>Q*cAG&jkHGK9LEC
zqR2DZgUK$)i=kCY(RNQxh&SNMF^%=+l<6`47fTqOJ<GN{1k%<~2TQc5B3%&mKVMfB
zlQ8d#-n<oxvSjA!qQ-Old}heSJCu8{mm-r;W!bsv2~)-lu*v}k($#q3Hjnp6$sSeS
zJOa$uqJwJVnM^8#`voIghbs1}nerBGQR&*nc-3NV;f0F0c#+P=yNRifESmyw5rXWY
z2csWr^l(WXu-$&5K)Gm<qWqBLj(e}jkSbIduBl9!gVMYPYyoiw+tmx_HTM=+$uAh7
z7dZYb)wnJsQak&4z2&#x?08T%6)>^<Pb+5?0r!U6>-=qRvuvu1OlSczgGA#b%Df7A
z&Cl%agIxd1z@gfj&L!99e&ET^+hLrr(joG3-%EU==x#Y#Q)jZnB=^CL=J(zOder@s
z;R9R1bobt6zbRvYO?t<V!+n||w;4@>tYFE(QK=;C!++OT9yLB87q|tqw95;^7JY^x
znoa;xPzSfZ>4MgZ-xnkhbU<W?3k(kWv_Rjsf++Fto{mF;a+s<hc60pfr=tcnvr`8-
zEBzl<yXbBu=&$URpuYGk??(_M7AWI>ck`Ok>O6p#EZgc^qD7OXpo{#FMtmbBHkYHP
za!Sy$++s|93O_RV)LSE|->NL~)r!-y1Xa?z4Df!szj}}VcAtO|(cp3?zAJJeVwM_G
zsa6FMgg*y#MS8M69Wz{3!-3|-MaajOl)%@SYLY!aczJ~vcw8wP*bXFw6fTTxhh5uh
zYyab*uhNYBtmqSa@Qk@P1lk4=4^Bb`OZR3>EF~+;emwg;sAP#AY#JhI?&ypDIiO0E
zIpUFdB-y|bUk)znUrJRZx|;KH<s8{Zs-tN3k?J*_z&^)%w_yi|=UbqMB%8uO$R`>O
z^0C)SLdBSE@nHnY)9b(dBe;abfg0#vksB_x!ta2q>^IdshrFLv!S(xgl5s06xLVz+
z!Nx6)WU_a(GUj|syAB274v0m`mL^GyMuTTDPWnrj+=-1Pz{GbHkQo&HaAnc$js^V!
zNsBWXu<}{hvo6s_u*Yu8{DmV7`j@HX++U`W{ZI3nkw41IqO$W&Ej6d5Tj&;P6&>nC
z3>}V$=rx^&asl?eCSq6Wlv1&vf63w${GbSyZpB2}^b4f)x?8ao+SWfjEcwA7sqKJ<
ziWEU`3+0)pZS;8E1l+%KrV{f;5*29I>{|wQnWz69IRr}I@*?@o<aU<a4h^)b@oHRn
zC8EcOGuc)q*4<XUah0;x4-=?ba*HSVpil~SDrTo_a#dZ@eP2r`eB4{#DO7Nw!kua-
z4Y>N9_^`UQR(Y?HtA238=+@x^rO+6oI&`2k$qZzCps^EnEM-l<Rf4i~WI-o<f7sfY
z6Yo6{=Hh@jvWcy~&8N<yT*R)fpiYNXaf<x<du&GYcdFrO<>1eSdLbLn>`d<(Q*C0L
z6Z97`aQ8Ig0}av>0&P;oF#hQ9PFrU|dpj>qsqa}}|6EPu_j1cnb1wFEj27kE>44YA
zL(JXA>#|VS+)Ed<B32zBkMcKv@LO6uP8N0mhT>1-id-;YT<PLvE`n5tdMXLdQFNC3
zE?BiZo&xSHna_Xi7$+`KK{MG#ztDSAd#rx@-rX|P^D^B9Q`k_mZaIS{IW;%ng`8-(
z$cW1AjSgG=F^|-bj>$!ME%=hfgM^jzI*)oVIrjU>^0VCMmt@)1b3zm+aOG$p1xn+u
zr*^}JK*-CZe$rhrK_C>TJXCQ79rCkj>k!kX_dRfubF*NtfgF~BJ9oX9SpW_%ARmZ8
z^MOfo9+TV|l2Vq^WTLC%#`PTkI`3GsvtK*|QAr{oBxc`dTfe+CTz&_5;YkNxcziER
z<~)N0*hVwH8iAef{~oecjc^nV-eVqq?$a@>&>k3G8IXWvX=G1}|EMd4prwnhUZtgr
zW{x)-fMxJMhPCfK+0HtIo;UeeGVs|mHlIZ336n~8)`#D*RB`^O1ICM8ZGO6!lbW}5
zceuWjthF3fmx@N8j40Kb`VDwxod7yoU|_kIpPH7{BVL#omecPz4KxDNhApfMSKalU
zmJ+fo`O9I%Ww#gy%59{X1W`%YsnwJ>WHz_SXuNLux_Q9Ho8nbPX+6v*HL<5$#ZwiQ
zdHw-w$1HKQuPvi$ec=^%bbgyA4YLwI<ueu<ztul$LxlrI1L@pG3+|16OmFTUxXfbR
zbZy~f9sZl<ce$Rj^0J-ou%iKa!f8%ao22BFYSYVDqDqS6t)K6BHRdu}WUbv9OxE(5
z0v59l;YbI+9Mui`V{8;n$QfWQTgt$EEAc;sQ%(#Yj7nAM{haDq@G+#c=gox9y7AG+
z4jmItz~@^Ph2zZZM}1Vj1EW$p352iCYguuu+?WNd*MB#SFaj2b1ZJH$*HuH;E8Y1P
zHwPbG%Qu{uvS)6cM!fo%M!fMvG)j)v(R+4tfez3bX-F6HEQ;;*{3$qqbosJx|K$>=
z$je4jaNWCAIoZZU&kXi2Ia*AAkzAb3t*o>2P^h3-yN8+Z<iwFbOsVhxQ*3PhqtiE!
z%?<7!&_8DN^kz77(U!@;*BWGRMO-1$aO-P-06>@Zjt_}Ka)e~;6(_b;uK@hwO}{TO
z1TD9o`PqITa+H}6@xXFxfmHIUB`$HeICWb<vV31qgwT__X+AuDu$X<3r%vpZCPiWN
z-L>(LM1P^R`H!~%kW0q{YqNZu>f1ZwR`K@92Ty%Fk^rF+544yxDenY{5W7Nw7g$Tz
zT%dHu9A21GBSO8H;OCt{%5=(a9tOshSJ}r|hqAAn)u`RLEj5o-NiO0R)y14&ZzZNJ
z=rblDMH)s(w?T+6MLGj(MzDY$?CAQNR+Pwg(3}0*el)?(YYiFWbXwKR6=RP3%8Sz>
zC`(nl4#^OmxV2n-26td`f+|!um2?rSB?&vPVI&Vl&mMGl*kgiZg!^PXv9X-q=+B;e
zAkFyjlj71Ms-o~<dgGHNdQU-gLRZO;ome8FFaQjG^#LtT<$f$DZlrVh+5|uFeRkKQ
zNFqSP@&pYcV>sXR&^WNL%yzM_&$%na;DqU>>eqiUWJX1W_$g(leEkb=oq))+QEsBE
zmumC2qKkC&mLDI^4*U>8dv2%lhL^6T@I!JBd}~K4Sede*$oNEd!*pTj9d=0Kt<8+G
zZuvOh#qqUQ{ha*uHqAbCH4#3Xlq#??AqOx_$AP*5lkStifH9CMh7B_(-dw(c{1VCN
zlAWJDsAD_Kb1FI_&Gl$#jGESVf^>kPK`Ci+noF^eb(#w;B6NBOQslQaE!HM1#~2c*
zdpLj`!wPNVC;Z%d$BcY?r~F`43?|t)2qm+};FVCt93q~UuN)NH;uK#^gkA=|q}Kd0
zH|TXuZ)S~2D0oab)Jumerj$4Hq#~G>N)+~dRz1_rs$Fqzmb)D0B7UsuK<5aj;3UMl
z_ElnP-P#~2J1gcjofQzGroCSni1Xwr`gI+IdCUaj>Hw&A>}X;dbpuKKyN^Ujhg}U2
zGK!h#Jt_)w6yCO?B<?tU7h--naRi)gRvc6+xXGL~zSjTCWXoDY<k#!KEkjz2t03l3
zaB$;zbGqxp-wk&l^&8}l;3sPMZQGcC!`9ThY1~(s3bK?>2>_(ubAa?)!CQ{jVpnrd
zeWa;|>H8T`x?vv=yeQ}taW(}X<kqGzKq12wCmMms8Pgh**ZL3V05dCqM!S-<k8Oxh
z^Vq13ap=a^%<KV-MkVf))Geo%j@}|sH1k&F;M9MZj^xuwH8&mqXVZ}o)S17f{SOyc
z&4Cm}3UAqt%aZ%@2>;6=-NUZugPeeqHHl$Lns^cOsU;sK{h4iN@%t9{FtB}BCd}uS
zodpBWKYGJR2m-0z3MiQNnC$TrjU`&AuA(33$3r#HVC}@z>a{A%<T%A$1emCU4fPmK
zEeS6lq<mOS65pKh1;#G<p6J4J_RM5lW>?R|nwF8@h0b+-iWc6ddn;Hs{cACd^bKOr
zN8(#*lP&*Amz%?tG!QE~uB+t4P7IS!zyNcpW-=rSZ0vX8&H#Q7$~v*h@$-%M*yFHh
zRoTX61!*tle3-WSrgx76H-OJs0Ps0<be^7>eph=V9Y|scgU=3`zt2!rC+-;!a$0bq
z4|*r-t)DxTX1o~i=+&i0i-bgd+}vZV*#T;k3Zh(%{h#R!mV1Gn6t|!J>POKKj7{(<
zco0`kfMKI@f`7-o?=Kv`K0LbGb}$&m0%+Yi8VSWd3sk#s4JS8~pXzf{Vs_JP|3%bM
zT0st9qoeBvLSx5$pIAlnUY+3~UUex@SpI+vvHG<d`C&(lWwa@9<`IpgavSrypk(M8
z$O>-${=K-_W|A;dlx#^_-Iom-LA9U7G_Qo7?xhI04tz|hJ@`1<JpJlAKaLTZwU)sQ
zNYS(Ze`8~|YXA4x*h|%QRQ2g67~Rx`v^Q2MjR*Z0Kvl1m2M{mLK9A*lD=$(1%8lhb
z5ZQs6m?0W7LaE~2##M=CH462x-znKb#Xy@ny&v1x(We~q)!slb%#LDOu5A%{8^RB|
z2yI+--VxfKT_q+o0d?bFhgMWrd}s|jK-=`1v-00<(*RM2$V_Eq8l{y^tB+dYTK4}p
zak0=(w62fsd%Tx`<zO%@<CpsF^I)e6rbyG&C%+|`6d<jN`M#PA|L=jlzG${{KW#4E
zB!`E4ByXz^ON#q<XKOg&)}5AJEptkn9)f5+2&uqA1Q2|06nTc74V#wjo4M{+*{vF6
zF@OE#--)phxsadA0l(fcY2v*^V{S%nzUC>_zp{CK{$Cj1XL9Y!F{BA=q2bpXFP&=N
zGHe%n#DCTG5Zg3-A;h3cG=aY0Pff|yNo&g;BVrIBSIABY_6?gj*n7NbmRsE%=ejFf
zc{suDHB7s+bZNOq1YbiluM_{D-jLy$>y&7w*Gzt_778CG_V<rPEzD?9xQ^O$o+&YB
zC$GY>D`qE(*4{*ZS0q?B>1XrvRG1eqz>}j)R<rI9^c$GZ66x?UMrJv7asGklXmwvZ
zH7BA{v=OGJekx*3sn$8HC+b**e@Gh;--j;iC!A4DlE-vxU1J&5#wuwpNix>hh&Oes
zH2+nBya+thnoHA1>{Dr+nE&Mml2pH0=YaFT2T{}nf?|EDCn4X$S?a@zQ29E%FrGsl
zWwG|JJz2l;<@Pgr3z0s)W=Y8FZM`8+Ybh?&OKeO1oCqsF$2BOrNMRr*o^fsu)lE!*
zJkg@W0X*OC)OQf5=Og{Q57*0GPNW}T_n+DPzl~CbyKUEeMv4>pQj-d;(@_TswTAk{
zF6;XU5ksM1LueB5v2&I)HgG0gy4(rOb#HI-<vF0Sc+v1vj{%ajU)D|YFNvTY*CG^}
z2CMxgZa5Wn{)pZg$fakj)<=Tz8e=)aSfvpVZUzZjlGt+8zmkCXQJv}^fsjvJyce<7
ztlmkm9<Pbs`!ufaJP=@xe&b!gW`YUWNlqxQh>YE9>ZBHpt^~B_ANab}f>)|fTZZK6
zKrM%^$eCG{8zyR>=XzJT#sD1C#i}&<?r16wrpZ7#&4QzJmTpXxt66#z);(9$>b5Tw
zi;UAIe9)h89c%becS$z(1Eakm%}h0;tXdQU`j^gn-A_R7!-N(rPkyx{#5I1WX#SZ^
z&cGft3l3Vmy`mi#T933v<(oqNf7=%K?+W&rL9!^><YMgP$p;T@U#ch2wAK^9JjFKu
z(poz&gxj0RjrIK&{rnz4#XtgK^_%;$|2KSZ2+e<xbPLcyH^}!H6m;5i={qrw%4R!1
z9I$|b9Z9>7m&zdTofOC%IxX!tnSkyRLv;Q?-7gmmCmy*TXt<|&f+ql1E()mPT&(2F
zgd~`drra>p&=^o)E#VrM3og9>;KH*HSi*&tggH0r&}A_==#0O}V4x))=>co2KG0AO
z_;G2^FB;h%`3AUgqIrNn#P_#S#dm|Zl&&oo>Mw5JAot1xJtz#<Jc@=0&c(y|=YE=%
zUNTYWoc{L`lv`J4yrQ%zNBo?!?pw`oI_b^g4h#=o{a)-4k7*&urmwLl4t(nhRWQAp
z?fl2Le_?z{PcHwv$4ap2Umh!MRnxMFV@|sGJD_-LI2R|a4|wkr!&MK5UoHHlBE@$V
zd*BI}%MKAY##L+w$0=F7zp~5c@@0Q_@~8`?K)pCImVOY7ePzJk70BXj?ee94*%KRC
z+}q&{d$a=J!3RjzyS?~yrAL%mGfk}p)Pwn@+&-fs#x}nVYYr^Xi{Q0s);&O8JHja&
zL?Ut?D6Djgv#i83?0967Yw$o{+gue0>Rj6LI0O|)#U-lb`wd&0oe5KGLBu9Vv?0sz
z8nd;@AK8_P<gQ}`8+6ACojFk@I=e(o#m~^nc>%`a%@;4qJ!bSV0ecF(-9$`r<O{cs
zqcw2+$f$KchG}bu+W-;psrlheS`MtWEr;i@1qd8W%-rFAXKw>^qi<gC1q8+HMVg?f
zrn=0~3e^6p2r9R$=q|69L|AE5E1oE74ll}vx(tSd(m20V^`*c)9(z-vA6D^tZJ~Ez
zM=YHIrf2F+o#Z%WSveI$-<v3(tu)0a)V{-xPJacdo^1K%(74Yf2MN0>D!b2&hTq=-
zDp{EHR>;59i@Af@{mbXRFVF@-KkRDJusAdgSEiIe*Ug}=q(mMEc$Z#nAx(>yRQ6|>
zBym#`pPnrW(*VHT?ZY}n5T9$R)d3FIvc_EkK5@XwuFX_F3zQaYu|IFGiWh*<y8?t;
zS%1&=5K3L|qz#Z;P4jVU0IRDfCgk(;W5z-L#`_&1cI4v;wA+2~?M}54VZfx{54guH
zkdDH~`|m4+f-!Z!3h{!afUzv6@W7XJL0DwyXtAjOn+AMpsIHsa%!xa8^8G($FUaho
z;e7?xH+liqFNZdNFHpP2u>**EeF;1#;~hU3{Mq$F@BvuKM^(Jk>5Bg(0;l~&1l}{r
zoet3*yTOu}mw1)Q-<$l=Q(mXu4|L~LU=q#-D}{B3@SU1ts5pN^3Yqt9UjgmamN@2W
zxNFjw=lLZ!F)_5opj0>E!=uZqK%Z6WxAAAqWp;}Lw7XNiltuMd03gZ8#TJ6a<w+~C
z{^m5Yo{tC4<FZRWDn|tDlTb7ekc!(=x6~sx?YU-;j{{y6p(SmS)S|BejB6jTRw)As
z#Af||1DGMXj&uRy7zy~3=>lk6Em((F;?-=0BqtQ&VRiOsZ{Z;RL_A`pP}g_Ey*kLh
zs!;hiUq#u3DwPxVP7h6UK8~|H1jti6f#erq0E0C}V<fu}qx4q(KzHu7OjS3aJNI(s
zQFrc>&fx3H4j&!HZlgR`J-FwR=P!n7aph>@z4OM(M&q44{*w>PN3diQ_=HC!;ze;z
zm80rh^B3uhOJ)Cf^fedQMXq7;yZpxt(vPMmp{u3c0ADB2&}g0WCmEPIY|v$M;p+TH
zv%#e@Xd}>S(+HZ)BsE+ah=oJL0fQVk;N{wS>>#GJVc&xC?&^3b@cFH*{tXV?@`NGq
zUvS`j8XUM5fAuXIXrj4wht@<R63MU>ou}h@U#XWy2d;k((1EufaUHLBGlOY1KhIvf
zp81Fr6sx$eO{=1D>g{l-8oqNHP_INa+J`?2S>2zD@k&agH(l4s$E;isQ#>K{;@tkZ
zA+z9oCy7aZjHdLN{Iw)VlQevfA6roUna$~&DK?mq5OE0`QY8}UKL+Nj=Tj)tzQv1q
zZ&lyzKNImk5`qbA!BIncMLQjG8X+F>>n}&w;ds}xfOJLA`Ea0Z+5%G#ghO1Y7n#lj
zO}VN#=}x@G?jSV%Ku2AlN?heJODQcsMt68fTKj}HO(R@7#L9|Jn0FB<?@R}zqzt7D
zfj`7-$D$%ZguVF3k)+ZVp~ZXqPFJ_O{NxW#t6#et#BuqvS5NwGU>RPBl`Z=BZ+(FC
zC8MSTP^0K98MWZd%9lmcb}iOd{3Dl(z}G&1wIk2WZ+8FOBgB&6aob-_L~Jd(hNeBF
zMJHKSIJ$75MJmDiq!K>zCvrdO++1s24wdi8M2?bxqL4#f6<A(wyS7CNeZP`hNaaMj
zOw2*z70KVYcO(hg6Q}VHdN+=<HwtO@Ja_97*d|qEJC%&-UR)h_sHzhH4W-$fVe~!q
z!5J(sV4Wb!Uy9G9_U~o5PNp|@>=WdVtKd{EPo9+CJSpHc^(oG8Kn@$|QpD~)%6*xj
z+GKDOgm5^tOb=`eopgm7Mfy-GQ)}wdqKd8L`}z``%P#+vkzdDb-vU;+gE~=_b%WW}
z1Vj52<+y=_?PdPax93zPaNnkP@PTb>)l>e3rH`}eA7f*P$3@$`cUKdJ0{xbo>><l+
zk1li6R8<9kgp;#gxqW10>_j(WebLrAQ?-Az4Dl^!T4_y~&IuS<ijB!THR$(G2AbOb
zNP6AM225?bU$HB}QQwTlIp&XrusO*M-aj|I1?npnL|q|%9TCp726nUK=Tv#|ewr*Y
zJ*5CiIM;@(ourxc#4}(mF$!kHOw8*^y1~@AY-hX4djTQ}>Hz20-I-|D^gXPW`5=+V
zDhFBJRK&5yLWu}TtudUtdAn!o121z|$GMKwESyh^tjt!B$E9_uG*N+TJJ3n<TsQlx
zg;icBbtoP_xw>zra%#vByuY^&+?X!zR}GYlmi8rq{c28}XGwu_KGs8T76;0`7ckKC
zzedZiu!v|1_~7J!yl4*aw1Rt17&K#SR3!WBOLxwz`i5(=m|-6f<$n~MRW)CioOH0S
z=ImjqaLlNn?K9{U!Tr`b+ZA$pCO&ljN$sPjhmzNoEN`!HRS^4s|1K94Cwih5Awn-`
zfBeowLD;3ReMw2zr&roik@QC7uCuKfQ-~4JEsY|I*Z4b$LOSgCy8RJ8{O(ni5+uG1
zsCCtC8hOJW))lX2#iLQeo?Mbz%IOmz0ZD3K(!TNg${>`{aBQzlr)op-)2|Ex*<SZW
zSy9=gye|qE{7yp*9P+bSVFx?Wqp(6nrg__sc{#9_XQwU*+9Bjl#RXo6?i60MZJxAi
zJ*Ns}$JEh*m6!6hCAPj6oh^lf);nFlOMPAP6^{{}sth)?PlQODf64T3EPD{<HLQk_
z!rOY}OmmrQzT-7y9s<*AT0t5CM<3YY+q5&&=KksB-7Yz_174|Bp<58r<{f4tK4?PQ
zEz+e)`iQx|45y#-TC5A~F%{quJoSxuY)hk=?DH=9i{9r^ksC1Vu>_g!ZHcdY?B&O-
z4A6lLbbNRErJRIRq(oRz#o^v!!?Y{C;)Gd>xl4nX(PYiCOAOTQu;1ilRbtbbr9sXs
zo*4L^)vJt#K2d8{XEE1!b3?11W54<VuiUG%At1QwFZf(^MZ?G|23y)LD=1-e=L8QF
znZogD5p9*+snLA8_7U5_!YKGrxK_89fQoK$!q2pC&MQ-w&?EG_r-AmMhVYg$(+7U<
z-8xLHUNnRpJq=$`<g+1B^oke)D>u9Rsm(eq++qPsEdS*<{6sHn%v9s&Ii*ZW;enQL
z{Ahge!*LGT_?O2T?g03k$<Aw-lFY`=+UG*$5|b@Qw`zK!0VW1S!T$Q0JRThN_uuKn
zPXEKqp`)|D_)nV=AS&Ga#}5d6F7<0Ha$m?BI)(n`&)stgU8##QK=v!V$7xC06)tQU
z^T#*pw?<CuyD(JU=^%uaOuN)sI_ta3#m~;b0u;Ve`M#p=i@963)i1Bwg&tBftmky*
zv@I0po7U9Y5Zv2PzvpC)GV>)b*N>t|rRS!+mz}|DcjzV~JQJ2Cc{0LnDH%(A4H3))
ze%Qq^-%)FfkoDM;PnL7S=c$sS!t{F#ub#o6pl}Zl36<d<tK0~;=F`6_^O#fvotHg7
zu0$n8Y$W%qY1It0rXF;k+r#Te>iGLK=-ue_Dv5x(MZgt#Tkgt0Kk`i2nLl3lw2(#Y
zJRAG$0-Bn-#oHIhsJV0Ug$Gj4dE!K_`a{fH5L(vZ(7)sbtVusQm-7YPFTu9FZp6K)
z1WI3G>qi%Sc!P@uHYCb*1N8;Y@znQ+j66hn9JW=1@6+{V&CzBI!Io1H#ELRzTTrO1
zS-Gj~ROcz~bfC$Avaw}@$oVdH(dTig#8sl-WZz~gSW{|yX#J(@2AAaU_DJVKqQ~}|
zp6+K?0!Dep<7&E_mh&hWH8D#pM|P?Y5nOK2W;^bB{ObVUdDR|-G3RHumU9tFT&cT!
zPbBZEuIw5!f3UOUZOI798=Z-_yg2|4Hf>d%??XljMlomj$WfII77P<U0N%dYhJg?4
zNLv=uV-FmHNYpcJ_Q{-^_Y4<*0-r0hn&a2UvH_EjiF;%7Vy!sL6r|bN5+ju2TDQ>#
z-a9pcZ<xs;h#4LBeveFJ_Hb(~JLM;=e@gWfrq8%7A&#>C^-^$@;xN<gxXnaU2(c6`
z*HWzFS0N89znADvi32mS3>Nh%^)5%h8zs>-xqQd7NzL<-SY?SCzw)B>4xSID4SQ0{
zeQ=8Kiu=VBT==B={hh5eBW@g=Owx*S@!POV%ugje;1tsmP+V@owS8cG1{j5PTs?bb
zzNAzY7ay!D<jU4Avj-h>S+trSaW3k4inQEbZO;^0OKzf2lv*D9ecG##fATv%QHbU4
zkhb^;VS+Do!{y$0?(;|!dD?)*s<#5gZwK=}&zcm%>E%n7$V=Sm*wway<-LQJ_NR4$
z{{1A5!cpsuH}PJ+EMMuCmd~ape}8J*r*@N{H@5A-z0T9OUhifN4-1*+RQ=S6pnF<C
z_bdIKjkw`5)-#4=_lp}co9lmUU(kCga9kv@BwKUE>vJEns_Rbfd>~N#nhPVs{-}Z<
zS|tapfHYk|dy<?ZpJaF>L|;n{XfI}0`48<~5z-xSe=SXcM|6{GeE5~K7p`Xr97X`*
zMF#*yo=A=03*gpN`hHWst_I7Y*~k^eLMJ1{+q3F<w+noz)J1%%L|J^Q^nf>a*&vl%
zljW3!m|)zma+lV*$*@1Vk%SPR?3z-K)HZ)7&5vfT+jf)o%xLw1IhxinvW;jgajlQ1
z2zz^yO!Jnp>n~S4>ZBh_p=LOJ%0<hYWkkd$Qg1sv0$mNP1C*M2(7J>zu`o#&G>g?$
z&Mwxr_xZH68!Av!51qJ3xM!hnp?OUKf1le}p&qp3j31{pEfIM1>salUjAyDqSJR#1
zd*D0dR<?xQ8V0N8?@nN4Rrd5agE7t9kyRzgmqe!SfND_k_A1&pjroB-#^*32$T)eB
znq{H%EDco@f4Ap_)0hF<z@vzO5!;t<8vrk!C#PI+lohl~X-Z4sOJ~VVS9F7$I?xc}
zEWETlu#`79)?IjZsw%r$oips5&gBobo#moimmb{!&Gg_Dwx_}*BWw_onOhCL(BU4b
zz+W8cWCW)qpHu%GqZrkusS|W*?}MZ8iUk+eEU~vMS*?xopX&LNV%C<q-gE65<kUP!
zBVF!5kNZ5Km-2}cOJ89GjqesU7{5O@+?lfg-NhLvYWOCTu@SbYgtV8Uj7mbHg40*h
zKY|8RyCmftr$nnPIW{FOleg<wZ6uD>t4(^5E_eEOn!>!y<am0KwZ->N@RwFA*+ytG
zh81DoijymZdwf^s466sYe^W)IJ;?Ql^l*{yH=Bh^wcBnh2-!CeA*z^$&N1mNI^nlM
z>W-B}q420B-1TJMMGf<H_Wd1{fTU5n;4X=`i8?RY5FB-biOIXp5I;~XGD1ay&6n1U
z^SOCgIvkoUtT00QuF6a1>2B|h{g(UIUm+w_cp}gmf`nW{wXvwychRt;fl6Vv6@!0n
ziupnfoe2X2gEh+&F+|SdjY^9uNR(I%Z(;Rh#ReEi&u!)58hp?H0?r7@i%JOZxb@#5
z0Z-2_Ee9>x7f#ATTL8*HvOz$|e%flofZZLHN3(}^)A~D{<X7Le{NW^}*2Y4NPbLas
z8L&TX*p`AAxDcjg2KG&Mzd+`g0al||d;nIXXTamDA={nAg`1@4v!kOTP~T$Ps)Y0|
z@7I9`dRkzrP-82ot`b!IaBTN}@xDh4SY4z>sUWuk&YQX&^xf*J`HR}S?2F8Rma1bO
zhZ*%x6@4!Z&G5<>McL#ArkW+X7o6a`S?;i<RPfRLKBn5QzA{U*WV1Nl*A(h~?iM39
z#E|kNzIf^Ug#kHK0?&wZpg1!TsW63Su5teIRo#i?`uXGJtCn_$NJ#dnyABWex9FDN
zheRzZ#{0-(QKy-~P%arpgnW;-kF(qTnys|q^q-<FJ-Jq4fuTRojhD1+E|DI3D`iwM
zxGav!Om&pXV*V)gFR0yb*mi0zTiR9y4OaxXL|*RnkkW0N5Y6t=;*zio5X=@`>C>$Q
zI_;^O(;((k&7gHh{0B@Ru>yLEM_8G%P%Wuk?sRp|AFm*?;sW&Srp}BZ<wVt3fU1{j
z^;{*$6mpq^hqR(E8EE^kd<!Bpt2)JKbyM$&e5(Xn+Tr90u9-CD66UcEaHVhww}@T+
zJTT^~TUqn>cJzG~Rf;3zme|40APH;p^jL+dY+?*gE@!TG*-#j26&an%Gb>+}1h$1j
zOZx&smCCs6IP9EU!2xA$R#et(ku(?l9)olG(!B_u=F@bFqAX#9lGXEceH)0|TUp`$
zh!9{VMmTI2!7CIxEv4<0>rYn&p4XE^Jw*nLHeban=8Ic;|NME~wrK$Q`j+RpTH+L6
z9LVv|fB9_s#QP*23(w(&&~-d>sBDvMN=sRyT%R7y-db>TXz8SkJs+x<o(Rcc2+0L~
z>b9Pa71k6GKK)5Xh(YN~7cL|Qu|MJC_u{PT7_*bPK6|BBfHRIVU!F58T5b87tt^5n
zb!b)^nOiX~-_xxcL9N{5e4B^3X24+L=#|PbZw4zbR2+HDGO=SP1J;&#bjGz=+X9-Y
zvbCZw7oZ*)7+qQ%j#tc?&C66^I)Fa=MS(UWdg>Aj@*NHs)2WODJxQ;u;0@@Sz`2Xf
zq~4smDIb?+!#%%NV9Z|bO(Cl0!Y}whqWjVv=EBggicn$;{XKbV&_TnooE7FR=H*1h
zb6nooTRC72jo&gEHhm#%X*VRk<Uioh*h=E1M-bKPHYe7{w~<~u1!sqeuiB*x(PyD4
zMJcn>f@81in0gm3r=wuT`bw}8^r<Epf``&#7&j8X+ktSgH}q|?hL%vY>8AZ!taEg3
zi`@7rdN1A)Tzp8tJ+XJ!@V86$k$&na!I=~&Vns0K#yt9}<`;U{w!gPs+oyN9$pceT
z!mVsL)Qa{?bD3{>he?hMO)dG@CRhY8Ag%~QM=@rHSw@GsT~!_Tlg2-QvtxsqQmRPa
zkMmaFGO6kWI7=$wGZ821h`WBj1a*zgSLMM9J=)I@2U*uNw#&a<Ax()>I9<;}t~VK=
zZ(VYB058#+vtpwKa<-b>=LW&eO?JbkDpz$t{U*#4L;^4nl3UZp%qIhz*SuF#l27Z9
zz19<^dR}=>`kvy!FD1Lw>qcs_jeQ@WMt}n96P>>PHV40`!-ef<1hz;`NmxEsYNB+{
zZb>?sBrulJ8EIg-h9BZ{qGpN2;F?P%Hn~G2WpO4E#XI(KYFsaDhJj(23nOTra>mEF
zFZdOB3{<RF^q&c-Q4g^Cye_lUc&@$BG@j@mCJ`a%#$DgP?yI=EJ5z^9Q@U~8&k7Qd
ztZ&;yXYF|DBu}>}m$H7AcIe=*k9yS)99N!n!M8QAU4wpG(`Hn<xp}k33mn<sjRz$)
zIEpw6z4@(Jz+Z}+(Oa!OjSn-D;Z12@2H>X#C><NtFluqih)BXviuZzIa`NRX=Vo$*
zm5sGV&(~{F_#=e5pFlV33L!I>!Y<1&=YlQlkv*dk9v-#~?>JVQ=ST`og4ASxxa;Th
z7jMJ$w!DVQ;u-4B9bWELLf|R~*Tr^Fw<LAaJ6AM)w_}{?8F(a9p?XI`hVp9n!%V;C
zF0WR5(z<syPI>ArTY_w-^QXYJ@}>z@4&Yhe9EtZTKo*dO^17>TlfEybtK3?_XmWc(
zvAON#Dd|xjmKBv#;1`9qQpKo-KO#-ry0*Ig{#uJ}Pd!LYjtGd1Gh#z(8#RN{`F|b=
zl3AQM4EtZMk`TBZBLkGj$i{<1L|W-cbR@)#g0|t@1HNhYk_A}Ep-q*yrl(@Pni@M#
zsNIm4ORKZbb41B_*riK(d3H1L8Tsg;(?;wl9sw8@&ZZ1e`4ygu$)tGXlyrAv#u?l5
zUX5wZc3gW?p5iMRr^0RSHb7l>?3EOLR^LI15xx_%P7pa%C4&}ng(aEEuY;phUODGM
zpzM+?b%W@OVNU4hcU*T}l)+fE_Rzs|ZfIU9ahwYb>dl_om}PsJxcN!u&ail~N2L|x
z<U?}_G`tdA?@mGq{_8&5p5TF;|8rqkUl`XifH|AKlkarN;Fy|$Pvv7>8OojFcgWI%
zwK-Oz-?tUe)x^G5FmH+?BA3%OtlcSQUwqbN4h}AK7uZ~wpgql#lUKto@3&5$6NZ*V
z<bpXpdSrS=nx20}DCZ^?yR_l?)ijIGkMWt0Wk2m{dofoKjqn;?_w8a;@x(l8MwGQu
z-`Bd%r_vX7G1LO@*!kMPJ2r!Gu<yH98?*1x#I&BBcV+S>wx?PbLPt7lC^#9?D4NvE
znQCEF{kgd31)%L^Ev>#HhU9U#-10Me=l{#g!OuR05d@uNI>V%CCPGMLW>gX>Sd?QS
zg#?FA6YcvG>ps7eXiRr-#p#%E2S+fEEZbd>__VOm5UI+ofNU+I-fZbfRZRvni*KYl
zXtmrsw7G>%WIQvnfKoRH3c}u5&!ae#hp!aO;6~Ef1d3i`Rh4V8YwTuQcuAHR>ODQg
zX6_C0Q)$Zjc3lU~+aKuOgi#IATep^k{Uo#Z%T`~{aR$7v)%!5Zp1R|-7i*WNayX7j
z_wF1OE$)&H5X=|tvuj=9A4QJI@-u}2312F5s=|*T0}2%AUrxpmI($3g<`n!%D92`{
zSC(*=qZc(@zvY4TbCbhuY)5k=jYJl=*|51lQFsN3R#BZ#Ycp^3uR3UuDm!nP{VsSe
zKHlhytu*MeXhH957IK|;atZMx60q`?#YbUSQd*!^?aZLhft`*_CGz$ut|b)#9YK%@
z2@C9nc_~yABMMlJbQgT3j8x>z%xv4GM|Y`jdqNLp0;}8%)cq9j9!ew}ng!Yd!3+Lw
zYWZ{8HUs!lVEN(M^7z?p`LdD+Zk54Od*LQ)o^cUs-kvVn&GDZj9y^}7C@#N0toRrg
z%rmf<TJ!JfA?e$n>fza+lDgxYqYfh4GMmwpFhaPO)?xxu+RMit&FYrpR?WFyH5_gk
zvv^Mfo**loP<A&wvdp911+NWJtkX9DT2Gejxnd!+jlglKQ`U~T`3IXx4MA8PBHjR?
za<}!kIwhC%d>~E%<M27;#d_1j1lHp8?LM=%;6{)RmkBhJyp6XcLoMn3Nlqu<qAqHr
zAsNNTT)CAoLYDScOXiVo@Bbc$H!E%6r)ZMG7)#UupAXgat(3B4F<op^WV4<pMIZp2
zNbP#OE<Gy4vLXenQl)bme$#^$ljUggFxYvj1Z<8zKkD^ItCLam;o~Uw!gLN{h)b=s
z!`OW)<-CV-kuMj7<EI1ohVCj}+}2{&kU0gB4Zcx7pi!o5$i56LRBy!{#GdYrqzIHK
z=E!Kziz0@SU;iaj+MNGBI@-y6+`>FT^>XuZDL|$ioI!Ub8~k+3U!!|+fM?0SR3k((
z{t%68GMBg46|^R1yqw~C5e4=C$Za$M#i%u{kFiJl)E=vM_L-LwY<wi|1TF(d?&d+x
zOAB;|#M%RxZ;@*DXf9+=riS_9$b0EBG|#9@*7RMJ=~hOSDY;=J6>s-)vV*m7b}E%Z
za9K(PB;A$^ErT$cUODL=CxPV%(P;iuBwt4QQ?VSLTbw8sT^ZH&nypf|r;P-Z1dIfv
z_SCM0s67_`Sl?MjlL^ksd42&8<+Q1SD+2<+_*S@_$OeRIJP`Z~QKl`a3|Hz44A}ZI
zHozOZV6HsHMsZx1?pR0P?Ah&VxFe$J0b6S&|D$patsA6Y5%#I>m;_31#Q|0lp<SL4
zP&tX=9!*3RtPFm1Soi^gvvBV!X_b=P8fv>8W1CxjZisKFI$J5D(Eh<yqGG4)2W*e8
zMPxu5?<7>bwK~_PfHo0Y=UduafbR+dV)W<zgZ|miZzQs^E|oCdFq*ZAg??3LPE5BB
zJb8)i=Q|o|mzQvCc<RjgUs=8+1v;_Js4`Q!fsrD1ufDevqV0Yl6{%C)D10~p=89!S
zxjMx>38iQUVYUFSzb5*O#*s`aYUcjcT|h<DxoM<-rOe(&QR=(%b9O36Qj5I1q;T`5
zJ7ZA7TkD%&%~|%7PKv#DE^$nY<4a$vb<g*cQBeRc?*v@l3Ap^abIg+v;PMIEe~PE8
zVVmG`Z97BXh>!Y$rx?-7D&Xfh=IKqzx4x#Ux4PmFls3YYg3^P4YhzhfY1iK5F@5?$
zv^is<=?9CyP`M=HFb-gAHY2s~pI(zWW;G0A3FDZqb3yNp^ooB%&p(k-`}s&LQ$Vq>
zhQ9t_?I*|j=?Pt3<;#AE+-nz@PqkHE{p)|uIuBg^*p(-^z=&0IpE&Vd6RYRWhqtjk
zCV_s*e;hx$Hsv1xIe_gJ{s*T8NUSk;|4nw~`zL7){EODVirwhwKA#AVqEF^tkL|31
z^BRx!j5`F<lkNML8@yb$q_O={>H_1U<5=2JDD7w3_#RD3rN{)Jtmk(^g2NTUYvz=&
z80zu}DWu6IQV7%m#gDH+HXhZ43Kjnm3~05drg<R(?PFyN7kXP0ul{)%y0DdyC_r4E
zu<_0LsqMe_ZpCY?cAz=&9$1z1o2p2pjK;&!h$~+$Qs#to)qg#|)^EbGB7(i|oBw~<
z`_8DQx^`Pcffo?4pnx<LL6qKmQ52LWf`Zhb2oWivBfY6uAc)iudX+AO-ccdcgx&%K
zr345)1PBng8}WU=@0@e*825~EexCgc27_d;wf0`?dFGtYe5SvZ)HC(By6x>A=-zl=
z<o2ZIWHP)z{F%VlXuh-Na2q@SKiH@Po$N=ej>g)R!Qjr~=N03<)qN_yod-YbGG?{9
zCH;deGgr7TU6G|bz8{eneU>x!K7+S7;!fC{JCIfij)bdrn{e<>mDlJA;&0MYn+Gc^
zjF(aXJ*nj@5%zbumQ`bQs1<Run|pdOM&>s=dg$%}crUBEVW0tkHQ)PKY?|fUX;C@u
zW(l;HKK_zG-F_+_oHuo1)qUORPf!fn7vo!>Lr_W=2#&NCxfr%zll2VCM7RB2Bv*No
zW^CJD>(=KUa$U|~=&pT(l0hcM!q<4t7QClN(|C@KNq|WDH+1pV@ZjTstt5R8v(1eS
zsO8L==E266?47i4qblQfIIpmd+MoYY5hhLZXSj}lwK!MT`{J@K*kc*fzMq$@wV;pB
z+!Mw@`DK}#Gqo>%{R8@qOaB4=LO3Hnd2C_r5F$znt6Nzr2ujT(#7{sMK>Whus-O2S
z728s$kT8m@a{qU<PeA+_?OW%5*i)q~kLVA&l|tSY9q#O0whZ@f0BiQ<)?D+sJAB8X
zmU9!FPk*;Wj@khGmpbDtbNj~gE+}!CdpDCPC!@6=97#l;7+)wOIb~bp>C|ofUfUOQ
zU*-197rdi#V^<M(Zbum>c05qk(BE)Z)%o9W*Qi4agVW;$uQ&2%w(oUym6vXMwDso&
zhAJ@K6a`KLuyX$iUS~q#hQ)<`YwNy^;K-L3EqUt;&xP5`SL~#N2|8!O%q$*>iYTL#
zFs^fKpUPKROit=#pT4V1V@By1Z#lyQQShO&R-RqbP$*1fv!zPr4#FE}17mpUZy2<)
z`c5~-dQ&&?ASlXkO<6{(CM1%E+?OLyk`08qc7z&<>X+d{6}ZbiCifJ7vvL#^FJvBy
z>eGM~Pd>BBUH3fK_p}f)prUN+10Cpm`uws}Zog?z)gzQy{{pU7Rb(%0q2MG3aWrCV
z|4yVfJLpABfa8W>dsa0~`I2WvW)$O-5iNgu&LCQelWWd7{f4vd(jc5xNL2m>=C2`P
zrGkbpFAjD@tGQ&A-_H*(*EXfHJ6}^Yfr+_b{wxMmJ}{J5Yjw|ggzLvhXsMaf$sui=
zAP8!S_(I6qulzDRK}x2%;)-he2owv%cErt}a^rZ(Hx4Ofdh8?Zeg!s_&{Afrd)Tz7
zC5yAUzb6Mte!rVE{=xP3NJpHGH5Qj#*!_zwtgAD<T*PbTpV1S7o9`)E=}@2LgcwPS
zaI}NxVQeZuWA9i)Du$0u<Ivk|lJ_8+*4+&Th`A$#q?5@#xj(+BR+FOL7mNX#7+Yo@
zV19JHlvH+89bn{(r?|b`Jx){8F7D!slx@Bzfhf;7&{n-OlaHo%>{{wfJDH17rHRr+
zOu3OKAB=8{_MNj}8&IBa8qjNQGn;LZ)hlzYEnp9cOgu;ffquC=#+IcJLYIv5Y0U|+
z0jfa_cPf(7D{7wpoxX`UMK7-3Od-&k&Fcs^utR??U4Sntp80aPOlR4x`_KXI<zek6
zE)l&+IAMo~a#OxUdAeIZI1*@`I<2muWT1s=kYBV;{1<L;K1{N?61pqVoQNN*uh-Ht
z-zp!HXk0{JEucBwKVzu)y;!QM|3X;79v~e!mpip{Spt9&%~!GOqt*4}*>dEQ2$7$O
za9#&o9qmk|PffQEHBbHSxulA1ucKDjz}53C-7*uFHiI}deCo`2(0db#v~|APSGtzA
zaO2Mo7SZV-`zSL3L|QAII^CQ#rxei0e(SWbF2yaMH<p<`BySn8pVHGD@SW`KsA(!r
zZ}oT}$DU0-%CS`|6fRqJd{<$G!4@iT7orb877#nLZc4;50BmbcfaMl%@OHIvp9czF
z(YLZ7t#VtmboHORvPrj`;oj|tz3Fc-Myhyz{+df;*Ehk!Xh=Lt*{-hL(f?_9$476&
zabZjfzE+)1-SN>Buiv#H)0YO@rZ44AoN`FC%*09X5Wt60rvl`^@a8koRje#e7T~;V
zVprL-{i-Z3rzeZq*6Y;%I29k?-1f%*p%yVqE<lOAbLydV3D1tHV?XfHj>b8mx<7BZ
zS|lm7AnY{6o{Mg@6IhRDF17o4bKAa+v7Sf*sjIS`w*cpJ$B1)8EU0rs>-vl-)XF`|
z9-Z0&UN5i}GddA6>Z^r>CtSdtJ~PkqSTWe(;wdlLI1m)*M-~)A^Cpy6D2@}&xCHdG
zMw^#`^u<mUSK6XKjJ}Q7gKcH6ye3|mf3^O`LRbXMpiD@tiQqK18ub5ciPO#KzDvs6
zVm5H<k+8B>=ob;RE%Nafb*WeBSK5PTZ*6&+&KsAA?_9^UZy{dPF4BbjQk>cZu%-f6
zL?$^g1#eujNF~jqB;wX)`5aAnmCqe?C9XO^3Uw9eI|5vRjT1+?j!*9iit7Sj5wd+s
zDX>S6y-vJqX!p?KAE~jCgf+ztFzY=##bZjie4*2YV43!B>5kN1(Oi30v$tHWFHCp6
zb-m^d?OL@YmFBhj!!bAt)+Kn9@o*7OchqQNN`e=3Hpms3g^jAG!bHgSt$bc2mt7k5
zF=d>M#GWsvOhu+IC&GE5BcvW=1Nx>C;U17X*!bskQ-w*g-te>6Wmdnxpf00xl@8kC
ziV^?#Cs7DK``t?|S^u25X`JCjL`*`k-SjD*!>$j!xweapua#~qWkj&4ieB6kxO?9n
z`+%B<>Oz%|UKPZ&G;i@=f(C_c574+}rdgj~m3mgCEqK}9yd7@<D6W<Wx<(Rv=Pn6(
z;xsX5aqT?WK2Dk4p)o_10aKaF&W_Ugi4dqq>=BaqLn%T`;tw(ItIZY0VK%MfQR{HA
z%tb}SN&9Y3A(y1_f{e;Gqi$GtR}C|X-Ngf|-Plrf&}tXBP)Dr(W1PCicTfYwE)BxT
zwDD%w83p6B!Zxbi>~_m+yt`RU<tK;*4J?cu4x39k&6J@wY=yq2>+vEx6fcMeIKJbD
zw@G2g2T(-uhKuh=k@UY^FfXl?th{Gc<22c`f3nN<VwB}kL|2j{>Jw|xk5D?1y#e5l
z$Nw5=w9NvvZ%YBm<^@3JS=+$K$I17vgFc;2enD6b7!TOiJRpTaw_mU1Wckz<18RY5
zX_niNZI8ek<hQQpk3VCorT%|VAaD}$k3gUw`kw>>Q*oqj@^@_C_Hi`Z<a#EIV`N>}
zv2~V#bR5gdIJ=5$yQ|1okW-5>Je@`^6Y`ZlroO$)S2}N|St_{ZoY=Wbz07UV1yPnh
zzI#^GW^-EvLz1JIlD_NouFe8aHpX^Z_t1tqiF1I5wOVnK|E#rwA%*KP(;pRd&&oNo
z)}9{NJL!>I%G&Jx%lJKotZOxWPZt<pIqJ%MK%{>8{GVglns<N4vMS5WS;o~raM^67
z>~y8OUV2pltEohvE-ri9#)v>^RAY13pKx<c!dBF!qJHr*l-{O)VDm3^M@lI>=jg!U
z#oHwX`PWz_M?`^$_Okll5pAKW=lj#IP@N0gX*Y{&9z~~J;MfBamaqbfoY$5!m$hqM
z1|ptQPAx3!VW_R<S^l<bTKi3!gPY#7k4Z8vRFS(XyN7+`u{<?bS&g@zHMbSR6Pu+`
zOJ&o3seb4R=4-SQXde^5FX%ap#NC?ZmN@HVj|wJtC$Y4(ulgwg1@mY4jK@|t)oD20
z%*}~l?2VS@yCNb$6DG#9W@IgUH<I!<KQH)#Oy4_QweG|>T$@Xe{??F8-H!sSWGr=k
zv&zv0HKU&yhxtui3W8N#<UH>*r~SoeT^eYvS@}I|C~LOj2Jg1Py19Uo0O4#O>cGcg
zD+`R9XxIcBs@|XD{c6+F?i(Guo?x9BTLJe#G-k8;&Or#4Ghwlzf$94ST+E?-DR*l;
zJmnKM{RD|i<?}O>j;?|&0e5aqi!XHP(op+tVZWQ9ix4MYFS2>wxSO-MF?GTGeEG6R
zPfDRMbITm-2kY;^knT`N>cC9UOi~g?`C3)rPR=AhuI*1r?4`-dfat@N{C=MMM|DHN
z*4+kNc6RJezV#ez^Va>1RLPQlvE9%}0tRFt`gwTUYnc7`Vca(XVDe!vr|aXj6-Lhy
zm=Ey5ygmY*`{k9$c$l5_Rv*gXnu;s#7F9ME^GX+3K9!kDe)X}Kxgnp>>RYpJ%7V;F
zSmUzn|1;PPaXkul{elOz`B+J97*z6Y*aNw}UR?pOx=9dDV~Xx`d>zoc@EXGGP@$nm
z(?3?_;?x)?1tOX<ZjQsdDsI{cKfE`-q-iRD$VUGDj^qbbemi!+ifql5IUc}H`4Q<k
z6$V!A5|ld7y)9qhzJ{NUs=(c&68O*Yu1bn{Ho~9L+;1yGTIJ!j{)Ef-T4#(uYq>JL
z99^})f1Q)!Yux_D1jjL{{}S)oAIH1gqc{Nr#`c&YzC!6Cafb^9I`cIi5F<mk_mBMt
z^?mbKpjTR+`ZzgqceoHwKaog6_{s%$r=#LFu;~__I3KaSK__Zv9lb0bR$HxEEjLPe
z$48vOaJr&)oO{vcVH{ZVCVNdQ?F7olB;evKY?~=tmlq}XLsmuybTuv_Wm!X45S1sg
zK-VVaz<r|;jc@bJ7GTJj{Ie?#djn@gatU9^J1F?ag-&Vr)XuJJdDN_kvYRiSoSO4s
z?tsr85u{5BQg*^wm^jYO9xq+1nw6U;t`U+P5Oeau4DVrsFMQ_%zi{r02fZJ|*fP0A
zS<mc1oq-tB#!iFmK*`CmupP`UvF9?cx+vY#4C`ZWazBcE^)Y?Z6wmecQco(WQSrFO
zuuQ9t{wlChds!DZH1AT0ruSpaTF=_!c$Nq!X)lgR=yH1Tuj9GjqE-oS|5BG;xGC{x
zIdv?64le!D-Jg49ro1le9bzstB_ySfNfVYj<}*WJbELp>@cpg&TPsypWV9Y~%g8Sp
zC%Q=YuV)9p%Oz7rv|8C69Nq5jUd<x<ltz)=_LlH(!VF!L*O-g>O_<F;w-HWgdRc7w
zd;;hFAnqO&5}s^I*P%(11)nSu>K+0y-+3<|m$yU+c4hQRfji*Tj2e4=S8UYVSI@|A
zBrWJpn?sK8cFj(QIZP74=@&bdU^z2dqSGUTS8l#XNNnnpMqI<x@!xG6;j1>ZVF|ww
z#yo%!;C#N;lFZOStNo%FFohNYt<Dt`Utih`YuohCYux$XC3zk*D?UspN(7&8^Pn-*
z@!A2nAc={sy5Nm~k|EWv<3%9#!5_*Z@wiXe>fFPM-MZec*Uk5uYel?FG#Vs0ZDomn
z(!j={JI!q>o9+gPZ+)mrUijsQcdjD~_z~Ks4Nv9w3V2yZH@{tqVIhvUn^&j|dxkTI
z4am$K@iY8zfS)0YH{U~p)mPM{rgLw68owo!ABX<S?TDQj?l1SniY`z4OCp}rFS11q
zVal}JNUkA~HngP@@2IMwmYH7rw=X5x+9~<;c%fWSS5{qJb5p`!D%BdAzFRdwT!8O}
zLl2%t?=xD?^>bHnv*-Ww_$COl!f04$x?gOi!+vEkxzOTCT)mHJJhZ1u_bJ&)vC~f`
z<^c#~KxTE_TIXYR&54xmO1HXPV`mh*8TnVJhEqrK6%tTF$pLi!3x0D<dMG4wjsP!P
za_x)8{bzq!K>Cv}AQ=9ml-bMzb!ln6u6-N{Eh>4ba{sGqPVXF_Q*<jc$7a`)QPd{A
zY69LOP6arcCsP)Z$vOr;A!dM&0adHD+(qJEwUYjsTvgxPNmkq$-1sAvXJ4Ae)Eh=R
zR#Ni3_-$9YTRTKL=QJbz4#seI`t1j`_ex3@kEEBt<zINDCf<5BQ84@@D5)iKg`}O5
z*Fmqk*<EY;>(%+<=ae9Q1J(Cu!|c059tn1KhrAKCz1>T!2CE+!n`??t+)%8Rd?Esl
z`)$c-5r;nQYD@ooo;*D1hI(%EwxXodpuS`@rl`-Oz|uR)w^7IzFTxj;00n@YJC4s%
zSi)T2xjRZHfsd5M)$Q}#>YXAQR(m=d++i+ZR#t=moDfk!2Zf04_uo3`PxoC*pL*gY
zL45XLShGgMzTWbO)8y%#=sUwhSsz@<+D&s)W3s?=CqY(+ta~NYoP0{^x}t9R?T^;i
zU!3*Em!>2<!#C?vC}jP6UhT?m+jhMxWV?kyy@Pim)b5ct++zqe6jP^sAxyhPdDTq{
zA(CbeO!in3-rZxe-MHiz7OI*$EcZiB=1oMBEZWU`)MFa+UYZ-T>o0e7K+JxE#77XU
z>}0QBLzU1@{_B1CjWti!Rbp&z32YN6k&jMIENXo*+U-V*_!v0aguLHpkgJ7cW9iHS
zrlp$P-pa}3_#FH+-R3=BtF_@>d~hU!RP*Y<*d=L4-JIs;bBdo|vmcNO<mw%MBg$Z6
zy_24!s}$r57k`A!{RZE?;^{YcDf(F##}jX|k`Da-#NYc#2g!h&rd---e<6+q_3}Q*
zPlw1>w{_ne9S<H#l+MsUt3WMJUi6_ufIX#^Fy}3A0xYco6;FT$2q<1Yx&l%ffzo2~
z{<nvI0D{<t%YK}^&Igh`+W|(&9n4X29>K#}zm=|l+AzopcB5&A8W<0dgg0i29EY-j
zjh+94K(Ndh-`_y6^zwfJf-Rb`eZ|MnhW3f(wjj>cUvD2zmiepjk5u))Ixo}Y0@}{7
z!Z7XDD3$v14i~!CR<^KAaFrIEqpOhcycC$nq9oG#^++$?xAe*th*=+GQzfJP$s#YK
zj^_iNpv+1tN$VHN)-V|t&54M2%n4T@!`a-Zbt2fuXV=x=^J7XJKd|6rI}Oe$vb~eM
zs8Q2d@o93Q%E^B9dzsioLqns1X1V}p>`EYR&MiZh{=-e%bG_xZ6?4cBTQ~(?<FUCw
zf{SOiX+dy>ve()cCDJC|s%xxvD1X9%6939%JmMx<fXZKe32n32V=fG(;JLmGk+Q)K
zj}}La^C4G;k0(7_@;Ja&4eg|mpbda?$yC8rshZTQZOzU3W8Gs;g$O@R<a8Vj_kS6z
zgC2ROXcClmj%_V5u76cc!SaRHT{e?jwGj-Q-@lw&vwexssEO+{<^XQz9FnhV$PIc^
zVy(<#c||8A(rO$0Qrbdu-Tgrd_lhL-(apTX4qMMilQNLUCxqx%VDE8NI;lIZm*sp9
zxeRV_XqfB=Hrxx%ZKZ%Qw*SO)ip5(${daf=k$b9p$qqN}ovTUNF6+Am`<HH2$R>)H
znagmI5OIO?+cmAdq{;czkRm{{s!i=1{L#W<k|V5+*HeD=Blsn1n1=3rOz5OL;T&=+
zl+jq>j4jt$Dx{Fm+{sFs^NIF!x{Uw|zIkA;7F0lE`~bBqtviOL{{w^vZi^a=iZi|X
zVPvw`l?NAEa0)NYvK__J=AQQF_z<YYS4*pXEIGX<^#^?hiz(e68x1J#ZrEI^!^9nE
zqaCA8oQ1_TCCKp+9-1uXg`BYQiOsh5${bB`w?{R0L!7{6il>-ovFoX_k|e>t^4k91
z%vGyg@Ll!m+a7SMg+(5^u{i!!x3y!%>4u9!leJA05UH2WUNA>ArP&+pA_m{2k8p(K
zLME4F5I_9p1HWY=ZBQ9OU}Wqo8Xg4{y_R3~ZjvsX32eeWeqax946yP<)8UX|al6AZ
zwKz4@b|RRFFJ<t5r-&`wOb4;}ZT3aW)vmJVf*|%$F?hWjg+$|Q{fH&2^S+&Nj~P2w
z%xo2|n(q*2w18ePz%^AY`B2m)iAMEXy`9fp>5~Oj3%NIyiuJHxmCi}j<_p$yZj|NR
z^PfY7&`m1xkHgTL-3e7`hF--V7D~l!*<=cX6R4ithvbikNpk{8ti<H*c$v7+Z^b4K
zwJb+C-8QyOuhX0rc0(-^eKy#P&vN@Xs7Ei_i_7`+)QAO~Dzc1h&kqpY=<l-RxbjYB
zwj|fXl1dyo#TK@tC{h;U|5HE7tZ1?~IVqkD%MbYa1I)RzZdQ~`RU&KauKZSnrh7Ya
zwZB+p+zW*f{yp^77g1a8ZJC^ktnD{IhH_63X?DrAMZd?}gtY|#$cFXxI%hGhzJ{nz
zjCOue`|I~Hu9=|ND~{2z^7ZcDV>?F8<;lJ*^SJ63w;^<2KH7g^?)r!4{}Mog{s<u5
zM0)`N<d<<hLl|3NghEoq&0+*e9;6I6E2N0@leh^;AZ<howdJlEnZV1&o^+BPeU(lF
zpjC6@US31eV=be?8ilvap$^UWNMUNmMGYsol3im(G`<=Skn|5s2k<Mt#Wb%HPn2v0
zCZ%;0d{So_gRUty-iXlQG`JJzN2F3>`T&Vxv-AUmiV~c-H(p!|o-P7o8DD{7n#dHH
zu4q6+f_@IOn11buQz}&XX#CDAz-H0*hvkgy>}664!EcZFbPs5b;NA*M*Yas~6JKcf
z<Aws49@~C)KjY<XSD#G1q3VO<%qg&#<+S>uvH8Ya_qvBQK_BY-Ho2k1=h4Y9V-t6S
z$Ky@Q{LwN#8`&Ln3148=HzpdEc;C!X3ufbss!nvquiC?nhn)Kp9@#+1q@EKG6)3D6
zo$Me^N<BKU57BT#Wr_0z08$=!gp@C-0Qtg5-C>$T>dE>d7@&!)a4Rn>cu0MoQT~OZ
zrKot+>$2hsI|KA%9fvk?Vu@TtyCL2D_6&R2rmV<7SkddEAwdx}B9r+E4;I~kA@x<H
zF`AQn@4n}lbH}YUW!xSSpZns$_btEaExg6>m86Pf^xW#ava}7UwBr^?=5C4;_mt53
z@8^Z$NIq%XZJ(C67_L7ZKEZvQTNG5zKX>fdl-PVJ5i&+e<sk29(tSD$84mT{`jE)x
zHK4`c?*z)^oXsw0BKw&Wa|m!VA7zb#Ea5#9yzwofj5}P~_+FIh`hML|VEm1UEz8Vh
zSHINDhPU32xrz~;huzBRQ)H#jhjndL3MC7~iKwvz7<Kh5wM(#$0+P1N-U|Lc)0^N9
zd-B~|(erFCFQB{~m67<GKImh{;B7gTM~#!d1qGIwM%82T9aUnDP<#(RvQ-ddBNvV3
zNYOv{#%<EbU+?&UT?00NzOJz~<6t|-JQUO+cw5JAU9kRlK62{yCcWHeJ7Vk0Ot|4K
zrP{tL_5HocgW?@JcHFyrKsU7&Yv@EGk@z|(St@?=K?}Wzi}%|Y0Fv)YIfCRVf)(1H
z9YONB@q{LocY`xI&~CwgWjUK>2E%|5Ky6o2HjRtxPoZU3_%uDWB#?*f%NAIEGJE+|
z3PFvC??&$CHMfm<iFK5Sbyhx*YeOm}Qr2K1z+Ws7$BUz=%^bF`2HHKbV~fLl*161^
z9_7^<+af3v1rh32wyUPK_x=ss`wtxh_cgYILGL#spFomzCLqb2%Wo4OP^HPox&V0#
zA%Ed*Rfq{{C)E-^xzg>u0<#;<d^#42pXFSX3)2yK**lA7RCgp|mClCk|Fpu`ENdkj
z#8ljjH%30FtZxQvtS3y+c=mSz0fvJ8IvElUh2WL)ZL#CWuB4;+Fp)bC1W_I`Vo)C{
zQj>QeA_L#>cM86kz>mbwVrn@gdzoN3jyMa{@v2NcJ)%!7Kgo!Uo;*+?IHcf}l2*@}
zyRUs8pO54IA%Z={k<Sj=gGKDN3O5`f{xT*AQY=))J!7VBv`-|`>r~E?@e?K=TQ|o<
zxS_iNDRu2@=4|Pcxy6C$x8sVQz5WZ|&(uw3D+$=lD6N3U;5FxmLQ|)yq`j-_P7lc0
zOnB{$^gb2|)eTNk2$yD=%^x!W^jzq$m}PD3bD80K;Q)2f4cF?desii*-q^avZzXaT
zP1g#U<UjSDIsKGPlvQpRp*r}{R30V6s2*SwTcPqT9W~Rb5`3?m#?!oPo5y;>RzrZM
zRSug2;%-=Fmti$-V7ws0CqOno_`(?ZfP8Xb>v_Fp9<HRMBhG*kE^v1}9JMFqXCgAS
zJ6@^eb4>E5j{a6gIesaHU%f7{Q*uT|lT^@rQx4Nv0~?YFa-^j)l{o*y<~S>IMhqp;
z<i^v+VIS3G(%v1$C9{#h?rBvu*sdt85Tc6*sqa$-OSIsfWOjVja`xlZ#;t5uK8cVo
z;pWL2N}XSi(xNE?AT4rOPTT6c5CDjE3ik>V;eHIPVbLk&$|AkEea9(G%52b2K9r&x
z4J*B}EmDAH)wMyRam5IkV1$*?ZJbQ5aA`0;G5)e<24*!KEh4}=I;2ts{_PI~(0G5-
z3WgRlpZLyFX%GYj094$LcY>$F;@eYM2<7a_tzx1zI`&z%o`I7?CVU9yC6Oz;GQc^|
z4%Q=9Ut#%oR^hd`Y~1|Vb5889(8rz=#K_gX_P7+MlOCfk-+{zw+U+=TVhf9QP*fa+
z)Xn+4x3U|7B(#AI7$v72dn_}j+8`)kP(Lv#%P69^?COm6aJJ%U&M;<mVVYvs<k4cE
z;gX?=?eKx(;w?8?>(Mdn#Q19W41G$$)K@QU4e5^++a<IOvz<3+p1PjCiA|DTV)^~z
zrecnhtFj6SbOYSmP~)syKK4i9CMGh4;dEGQ7wRktuTYByBs@Y#5}xCior0pV9mEnQ
zsDu=#*5%mjM?3`RF+Kl-A>W6J25T`X&E-J*3pPQg7qJ&RDoYv3F8}=Bs~Z6iL`!RC
zbg3IJLftY<dU<`KoIgrte<ZNQ7JpZi5T#8v^$?6>emkn)k_`r05gcDlkd42Y2OxwF
zSwyrvKO+ukYI<7}gh(C7l!Gh0@;BwsRPvnj@7wfG*=ADz`$H(2O#lByGJk)<@2W)_
zX_d(vF%YC~QCU)T|381)54Q-M^G~9nFlN4g=hh=h*6OsWbsenu7qBAbm|~KLFS52B
z`5#m>)Q#hjD=Oo$e+z8hZt(z^xoHB;(zhgm?mcQAnj|(r^?Sr4+&|mo8}Lt!WYemI
zpQy{$-O5dP-(J!$$U4?(RX47`S-m**^z~)-`Lki;r<k^|O#-+7kpcOqw1-3@1e;>@
zt3bU{r0iA6%z7U46V3Mk0!;niKiErD+$Z~)cSgugzy;Z;p!=RCvpK<$@~mzueJKF0
ze(MesX;n~_`Qca46I<eT!m@EftYSPlE2c5`z=10)v3KK)UUUF5<kwpMNQ&<8impbK
z1ttl?Y)aB?aef2B>tEv1h24M~@G`Xn;d6wE{QD?SBiVb-(>?3jo?xc=!fMRuu#jm&
z%4C83Sy^OuUjmZl$v^d=fuQGOmNVK^0UN2^d~Lxh%Ub0DvaD=WFR62A8E2o>6lj}R
zh;F>{Je6c014uINFWSD|Db?Qt=v2iS^s@Z}<o6pAFdC_k2;r<fXUFBZIM5+^7q`va
z9#EU&oR=JAGTACuF|-BfyZ$We{#5`o89e78)&Ib+B14aQeN(&rrT1)TNF*%!*OY2L
z(>m!0ZdUvtqJ_k|&pc(x8}<AhZt_dPix<x11zNdGFxxI)wgU7USvG@(sD}ri69J6{
zNvMAa$=5-~BjR(n0pN)>6b%4`%hu8gUq)YJ+ryzX9GQi_CnEOiB<Xuq8tb><3vI<P
z@)L$cJVkB6$C!h#C>jKIRWhBhCGkK3lz5|~@9N*}6e${!na%g4YL|31@Iw82denaU
zHJlx`^eP--+D9O<mR)yPOA{T(ocu=Wp|lv{Qg}oZopL+x_n0!Gj4O9(97T1*cM;ay
zkR~isO^O+^CoTY%3xm>{7s5WjeR=*FX3Tdk1vKf;=3W4*d&H0kr||I=4u&QwK93M`
zLRAfTsr~Dzi#*CrIRS^#S@d>Vg0qK0v<jN+B>K6&7^)a=>>9z&Io#6KmSHhun&6dl
z(wx}zjNrM`&}36Fh`!a88GYwU(sm5R4LAf=$R2w)ZWK8c2U_v>a<8Me9*sSTyV=Cm
zYwM^yJ#%X~!_`*tpX%bZJ)o>HF(utZH9ONecT0p`o$hAf>=8=vwGYY~*3}hTra<Ku
zF^H8VH2pHQ;rna2$m4VpX?7*f7`b~_nk4bSJ+OZN*CkypHxn>M<6px`X%CJJKT&1Y
zufIIkWXgHnAucjFlAW{1`x*eB8hGD1!5KSeoG>=hd=Fy}Q>iTVihap)C~nXhWhIZ3
zpo3l7=){OLLJ=N<Dduu}Y6MrohWv&fV=7EqH9uC$=pSysq!y|q|Ggd571;u4{d*7J
zv@9#11D8*IrzNI6sdyfLUAy)3>HW^x^osgS4N(`b$P?lv-QGWSy5}?>`eqRPby7Co
z2<;Nk_5RgR^0?;3L-9^<hmDBW0zZwqzwzGkLu7X0WCXo@r)oO!*QJmsn(|H4GV;7O
zyO%w;fc&Wvx;qs!4nBjBCjsO0`P<rte_n+^Tb*-Zk0ja>CY6QaWzON*z`(q9dAD+1
z7oz9UYiU#C5_G$N{zDO3Nn!)Nc-z^%ku9G?I*9I<T-L9>4{;@HPQgC#EwI4Yp5<&@
z3D)8Sm?)kk`JvPWs#{e0Y$h$Y;j_Zs@LLW^kZ>jXzqNu>q1xxdx@<m4)tQgzX_*?-
z_;vfci*BI1>ezPSeUBQ9s}r{1TtAD-K!&3FPnL?Ac?bRCI*$%D+v5u8OEw1HPpP+O
z^zn9g?l+iC=#SC}G*0d<13m_70ulnfr_cb_j@&9a548PwohZdlMPq^*!?Q~8{8@!5
zPMxi;dwNh$HoYh`TLDe4A?`0U?Volv6oDZ6=C=E=+42oo<SdRG^_z4zP4`hj03Ab*
zuPJw+OGs{vzmE%WuywW)EC2oj#VTNDo%<X11JfENFYtpb&L<DB$CwBF(J}xg6?6aW
ze8-jWe{e_OBmNr-&wsw~zfp4kzh6WV{{zGxmI#Pr{R8?vkA(%0HF)O<_x-)_$H>(o
zkAv6>|GMkUNy8|INuXI4AMKmHH$Do0e*uWn%9?^==SF}7a2Bz!Sa_|I{q?>&=XilX
zl4H`Z4I;lNQI0+8bC3J}HaN6A;w`x@JtcoSw{z*<woD#N@|M-~(7F?@#g4R0h8;N+
zAmhy{k8os=TPF&AtgEHS-Ud)orH4UhZ|KbXdv73V0`EJu`V~tRak=d~)q(eFxH&OX
zaeLj{S1Cke?*~WXo|w<CAm`W{{k~A1>a|@NS1!ZuYfag6_o;E2#!OT8nlrAlhwWF(
ztBbnGE(5IW^i%XlXCXk1tABWc3|T>Ew#cCtZ_*~8zH^Dbx9ASZKgUyEm;+;6Kp>y(
zE2=Rc2p{LefA$j^2*B_|YN5&+n&&@!lUR<8StU$DMu)pkZE!QMEMj~jPkb@eW-<oS
z)czx5gzQAeBkO?<)W@slcYYf2KABalh~`=k2cWfxG6t+jcqEOM9m$c>^YF2|O5WzU
z{Fc1*VTHx+c;^S<-F<#Qs>(X25wfy{apGdjwy(B@xu*rIaRps*<g^Q^kZ=um)W0cl
z1p?+{3Rc9T6igb{%JYpdw)_DWW9tm^{`0*wPDq8dyrb8w0hlFMKOWd*Ln7vIf7BqJ
zA+RO<r$}97a>AvyY-Z5*Ch4bce*~B6&)eTc@50FS&G(?|sE<D#D5dDngt$I+PkT}k
z1V}4-<N;{~``0pWM3-ShNdY<4W+Jhq_|0T9HL}~`zO}J;%H`IjkB(O%3*_WMjY}(?
z7@v5D;^;?>3ScTL;5qtXxR##Ioav}k(P2LVb>S}$FWInki&>2ay3AM+Aowh_UxB;n
ztJiKl%>?sLUsk<$IWE+VL6BaLI%s%=a)e($42e{$l+YD*@$W(OO>hNG(*zndt?apB
zoMHkPj>5xU8p}F@>5*RCw_mV7IR!z#<LR9tGTc3*{HU>mL2POtzYrA8j}LL(lkgu+
zwLNDL%@!Ov(W)r&HniF6R`)$XdBSbu5e@htJj2#F_+5{__N%cy#&FAHZ*B=@mqvR|
zwm53V-=;)K|3$F2%f<Qe1NK7q#y7?u9J&1Y7tDbMiC3067wWoOs)ZcatS9?qGpDyC
zu+pGo+oKS%Io}<BK;%1k;%mf%vE&tbr)S-UeX|ZU(#nQ3blP$w$b~BB<ZhNB|5@ol
zY4y&KR$&m@yVEU`s3NUyTB)*~MtDMZZvowhNlUDtY227KE0(hr;v&v*hLw3$g#N_9
z4yDG*5R#1ZUxg`0xl>H#-=4iYJbYSxW3PdBtT%u0-zF#uIL(;fGisW3K|n;fa=Bn(
zodbaOj)E6r7xrgV?+YLPbydo#L;lCR;65$7F4eKVAzNudQ#!^@mzvg(5W;0E&xVKg
zvn|yF`79|x3<wSv|B&qNTo?QIRn4|f5~O|QM4>|YGICRgXS^CriHZ6ydoChuE`m3F
z?>w_KH>ZZ-`4L<W8ZpHUx#d#XLVaa}cdz>yxTpxXDK<V3a!u;p$^(DR6f<99=f}F<
zBPdiTFAwJ|Z;cd#ebyGyZl*7;cPL;*BpM42!aOa)hEjxx1#L%?6{{5ZH(im5?a=@p
z113-Tl&kC5j?%skIaK|{Ig=1dp=F>VL3L8~V`I1olz@_W7fE%`O%ybt*a3I-G~pKL
z2LlLx?XHO$T66aRy0G*_zp*BMG+W*>tFu<R^JRWSzi>wD!8Vfvk-?=RMU0SI(m9Nn
z6nC&SMz-4F1p(KB8vCx(^H2Q+_4Yd|XBG2)W|qGizyJzNV08j0ZgU!;x7I^OtflWN
zxLhLn-&;Y~G{_%Hub#B|2ih+)JT>pOytED#-+vNfVAJtrTh@P%*tMD$Z)@R`_*QhU
zPeAQvzN9R2PR4=;P*WR=x+F80;W<GKf;wjNA7NLaA2WHCF)TC_VjtnCC6-Q26S<~_
z0+TJwfUSv)lW-BQal?r$wSy&C-tL>>k8FY-{2NG)*E8ctD~0&Wk9BS~G}dL>t+}*`
zLc7OBmB+n`ChjHRy5#z#nFbU(!bUv?x?#ln3QZGT^H;sgOJ|WUw_;=OJY2w%802sq
z=k_DI5m2PENF64`sxh`s#fk~VMq2JoAtkd+Cn7uaX>R{=)$mz%chm%QHx?MmXN60C
z^G#6B@q)zZI=i$iDgIPKkKUn;SCz#47@PeWn?-pz&iGL0Fl)@tuYT*V7mK6m;fQIQ
zyJrLxJo^3f$(1no1T$L)WxJ(-k6=$eBj;&Z8?H>{S-b&6-PEMD#2F($to+DN_Y&6q
z43FQzGU+p5jGv-momh4;!kb-)RZdgEU~lsV&4sq}wo%V+9CYFbw{?1;iEien{Zy`%
zLDAO;mD#IWn-Mw})~&O)a{J2#BCfQu**SIC*ZFW9<yl7LHQVly${7LO0*6Ga>G6+8
zzzQH0I9xQZ?{#5(E#nq`z9wnA>_E^zax~e!ho_-O{6W6LoVKR3Q)-Ma*Eo=BO&H6A
zL)=bNXom1167L9`ihF(HJkZa9a0J#isNsxaL)`6CMqO&YVyOmCUsR}qoO#PG2?Kxi
zyDlT2kY(}BP&v`W-(vTT<7<~M3$oHaqSx?VSsNwd?q)Sl1-XifW0_4HezEFvNHYuo
z1_S*bT_P~7?H#bs{pnw34QFF*jop8?zXUUhojDXm_?{x!m%nKP&=j3Gb{xkq4pU<-
z64cwmY3yW<ih=eh2|G}Ut9%d;n*#w@Po;=K9k1Uu3w_(;kGkxCpi-akm~ZV;tHqRG
zxc4P_y5RSs_VikGV=I!b3*9<huKkcJg!b|he>C<w^cL!3UiPX@5oSUrQi8M2f$?5k
zCpNm0^RNB!<oV>R634fNWm{g-@M7Bl>CULTnDa6o8FRPhA2)CjW9{yD%1L{ef<F&O
z@3W;dIWL{0D%kHF=14ImRmmQ0jM9dpE{PsC5zTnb6Qt1!X4U)YsglZGcgefT6VjKE
zOaOx?>|uV|Qzg3i@>j0u%X{r=UnrsP>+3AyjIv$#Q+L8WTQAw|#Ek;);D1V0=?U=+
zx?H;v0B^dO4dh`(v$mwtGBqW?q!m>y&ZLDkAS3<2sFES_p?s7UR`Ide>wTcl=w}j9
znY#`b67$kz*mFK8jP=E&IRx)}$<y-teQdHJneL0@WFF>`66V&MlFnO;veM=dTi@I|
z*6CCm-;n??^~^7cNiE6!-z&(VNj5_5dNt?MKgd440Y<}aMan(+;Om6Gf+~q2Re-Q=
z>KeO~%gsZ+_TQ;__K08$T+0_+2~27D>{stn;jSgE!_uNYYsf}!NH|IR5>jpwZ!&U;
zc@0urzFfHX_>q6hy76pht%nh;AtS@_&5CQ9{y|Nh`AZRjm-huCgc`oEeg6+DQK9TG
zF?@T!u%=mOQ(sSxJx9@QTqzbm2Wptx8Sg;0+_zi*Qdt5tLobtX6|SvbEU00ABrL)?
zCwBKcet3JKqIZBh&Z*C%*IjI${0T_lbNAnE;M5Mk+Eg+m8B*8OAela3Ux5os-Uh(J
z-6OtZL^7b=OGM47@lzwq+^Mo&^D4l0_ysP>YApCbLPDjhv~2f#Di|+9mwtyqH(!&}
z9E&Pm{*}MI5ML6V>#?(oCDP~aDRe;2=-haJTYgo5g)~`pl3jPVqKwa}2PK^0{Cfd(
z>leyKKiBJz39HTCI>d~hihYDW;cxybcOrUzweY}nCd8A`8CsZ^;p^J45C+}s+M2?q
z+yEwA&i&y&248okRUv}7Er)sY#+XGa9Es0QesE!GurrHe4`)X2T6il6x;3^@@C{`d
z0&6!m|1}{Mg@|#CGItk<=bclQ*e>;3*eLoS0^RzlEr4-S)xgUcXUwHRNZzMXT!d`p
zkV83zK4R9X*m#vr@Vz{l2*REkeZH>(D&w54$w!U)@N;1tg-Pz_8Sw5(4QKO5?gEdw
zCpM<NxBbqS@#HkP_eRo4*sVBcWLOwq0jqNnODFrfYFm|4NeaH`Rp-`A4c5{_N{pV3
zb@G$eWLG(-BFd$+NgFI<&Ks%tUa#d9EDJ>w)~?U;xw6#wzNVG%{NRll_S%!4{O)d6
zX1+w_Y1dg-9gzLrWF6Q>Yb=Sf2HB#jU4J_BGqli3?;-Ha#~MN;kbS9!>KQR11J3q|
z-u<={T4rOXV`#__Twm&<Txx#@+GL;+p|+fTCjdgkUgfTHedlWbm~$3R&;h$;;>u_Y
z1F`}8KHtg%?SK^59=o3dZ$fBWzoJ7Ec~W+Nu17fGgdbq-(h8;cdnPRiae&yof@j2%
zGABF{=fEOiI&9@tC*0V|9zb~Og4i$aOP)0?_`amxnP3Z+tRQ|kOX#WfRS3BgTTvEf
zg&!PJiqHDhn;U<lP82Ptb6&G_^*J?8(&FC^6dqPN3k7_Qy64^qY|6D!o=fYn0TS#n
zGf*j9LrA53!Z?Wcwc_sC=PX!9O<6Hu#rsJz*wNH~aGII`6u1~q;d8Kh`Z|+*A!1qO
z_auYIEHhK7jTenBNfB$62TlEi<vZuKP4UL;D^SO<kHr?cdb_h<EMukS(4Tr)u!~)Z
zZ>pNRsGpd7oo5jsEXB-$O#I3@<LIyZ`ItPe6iec?p0n<8`6Lewmw3_F9NZeH%hRan
zR+*Fq7=G0I5@sESJQdAZgv;UQ{8k<x3T-qiguX%oy@dV!`q?-a*p`Eae?R5Fx0N@Y
zfw+>F$gOj{p_91mXGLnwL-qA=;W*76NuQ=#)~^9c<FgB{LV~sobY%lN2QJ*}=1Y}s
zmE*7H#$pAe4VZf3aCfbT9>&nV%}OWBJIWF&EjlgZ_q8jridDiet7w^^ExDeNAp!L+
zwbhN~FO{Dw$_lLT!$Y+N?q3zFW5EnD7N0GZuwfxa4at6HZ=;5OJ%r>zSfzijgaWmQ
z1Zlb=NsPF<$M-@**AJi0!?a5?xcQVT?Q+)UyWm|{bDuD7Z?rGh^Maos8>@-wwXM$L
z=>pEq1FtQO!yOZDyB*w5*t$H0>ck8noh<kuS@#8Y7RqA#MYs3T2>nJ{BgqA-uOWkK
z86M}40quNES4+iSD^S>4BKXpV+{+i~G8N-TQ#%dTz5nRVjGeCo_2>A`e?8bY7}xV;
z&?7t1NGEPC+Y&&L2%5`AHmQ%oEg1EYCC=1366br)E^qi}*V~@|wtHqXAc!Gj`OECZ
zjXAP5yRnUjPeG>j;G5zxo%hPyd51%!sYJCzD*=K-<L{j)Y^w>6W+rf`?VVU#9wT%v
zqGfR{ZMc$QOqov#yiZf)K}}pTZEHqn*11$MCUV5p?JlhpyQlZ<^inNk5Tx7nO&W3G
zr{`obE*SrA@s_u|2c(l5S_nZ!*lJ7w_ec1c`}*bmzg-_#qiC|0Lh*6Ce>qz{r?jmx
z(+X0(rF1^umrJ^mSo3gWO`GIF{a%pMmK;$CGF^RM6f#PGc3>Ml2UK&obKi(w>}IL;
z#q9Np|K7H8?OFqy+a6+ItE|3i!~E$}RasfYjKsy9O)c>|3Z0-#o6!W=1e0@;%c(An
z^+#BlTR=38^S?($&&|vgTW01Fc3p0orrWMgV(iKH#;(>H<j?S7lDqv_USdo@Y|tGf
zlR~26CTH<s8&HrX0@<*9-k_+H&p4D-T@^o@l%SQHVdd7fdv&zyk|hpt-|VJVVFCQ7
zczs<P3eGkn*G|4D(?@=F+Fo|H>ApZiEwP%%AX+WLiohv3+qhdja`K8Um#hQW{h`E@
zJOKJD-v>-zjezYd#b+=8?&pdE&gYV`OtCeuD?1f576I8(F`*e`X+)pP;_cO8d>w18
zO21ZEjm=#^+532qSG#`R#ks83eyW)8P~bs?-BU0p&i&~|PpbUR_BiA~u{oP}ZODbt
z2X9t#Xwi8=sD4Ui<z0p~@k*#9D+Ir_XhxjfSbIeBpn6}2flm3^D`C8Ip!mnx`mJ0$
zvv}vEHnF2G7FDsVHpacd`D-|U=W4Lm4gsV?(@905v@XCOvlwb_WA5I&J=UF0kh}vg
z+xSF7Z{GJh9q)<`+q|?ICOu>sZEU1sy&4h&LcL0g0;$z&Uk^_&NX&@JFJSeuocPG9
z;>vZh2@h9hy}^94FdO85m^LAo{^Jsd%8ZD1EvNZxU?l|E#X+%8z;3%0*16rDn=T{C
zkgc+!Y=DvqLye+Xk$o*rJ6;If>T=&A>@QfEbE7ZU$r&*kokNqppC(1M<IJ-?(<A%9
z?>&t&*K1ms@^#um^+BClkK?UdZ50kD%?DeQwf2szST{fnB<q1Uaev;Mimh(B(7K3l
z44=fw6c#r0R?{My?qfZFL?7wOAO?bP@vp}Rds&2GQeeiWKq+&~tR-)Q`@`7PS_dz;
z>wuKb18j4xs<VIEWaKf@%RF9j<hjx~0(S3kpYWiVoJSG0#{hpT|8}7MR9#eF<<~yp
z{yV7Dn=~5-r0uu5l_NFHb}zuYBWBI!2o6!u)!~t$iMOcl2D1v!MfKMdDh?d#`x}R1
zR%OUD0XYFM7E!;$>tu9cfWmZKG1C1i@?K}b$8?a>PZ}Oz<zg<Yg%drw;XG8H50P17
z?_UB1Af@Zx>*po%Mu3iREDy|T$>DA)k`j=&7s-YuoKgIT*uy@FIHErpIjD(yrNIg=
zXy%u&E%4!3C5sdCXVKw~dp|!p+ZXtpWdNY#rdf^OT!cke@!GRM0c;swyZ^hk2!jqg
zq!hrF4;)(}Jczsl`L3D%uPv#O06=@~sE}nU%VASwF@%b-yL4cY{;14KRYC$ZE5Xx=
zR!I~oz}5WHP@<ri2S_`>I@W8Bmwg5IK}>Lq*(T!dJ$_3;5`aqY)D~=FqD%omATeSb
zpgLcK9#L6!Z+^SB{(M8(U}H-6`86f-^N~vbtM8HK^uL17J-@%<4{Uto5%A|nivNyj
z>;EFizwH(N`%|}{XF0R~Dra*K-Vi0=0JzciU?IYwvZU?T-hc!p&?v~c9j|K6&rd_~
z>Txk!7<1BG<-Fu7;AVRfFu8qU4c_~0a(j|Yp4+&Kf%QeMlooKGa7e_E(<UjJ&{1Sv
zv~zoeSLHtC{yEmTI>v!F!7jk3O6iEYobTbmpHcu6uzu)UlKv!;TIY25XSYQ5Pgt{t
z7YYWKgjp9%Ay&yOt1}lmPPm^?nFdTV>QQOj=v4s6=3g!7QX6|*X-ZN+-kc!I3gnb5
zeV5capEO6p0ygR+9U($G{B-CtQ^N+dPop(RCPv09gBw7E$zp`Pw~zE@JZbJszEnIG
z|Ax*#(Du>d7Msvz_uM@@)$k?<E%(}o0~m4RO;*}~7v~lAJl1}_rB~nm=9N}=TP{_2
z3cL?=KpvLcoF)m;43%C^ub=3^?sI?-QzjcaY+V2ewq<(~bVJ?b&>^G9^0{!DQzL%^
z(x)+Ylx5=1B*3CSi{xY2o?RB~+HXi}yuwNnw!sMSqCEI?j5Ju!uT61K24ZKxv9|E3
zOj>2b`kvIoiK%qPVy5P{lnuJsz?#66Aj_FDZA+oFYiAF2ipokWF(h~jLUr;tvRaVj
z_r-=-)f!efeLpEwGfAXEgU(6`E>_>~?06VHB&9oOGwPhlykg`QgM_^;^UNd=?)YO*
zQQi=FCcLYr#4aALChrKSjR?o)QGcm7TuW4=ScBKS5_M7UR8xI;ptT<I>+5c3^>o|}
z6dpc7QZFIOEbyeS*)3~#{(7>q_Ks5iQQS~htXq}(Xz3{hrsk){x^XcfbyL1Qb8eM=
zInm<{FKnroYwE_{pC~9_9d*$F7iY1BZ4kSQtzGj5?z-ee?d-Y8`<=gm`6yU>S$|pM
z^X!mMHW1J*Y<a5)&*QVHbw(MQpQ&T)(hoWNTQ0Hg#1&gY+ciz~+oM3R^;ft`PfobC
zeJqqv1_Rd<Wks%SGJP;oV3%XE8hCS9-M;GCDl2154tanvUKkZE%0hfw!pBni=OG5D
zSM&Wa&^HDRA=?Ty48us?tBu>|3w8W`Vj<U8<ObqF!?K!e732i}U0s~PwF7wHvL(^o
z>tLYqL8P`opYLGY)qLKu6$iiINZ4pBC7U9106XZWU;Yf>IY4=GA)N8o;bXnK)O<jb
zJ7TBEVcq;a_+WTmvHK0btitUkhRhu<3_sc1kg>rp8dr7a*|O1;*-6F1t`Vt*-N77m
z`P~gAX}^O?fe2bBZ2C@$T9}KdFETec0Q5}&gsO_@-#*b~=e+=k5qeoVKkSs2_7UVU
zY6tcydAs4Fowv<i6`*Q=DC^p(E_^<PxmtK5G?X#$u_C442^-J=r?t&{^w!<aJ`cCs
ziakVw>JuJ~r;cdVY~D3DJA;4$ye&|^*Wyq1;(|LCN5~o?B+9OR&^*+IgwByoI%WlE
zF6ubwBQ|D0HXLF5;`l1!NDcMfa7BE$pG9hb=-!qiIAiW<Sj|)-yv`q-=o;0lB02`8
z+pkK2<Rk|eO-gHmmN6cZx+G}dYW%)*nAV4b#=i=xgZ1asKa=fm0aG7|-H?45wq=Ks
zio9MQO;V7*fN?+225?p?L2Eg~lUd&HJ^BmU8-ThdK_-#e&NW(f)FU4>q2Mt4P~XAd
zUyeAJ6x%yg4cIXpS^>&vnc1m94Y`e;f#I)PomJGYOz~FRrxT|ecyU-huW1@+Z*XM8
z^rN7OxL0w-hsRTk2d4HpAML_d{YIApl>APVnXNe)9=7s|UBT4P`BTUk`d3^O+MZT>
z=Lf?Iz!;radRZN^a_Y!O1F$&#I4#+EwshLWbfK%Hoqj-nwNgWFN62%eM@o0zW3qqj
zW~d?)V5OY3sjfj1-%tRw8CL+KQht%A0^HAmcHiAMdx^34Z!j{J?Vd?N$GY`}NaDJv
zV`Hr(mLI=*2M>#ISYLfq_Os=-h*)Pt#KVSKy=A;}PLy!ysIN|_-o&YE!@q8K3rH}k
z`^&i(fGPZP&L%DMs;1XnC3#^|+f0Y$ig%J&=DDvKDzFzuT*?=g^-Bq8lo&gz@R4=6
zo|0t5kT5^>p+E+a=~33;or!OGL59X)B^fuwlUFHYXdX6z>wPPnx2E?qH!hDE+qR2G
z4a-UDlXfL!zAf_U&yLmv4Usw{nxl23=D)<fkfCl${!4u>0OK(Mp;gc09ot~-DoEpe
zv$6L>tPS>gaC71n$vTN8QQoG^T^OrMSF8i=FLp(xl|sE|mg1Zl*1#R9TvN_J%}jTr
zVc2+zZ6WIGsi1j&npW9z2nZme++bc|d#~J2P!unfheW1F-E1V{Q~k{$OWL*arb>Fu
zNjnqeW)&iKfg=a8;)F9EVxLS0Dp%FqZ$Z&DR)e#jK71^I$nBU9z~||2VK=O2jL5q&
zJ%j6=sV1L!D+;^8bCehQ-Hpsr#dj6*8;jP<IpDe2ePbP)Pfr(1^2bkW$R{0UifhfH
zLmR{psIwL@hBUcGW%qH#lWmoCzHUi`4_L3DgImwa5s@3BCgzeZ<@d=jjidQqm{1YL
zKw-Itv<dN3n#sDXg~c{yoVr}<cu(yEJBe*=t|e*nQ&~$*2??V^U&rH{@f>|5BAZ5d
zzrQ=yo{~#`B%bDl`@xcxlhB|cRwV8;%Yvo}GKzDUCLdDd0GTLG-w=SJZ(>pV6~AWV
z3?P6*UAa3KR<%RaVWZlAVju46XCS{Ai2x-Hfy3q>#043Np4%3hTLp~pdQtJ}wCdgw
z275$a|K;X@aT)DH3X0;XZVSu)UYLg?i;QVjB{J(+Nb^xK4h3;p_Kt-$$iCMdCDB_D
zDF-t+d_@)!L>8)t>Hfxfi>3{eBq^;sV;A-oF4><>n;m(%{1}(B3&BI`JBRmcwvTL0
ze+dQ<l^17eTDo{0k9PLqS}3tGDP-ab9+9E*auU}9!U7XE%I~G8o=e}{xJq!y;Yo&1
z0wyHOf7LGWm2yZvY>s)SODf1Y3mBGFd~g5$j;w3te>Z8Kw{=UJpk47xsBf3*`pK@5
zHTh8W8z88+`In|Q7RmE%ppr+Ju3$;6GE)n>e)(kkl~8e;m=HUkVcEg1nF8VUm4x=u
z4wc3_Z>u3OM4{^rmHv){%Z|VKnsR77MsAPTK<zXyzDB1mUM!^4ae7-wF$3tCAm<ig
z5W1)}PFiQ*Uf%=4Ds!MKgSdog-wI&sX~@!~y><JXh6%TP<NxXGy~EjT-}rwW=u<6y
zx{MlawW_wJ_Ug20Rcq9UR#7El6C*L6R;xzUNQqh0rdDjJQYt}ek04qj2!dEigzug9
zv%bgoIDUWp{F%BN_jzCUbzSFmzR%Zt!gQtkE^}!_Y@QhfHzAIsHQgqzZ}yb!`E3SD
zMlz%POk}GB0mW>m6!;9AyLh#pTe7_SI<CZ?<XH{GCU;V*IeelE%o<HVa6TWi;Vt&+
zuSvUxi|Q;x*rA2=FwZY%6qu#eCH)jWH$A@P24vs!RR9Y+nx+5}mLwFmNZ3qvZqh>P
zf%jOS=J|mc>H9*n-(mHL-^y-s$KeljP-qi61a}JHkz5uo#VqsXrF?M~bvn27wNo~4
z@`FzognSGhdfE}LvoOp~@p@gCRtLtoutg*<&r`c}=c;pcl7BY8bm{_Jmz5JOxi!Xa
zZ`a0J@v^OA_j>GdjJyNpMkSj>uqB*b(Y)l;^d#F9*aJD7&^n1H+S^14=i0|QX{!^Q
z;Kj$V-cK7tdpwCd`cGIue4P?tjW7BlZ#K8;2WVTWw#hX6SzzcO=`_&hu1uTefGjhn
z!jb$uUO<w~k-^2_1m?Y~JjRCx*vlwh!E=g3bqlO<6_YeEvzXf^@%=t{QvESl1H1XO
z4W1{D5wp0DDL0or^Z;*a3No44Y@PMtj(*TANDq{#EUQ3zd{^}I1Q?HZyGI_8Tsm-Y
zI)^9~sgQttNm6Yuh(7SqIwN<iao2MZ5O-*d)ypQvuG~0!1!KQS4lae-0)xjP>Ti{Y
z#7nrYwe?RY$7?=d?9*<i`nJymawqyEgl{YR<TvMtYH7FD#+v_*SgBo3ZNAPQbu^aq
zFEU>1?v?SB3>lDD`QE1vrePp8#OX=!5q+LLg+n3h;hE;%M>bF?o6%3UV6mPSwi|@A
zC7kT$1;*BR;~yEFoQFB%1Ws(ofVJUNOZL&$2pkkydXDnZ_PL6pquhJb<`*Z#+|b+r
z9h86zHFw;zSPnJmSe&yuF08a3R7rsOV48fY^twrnH}*3>kB#cW_=;rwf3GJ)b)5?$
z77!hf?BhBSUhkxvjXC}v(eglK`!=6e!F5PVe{}4VTW>F=G6rQc{C`W>*V&)vE^ENj
zZi^<CX)!ehb2IA=Kk5qIRzS?v<vRPDRMyldRf7P{DR4seuH~@r9h4tB0@$TlmK`@R
zKL@rOc!LBRH<d4rj1B?4ztvSf_X>GdxO3ds)yy!HB>G%+S~<8!A>>;X^k!2|zW?Yj
z(BaJ@>)7bjLQOK(=z7{}_UjrESXFGOyzz2B(!ZLsofZ0*CCN5>Cm?9b(Q|Bk13+h_
z&zR%Tjgk7DA(jlS2Pv81^Yf{Q(;gZ9wkv$Fz!TeI0W2wLE|H!eYGxk+^I!3<+gJw8
zL|?7%?!@-Lr@u5UW7=G}9P)bB{0m~EC$<YP`vu7DMiBWr(}4ka@Xipbdm4zXU7!5e
zs0EC#K*j=+Hb{x`a;u<o8WTDTINS3~5(f3wf^HLZgB6OF<<?h4nZ~zTh7=6+w60t;
zXl;VI{`BGox_l-#ww3AAQmZA~Hn<GS!jeQ`X}hYSm+wf!Gg&t`Uk!L+91A?B-)p<n
z!8SQVe7kll7Ng_P=2qMGVft=cD!(#_y>y=*m4QLufOJ%&{6)OkQbXN%as8KIgT3T^
zTKjV+!%5PKJi)^h4~aZM<yoA-w#VXEaj?}Q>|jX7DxffY*)5+b-iujx__u_PS1!oa
z*c)JFI?0GM2bRd@S}ZL*h8?K{@uYfwmf9M|%InX1j-&wwKP@kJ_o!|<=CJr5H9>W>
z#C}TIY|xCpvrc_JLl_|bUe%(?5R+kn<4U>C)jg*fNW1-O(?8LiMAf~KH|YG>d))7T
zQ#At_%Wcl_XEUe=nc;JU9__AH3!k48Lv_p^K9X1sVils6w5#F`35}Phzi=ZwNuY=`
zdtTf_-iIUVj8ySY4CKSJHOU+NX`~;oXURH<*laCD5a$G~5{f=|u-0ZcL>(i>qc|En
zepT5mkYdhmIsx@m7?8g%bpYDdRn}IZM6BxZoOdH@?uN%5oF36$7*-+OHzzCvx_Ooz
zF!Z!OX;#YbeDckf8SQUipA+eadKM$X2cWU^C#N4ef`0*n$U`KL(k}-(IkmZK2<S(z
zjZ!PyrY?*R%2r9ruVOh4iym`io7&7m_#*(NsKfqfBU`-7vGZRB+yQZX2=Q2pZiq-W
zLM8D~_iY;gRqD4SV?U-;>L>m-w$-=5MNYXSJw*-ViK#O{rh4(uHOE=kRQ5sPtQKSs
zY12a?U~XosRiR9@0mZw`Q_lYikTxE%2F?(s`~ZdS)RjNd!}N#giJ~Hg!s!D9JC^hi
zBSMeZj^`AOfVB&MnK=;^BvcFZnSPN4c$)t?%lLmFlK=Ua|8I$5&?)Nle%Agrc>nJH
zX2w1s(DD=T*T`3z=YnvuO;Hqv6xfE7wl~~1U-<&xX#*g7+FpfqC<P;|gbzT{jjfa?
zjp&^0=loJI__&qy*=<v}O_LC1bZtB~1|lb(N@WwEr(s6OU=qVc3te7%pcZ}lciTlF
z?)JHclFcPNymZo?2po$|e~{rTkZpeShuE`+s*wj+Su$AJMCW)(_>-=tJtM5zAYt7<
z`_eMsc~SBEzaBYi{7j87zh-8IsDV3mc?nzZP1@y96ha7M9uN~p#yjB(!NT<CF7{ha
zTISGN=mEA6k6nMczSpvn+r%cJ;NC5PF897$6VPV8<Vi9AXQW?z;qydnQ5PFLw0w--
zu&Jg@%G=uJ5Dm^U+RAoZ$^Dn!taa0k$uF{E1z0300}sW3m&j5?iEiR70ab9?9t6;q
z3LO@W*kEw>wze9b-<58EL8c5*rcyHiqdS`)VvHZtET9f)3qDz*HQgG_#m07w<P;vq
zufYdu`97F<l19<dZFr#Su;qgbLoO^Ap-W7YA`;vp_)G37Mb5H6f+KgN>N1`~WJRBM
zL5E-Zf49q-OrJGz(MQdiYa(aE74rhohL^MyP0VnMC*S;T5QtT`xr+~Ncdg5&O3N${
z*PYHlfuF<=TF#DYKEO9_6--Zzx+$%XTs(>y%DoPK1g$?dHnCu#htTSI4m%=_jiui7
znAg1WsOb#DAYLu4ZLPkq0ywCA=<E_Jp4sYYW_$&=Bw*=Kv?#4Cg<A+Bhgms>ZNuap
z34iY{-rDZmw(p(vrb8{=7<Vr-ztF7@S@?|S0elW{M+%hT@M*Jfx>!)X^d+#w^RB?D
z0gNnoo|DS4fM!V85|!ZrbQzvm-7)^=C*xSvBt|F7+eqLYsSS>6je|auI|byI-zX#M
zWe1iNJ8RLpqDifMkUlj`9f`3kSR>;dfZ2<(gnkvDMUf?x;P8&R=ICLuy5XcjJ)?R(
z2NB8$H`}Tbqt|Yxg9LPhDZEA_oii@_^d8y1`}nqWk^lYLA}BeBx~jYJrRidm&<kgx
zB;VDT<t*p}3E`VGv7i2a$6iBhX=R=>!zVLJ&bPm#Ehn0XT6RQOnF<#wWkDZ8=?uak
zXRi_I{-y^Pr&qsqe57Xvvadp_?RcYtyL35c5MWyLMl?E!b<zN}amPh@*+z}{?vihD
z>a=x#8MO5MeV)+qrJiFWfBGlDI#fLyM34oQ_zC%?cF?NdO$=#xY!l*$y^J*BS-BnB
zr|>x7*r7<T>QY~k!Sz7FX+-xgyDA@3s1_^@311J$qKTkL3Lz6>{KH9q<DMiaE}+1{
z+;x2;i^CDz<Mnb5bAuW3i$;8;)S89e>vN$<qDdK4v`TNU`kPGbKc%9vKsU`~H;+1K
z9W#&F!@b*w|6o>qLKIg;SBj9XM9ygb;wo9f2U{L`#nSTa@z6;bh<PpF(3Cf#983~S
zWf>1J+~Ty7I}m#dx-LQl7DaQDvP3a+<PTz@E*lu&7dN(Q`<;uhQ{esbPq*_qv^#x7
zH7mE4A-9K>vk0(z<o!?0U{9HEPBVU6o0vf-=R2fj1a6xlvd=6Y3pck^`@CpDl>4Ma
zuZUU3Kw?j~KM8)9SA&5!B?s-<1O3*Tp!JQ^0JZr512^vhS|WP=0A-!}Dj7Oq=1s#$
z@DxY|R#FiAO68PyzdrMQNwLSS7vi_2(mKl;d%=s$01+z_8C!>T71^91k=AM<Tbp@w
z`kxbr8r}Ue!06?;Pmps8<l3<wh0J?Gv42Nxl-ciZ;q-8Y_C?E8V6H*UOBnP@(GFpI
z-Y-&}HshB-?LL~L{#An+XtitAQ%j@#3gE*bNDeKHj<CQ=@@_0xHb4+HISD^5=M1pL
z_h5ktB{P+%L9r&P8lVzt+mni?WRq1DhE$-`TC)I>-!7T(4pKp>cJMOMI{(D+5>Wkp
z+xSH+>vPoxR41!3r1NdKLK++%b9XGk4IJa(s_88(Jb2FfiAdp1gyxZ6XM{RmGFi8X
zii$iJx^6}0_KwG>jv?~;#Rica!Sm;sEwUfzPbuHS@4-<{t*TqNmO}dpvC1hJ&3w?8
zVha6seym9&n|hwp%ZBy{qu7A>%Z!8?L6<_w&)me@{Hifo9`~MXy$v)Z-3?}FKI}M`
zu~{~nTK}UFTWn&^5Gq|>UHp2i2M{4y<1!XOT(3xh$LGyEcHvF3IzoXp;Ly0FM;)Jp
zgG*nIc|n>+2^(rD_ux7zL_@2029=dJ%J&)?m(2-}i7sZ~7<4IhKYN(QTWh<arHFyQ
zLAarp{kq6=e319a#Bt@|y&p_VOY`(RWs=J5P@L}!neYSRx#YqHNJHjzkk0rym$3{=
zK!{EAg7t#a`MXy@Ar=KdS<_MBd5k(=rRl@Tp$};CE%m9F^*N%_S1_Ck<avfi_VnY#
zvf{Vdw(xha0?8DZ<8L9)G6vfd?gt05as$KkM!hFPYpz8Z4oJN}3ZBrGl<j<_H=x%m
zJ(T(<olh*uehR%X-BKI;jlSjE$8-%qJja_d)IYbo^1P=QWp2Yvu016D9Qqbd8oaSv
zdSH6ET-J>^JOO8?RCTGkNpX4z4!UHu^Z4(;P;eb`4@H-Rs`J1jy!#O5%$n*9_;9fs
z#XoJoj{FVOVK#w0w?FPmpjgai9wUk+)4FMi(!-yb{C*67Qc$R=n~<}QilJYy7pbGH
zs~v-8V*~o?rq~tdcrVlBg2`T$g0iXw2r^RO((o0TLTT82kQ*J`wOd1Yg~?~knQO#%
zUkE(k%-lY9juEzEM`}^S(1``)=!n+woX?B?=$whGnSsR|smVjv0>cCPUkTkIBM30>
z=gd+uRh|ihm)%BK7W|0;tdE^4vih%HsEybAuQv)r4lCI0s&4gVQ>c#@SJIL!JECW&
zYR|sd4T;F8S@5e`*^TjT-aMl1I5J$!4Gs8zZ`O{#3NWTbQ6;TPpD;ixzL0sZ=70Wk
zw=SLcoRa0Kz&6U(R)Wtm*s{f3@v^N3*#RkC{d)dc`of5gf$$4i!ynE?m=Wo%EP>Zr
zY8Ky*eMs_J8&o>QqCMFn6)~NI@r63&#Cb!)y}-I;`!CCB#LG5MlQFga>3QfAjQb$=
zU@Sh%Psruz>eKXLivXHQCwUkpzj6Mtv{BJ8q5lr^+2*7-HLl;lQ-FTFGRUC0_ZlP0
zlsDyGa(U0&_dmzD%Vg@;71CB*ZeUbbf_s$g2Fv{zXZQeB<oJkh3Rt-9Ut@sUNf9Sm
z5$$ri$ti;?%cWE;c$5F!h-&Epf#1)5XSn-)9^n!Qqk`X+KR%JMhE_Y~-h=d5{}JcA
zPDVLRp2NES>%{=}Dk2$G8g)o#C*LRN<-b!aoC3Z_ppj^=1@$(hg7n-AyK5+3=!T3<
z=KAIGC&Kfc)W=3v*Z#aD)4z&14F2GwB%!e2u=@MBfU(&OmaWvC7w%?QHBpSjkt{d3
zqR6jG0DN+g9R1ODmh`1yx;ydXjy8#zqUj*KLCSb4_o(@1aR4&_l*Fqbs;{xyu)cIH
zJdRjryHX&$rC|qmnwz^_RSghY13K%gG!N%Sq!-yGX^O@7xM}%fT9p-B^SUfZcRSr7
z5$3$<Epd@0wq`a6?x9oar3Z>tS0>78><8tunVx+WixxFPWiDBBgIXJ_ip-d~c%n&<
zlABz`wnEVa!G;8$3h3SccNmG3=fmA`W0xc`Dwv8{fqE@_h)6*YPWeA`3u{uoyl0HQ
z-S~xrEpz3rtzCTN+gC;Q!}ue(K6UYthc%$;xeHl#MQ+$@Lw74uWD0D%9&bxdymexa
zP8}lsLh&S71L_4AOeV#8=7NZIgEF?u<@4wI{F;BvF8f7}aj4oT`?$ItPJ70X<}jg{
zLF+zwt|WtLmSm7GeF*gqFYW<sIk-H+SBmWm`Y4}cjHu$3nxwmxR~J9t1JDM0CXGqw
z1J`Of+!sH1^?-ndbkz-0TotA222|)>vVQh!ZLDcCjoXU^wUvoLOI^Fd+n>N#`U)&j
zKx61Oq06ILz9-AkY^)snBuB9m3+);7J7zb(7i9{yPn5j^z-_;N7_RwAc$VSvcCSY&
znlTITG7yy_+}c=9c24i{>-+Dmc2D}`CfchR*bA5TIFwc{CtoCgwNGkw1RbJtpVc-d
z`>drU*K`D%rYwRJB$q6C1RK2WyipJj`VJQKQ>(|`xJzsaTP<%Wycvila^7tTO#=Bn
z=BB8gRqVKtQ;n1Cj-fn+V4xPz_NTJgI`60l_+5R^I$ZU*nCHC#6kxjBD-E|eH2TU5
z1~ik;I&eWY=Ym}yJ0<@$PW?D~+Q;Ot7yc=blAFw?y`;J3#?k}LK2q*IFyIS@8VP_W
zGgL^8U`gFRrwM6_^;`72ROZP?51Iz{JQSC+{S7vx+>B;he-RXQys8ebMD)tHbijxv
zwpu{mdmX?3%Gq2de9NgT;9PEy3Tz4udQwa#h(jtmSSvZmn}w&%bgQ>}c8}odym$&T
zKPn)5AAy{<vGr#4@$il%)^>hprLjMzRPvSlAyx`Mc$YsjmOt!qG;xH4xm@_267FUq
z0YlI^^6r7Jp4wgK51btj`72=SpJOkCuz1~MS+sXOnWXuXw~;qIAbdD+bC?GD@wEQf
z=hMy!E^OkpXMgu<f8hxLe9W8mu(_;;qXqhx;~zdtXb`%uMlY-{QB|JU0bg%4x16P~
zzve(TBtg3VuS<UdXPMNZ5%)u%{)iobbk|uF9Ci{&cZGp;SBq<tLspkvlrvxjd3{cS
z?X)%5@5a9=&Zgb}^Z>Fl=Kr6);s@Y$zrAbKJbrhOeTpsL;=+iGU1}7lN*fNXzki<_
zh|?;{h&SqVcp+3ipVhV%ve(r-pA>o1$=efJ<I%UL(r3f%_`8IUzygK#@5nCHnkQ^m
zD8|pUo#&|H%5}G8NZvjiAwN_{I|vuYAnu`p2KSKDl=-6GjHVc(IWM#4LX{ryMLXTL
zs(-MwN?1)qs~#irEXg9_o_<YPvR6ZN{Mi=Wd98+zP0Mye+Cr4^U6@uP5tC<h5h#ZL
z$XJU~nWG2$^il+Gp#Y&t2yrPlJPtMWlc|(iRpz9VC77>yYfPnvpzNIMSIa+iXcO!o
z&X|PXKJTDE?nTXOvV!OOxx{0Z1#f|mV*_LB;3s>2kCDGxzVpg4>K!of4$@23UtxAr
z?j-$|-1~5Z+g`CBt}M)M1`k|wG%POoryy%Yyn_Rlz$~KyUGDZ7qp@vPuXz&7S#-P#
zM*nx}N>$HE7{6ELD-}MPw(Lzz&}5~X)!yxkryZ1b7MSznE1r>U%>YdYCe^y2-_RK&
z*-b@x9SeFicF`z>WiTM#S6rI#3_(s;aR}pI=q}XC^1>(xhf{$R7hubtO`@0&+;z7|
z>fg+_$Lua<h83G9BhP<~Fy5D@mQ(nDbgoNgaBf~$_o#`vH)Eba?wdi{@YWj|30Gnu
z7}}bPS|hV3NllgbG$3J7nuQHE#IKiUnK}?zNp3~O%t+p{Ygn-+rGq_qV!&w2Pl6uv
zEebt+XJ74;0nV#I0B=N8RWxK&tXn;hh7vvV{Z%o~;}zhv>i)CNV|2Z;Dh;}(W^D80
zRSm`cdb<rwF*Ul^w_P(<<<%17644&lcZT6v*=`+%5Q??n4FjVD7gWK0KHO{c%45?4
zQ$*w{76~$&+X!Y2&p;#LsGv*q_U4Pt#rg|l8mFRMdu|7)C9DV4s<=sjW?Ot)gzt!X
z+tl?vT>9$znE{j|meDa)9f79smsq<kyMr<cS0`er>NaOTp9t@BH{UH;OH8~6Ls8&3
zRSpu%NY@AP^3upz=TV;}AVG}3$gyMYWc##!XTUG9lhT?Nt8;(+^V(vV-|cCluQJ0Y
zr&?V>6nbH?wSG;a9B~9V3{>T#y`Gi|NrJ`>#B$9M2e;x~)Rg#rq^gHrK10?(RURHs
z09SQ><zvc9N>@_II!LfV6HVk?o(%8_ATzE?lgHb~WHB^5w@19|eYbXAn|4eOev1Vx
z66$*{XshT(xv(|JuE`KTp>0JZz37^(cWFY5N89oJJ8sEM$@#@%=vKUBi6jEFa}6TN
z<1F~&NePE0)h5$%=XK~D%&m`nvmjMv!*cg2i!#X<0i4yv#lAi2xYHB$P-zK%XdGST
zc(%_^pclqSMrZ=<{O>cY!beAM<{7zpkDm?4lmwj*PC|cJNWi-t-G<e786sozYZCQv
zBT!4m)c62o$jv*AN!!UOaaOA^Oj6azj)BnzJ5c{yn^sC&%e_Wu<D71s8n|m2Nb?*s
zw%j%v{8y!0HDW#Bt!$N4?soM@1F1dRq?hWvMiZTi9;_<47rtS6-Y&<GablKHF1;s>
zpM5F*)?$y)drz-~BoCZ(-uy}dOh)Ooj#LGqTW9R_OSg<N5kSxr!Bu)M<H2V8xP@$s
z;TpixAbK#*b1|%U-?OLYeGD2j2J4c)MZBU#F>!=XS2mLf?mlZLhiUIiHL*j9Ll}*X
zWrKP`xC|(^{92g|5cH`UmNb$>ezxYI@Vppe(B~L>m!l!l{Bc@DM@eI`#aP6POv_!M
z@x66BeC4*(RiEUHigSaD*PQZz;n;w634ZlnSp4y95&WAke*EF=d34F9Q>E-9t*4o=
z&wbtAmhGgcH#*0ldW4WbAZskb)W+F9tm$MH2H>eJ`S!N(?hrKqo5knu)a7U9$mQRQ
zKt-Y|dZ=M0ybjAs*Hh7^?R5neKFc-Xe+fonpomuS?Bx%^*xXUeFWMyeeL|ADUs^}1
z>X7p(km>l&eWFRq!Y>={Gq>lqYXD0`L4`zb`%O%-raozqLKpKFwgMj07KUTwkC`T7
z6Vn9zX*z_TOXhi-s<D@e^?ccs_>4_)5a7tT(n<Xg2OOVJn>WqAE@Z3rT`N5hOa}}#
zr?kw3KOunFg%E$S165)55dlfQ1$_`8xP?@fz{MVm<z70v&{#HZt$uTM)aQ`ZaJm^X
z%Zujk^8@2%0e|R-Y?W#7h{Io$)HKC`L{(=&kR66})$v-Ny>FQ#D&Zdk!l!IH#&>-u
z)HmDIWp@`p7M295K=KM9m_Hb-{LvGSf^iR!mYxu}Z_(H1YIC1tcuvk^O7;kGq#o)A
z(+<bE?=05~lwME1>S8V=m;d$T?0v=O@R#5TF(1pSyy+PLn<&S~Zx1Ni3Qid=e^i|)
zyDtFPzuIiNJwbyh%c)Gy9{Wz<1g`!jLz7HQy+#%A`ioh5;P&9R@53-+faHQQ`XmN8
z(C4|md3j}bCpwYFWZ|54-HvF#pn5G$D*mh3%fLUAMd#;CUjs$9b%zD=XXk=SWLmLH
zmvU8Whn&01uZt&!!%gmL#$#0~3p?ON2(q9%3RTGJ-WV<bP&U>;Ag)0l1bN`O%Ifo9
zNP?TXi<Tc&FPo2yD2a3bj>5L5iZ)>ZAfl<nCi@g0gmDkTI~oL~L}z#*@8iK~XnNC`
zCOUAdE_Kc2z<mOpoAvXBeq1$7x|;}DeGZ_c`M@-SCtPJtmC2g}rS<&!AN!?n=-qM=
zZ8dVq^oag0HmgBH4wt19I4>;eMP{v-aJ9pt3O>{C9@LgWac^3ZU~~PIvWGr6fgS<Q
zDb<)Yvqj^E0sQeoty8WSVl2CVh-1yAuNLP);uB+QN0XqFkCf1x$&Z4!ih#$yS&{k5
z(#)|gImjSzi=Q>o2rY=^cEEm`@;X3faV}~O5QYM-w5{$W3mt5*xKJ@3Aptg08e`oy
zf7FI*EFgA)PWJ0OZQ@q|7!iXFf=9zi=8Oiyj6YMhDE39M5}u4gZmrMGAkBHB=nK8@
zCbZWUP9;{QeV#Gnq=D^n>Kfb+I#()V`wLjJJX|jeIr9&PJfWY&H!4LX!T^B@KqQ(l
z;1osjDxg~7A-<=|s#JM?lZFNUVJYb+HwZe*>aS66#G89z7y$gr=EbDt@$DvtpLLM*
zsFdZ5SeE^4SuiVM%R}BWC`!^Vyt%JEw{3qn;!6PNvyZF9`3ssj;fpKVTNc!?SHPjM
z6Kpzd0BqRngwbcPnK!xKUNIkl?45^`cf3zrVx*dMH12LS=#AQTZ%3-vJGKHayr9L8
z#Z9c&q&OI}*glO~Qc}c4jB&=%n?GU5QkBVdy=NKCfS=>oKc89o<W^u!nk_(brhWj5
zx*Yqz`5z8CsAmb`*Mi{D#7VJ+sKNaTx_HlKNzK0_w^yUX*c{Uw0;X8)KU1uxed!)+
zik1K7lb#2R6kw7VyPeV8R-_@<j%}>}Jx8nnyE=p!?<NO9prf)bsBcld_VQ`h5`Hc*
z<(~<W)8ZV#?7>o&0^7IiFMCa=TQ_8AEu^6rcH)8QWPPGatD!?>^%2k2d)eZQ^!XW0
zIE9G@mf%6ct?~2o2mI2Cxm@n!3hd2PHtO_x5!r);<{zZS)dDVKU^(Kp$9Tml^0e0~
z$989T{S(xl{FZq;1WH%{4mCA_#XOZb)xg4otyj)9;%@SZVqOx*#WbRZ_PGc1%ov&+
z#!d3;xT3do+ORHZxQOW)6I~8o0dC%(c*O(TwaVvJzUZG*KUJk?rZzqA;o7}KjrKY|
zwSCzZsDMg(E&hH_8iu*GkFp>QclN(#Nb<Bd_N#Kj)}Ii}Nqq&~^7%`$xjI4F{hwe-
zbK2wod=Gs>=+X_;%XR^Fty}h2V!%lbNyERhz{*$aiMZ#be_t*|CBCY+NalBe(AFko
zy}mG=Sc%?+ot}ETp>l-EMuQ%|y<^URC2!aIrkh$9bERjtA`##tFV9Oj;S~g}oRXbI
zP^Tt-MGT4wERN<a16B?J8K_j%^kk*@NX<(nTMm&LINf#W2$OXeS~sMx<2jRJ8vMe5
zNP278oN7-k<I>*8Ib(os+ohTgAQ;bqV8nF)0%=l9r^^>_3D8uD1s&xIfHO&$Um#$`
z12DmWVEozO3<CX{&0*~ahW^>W2{*dkQFguw*k`>WLUX*R{zxV7b=ccYRBd8y0p3lK
zG5$GAmd5<C-4xZ^V-JiZ0F836<@+0<3X-yd>`$~uU)Za2o*%v_lRg3E#&Eufb!yGr
zTj{hFx{8Z*Whp9Cmd2Gy3H*CHjirVM-e2~xxYdTtn4gf64Xb39o?WpN^AaOWZ{`jC
zW4x2e^{?MZ`$JB4?o5vIUA@S%)>Bew*yCmkCX&DW@e%s5NICjH+^<+5!C;$WPjI~N
z<IV&G^{k&SYODlrioG^0#>g7M^XJ1=ZWb59|3GP`{i$IhYrz118V|9g`nDmG^CNAW
z+kZY|qp$&*{1D!$qpFG?|NCE7;9T`00PxI<g9ol&?oH&zC94HP0b(&z>)@DjO{`Nn
z=0QSEu3g8ynQ5P0YnvHIuy;H}1~5z8PcXW{W2v~E9jLZprn;|MyC?kTtEPl@6I;H|
zt|hVaQgpVr#$64D?~Lzit$2_1R05SMAfh~UV~V%+X=h{uXyWM_OjvG;Lz?qQ2oNRj
zw+nHc#xFN+t}WjoCyO@85h5gNA3Bg)mz0Q;a96-f@0U(_+4VQ`^<d?G?t4i4@E%os
z)$q##azYar@zzUf3((fO%x7-LpNNvOw0)-WruoKxV6ageM&z{XqI1K@a@AWh<#+xJ
z)NNY8`~Zeb;`zv#8>*A0bmN0#LkD4!^=-t2Np3n(`hXmNRjT>^<ZP$i%u=Rf-y;n1
zwQQB=?!=l*P2TV7rwMV-$t$bY853hqi158JloYw1e(}7@Q^48elh?!Kn}x^EMU9(?
zht~n_bUn*jM_RtWFIh4u_O0M?V7y*f^65d=jXnCmADj>&7f8pQamVtKNA>{5@jE6%
zw`Rte#8f3}f<?p=DkapK^H#0pJbfc8DEa^_QcG(cuT|p|Sz-6o5b3Z<Q9S)U(;P0y
z%}QEsZBx_DR4r8S+ccv!hvYnIM@zAPGn8+sSQ#SjwI1vKIqw-HVKG<ww$+I9qR1ak
z@^shZ)ZmE<OXM1?a5U%G5$&XBExS9n3>d^AL%o<{2+Z@pW9Wk%pELWL885!@e2Xo~
zaJ4~)#%d=8xT!?HP~;S^WXUn%MM$^di6@EDcw;E|8_~R^EWul<<Dmg$qD8klVN@Ph
zdXfz#?dM%8)nv>P)T8?L=I9O92a3I<_InhpfMqi9@Qv>6(VJI6xNEA9sF_%hm2c2z
zEkEgn=6+|>!=}fHp7W=oHVWS+7DLyPy*{~C4qn9`e0x<?D^EA*9>5JD2M$CvY;AqQ
z5|63|0s?5S$_spT-@~?lQ6C2FqXly&(j<MPw)o*^LK%j>aO{?Rwoic(ke!ob`B$Ml
zDbs&*#EQFXzQ}AaT3lajvnpMlj6nU*FakVw#D9bq6n>DAOjATflk~GMwV+4JppT8^
zu5s0#4z1i=C_QNMAZ0Q6{-S@1qsAQQq)<zdE6r~&?M&{MOzyMweShVCnYNzM;e0H-
z4FX)+H_BCiC_JdVLG_et$>FxCR56lUj~o7$ev8H&v`q4J%sDe2aQN=$d3P79zLNCb
z0y%6buyHBj+I0OAAXlB1M){}oQ*g~_!Dw!*yE)w&kS8=rRqAlg%;HrdfC)~N&+_m^
zQPnbBOH+`dvysDArlyw2oAkR@Ch8||{;J}9Vv>4JA<_5Ou$T|b=RbtL%olV($(*aK
zobQ0CyqxBJ<GU110hf<qk&uPhc6M@>ExWH(Iw3sz1Z;3s%g6^^+%Gzdu{9qLiZ1#u
z1?ML_o7K<XB4EAvx=)TvxmPl~v@&FSPIuxZTfV!b{#b6ztFGQq)lzhU<zKgg_m-T^
zf&G|BmsY$l_n;dcbiWpqu|QXVo!qQfBa(@^;aaddF~E`oc}UNS=6QhzTY$Yu^>jZR
z`GEE{EhHw9ADLQz-{)Mf131{kguxZ2_))r(vYN(EOJt`MPi16<`>a_a2X_t2cl`dY
z8ZTtC0m`jAN|eS`Ev2HC?-x&8?cpWAG|pvWEz@&#I0dgp-FA|Dpg(;cPX*}CA_td}
zn^Lzj(<y~plFj{93YKR(fbwWQuPrKJ!0Yi_wgydGe0}M)4od5PXJ9)Zd;);}m;2~!
z<;MzCfJUn^hscYT=wIwtCB$3xtb!)b0DAv>mN)E6le*wnPJdUIeaNQE?l`bk#V!ic
z(+!k_N<<5_fxKx6VR1$Rec0UmbtB@|HdicxS30ew0*<Sy2~BC5OSSRSa==*|7=$~x
zFJxw1DLAC?qCGVb4XBW%&?RI7)_OKm%S!@tC&HAsr@s2l)~|q@Cvdw7>^IY5Ill{Y
z-9Ox1;gp1>bW@*W!AY{)jX!__lv^ais^$B)VI$`Lv^F?~(g|)eJ^s4k1I0tI;RVmc
z0Z*quIog>FAL%lu8a<mUy?S_#!F}V>#LM&P74(f6df=bMA;eEESgzu<vS;stekDpO
zck0L$$QmqA&a)(g7nZ37IwOd6Qgs03Z9K3M@_&BxT~H~{={AYi>q2Q8BuALD!<!fF
zHkyC5z)qEqMx09=>MU2Gn|cIOc~GwgC{`rOnmp>X$z0gA>H|y%UCOR!8z+{h4Skcb
zqIkOR0N#U7C&SY4ZqtJ=|I&8ssFFUI@{KFZ`lgZapbTm&o0@~O;ppiTKMA=n%Wi8s
zVAyF(dfBk7GTYJ=K`-;te)a^w{*85g=Pv}XfAm>x#k4EKv}d+G14SD1Okd6;Ou{ow
za@Fl>;=->no&TBG!+E*2(}f6!0sF~6_xpgopaRny%1hlQ7p4B4pHx$r>?~Bdy(bSE
zF#KY``K1v5s7o2^MAe$tnsqWNu2Rz=NJkrQ;I$kxfVRk_D?(BR3y9%#zwcbyD|;kF
zA*D7$5F|bL{!wC?TAH6hTr6UJSu8Ur?3UPx5t%X(ZI$bEB)8OdG1r->sQHbuUmuQ9
zcMbgl*RL}k%nmEJJ|HP^$zgpjhyQvp_r%j?7V>B|G|*Og%r;&<Kj^H($AgV}wSIYn
zd$i&AzIh=lwDqn1kNb4Qex6|40L(M-36M`O-=qw~$qgt3>dwDugqaus`<;?$b|-I)
zkn#5Nr|oQOflgqY_wZ=DIiuAVkan*rDbfR=Q)!`6D0mE5n5|!y|Ek8mac4%8Q(~-T
zdWyEw*kgRV3wKf-`14cSizmSnvRiToK9s7Krwh$L)^^FiygZ1+iv#8`clRE`$+J^m
zuSRGGya_gOYOEh|gox$3gAt5%39bi9)mOBrQWla}Qc*KPRI>7VPVkV-T~x|KMce_b
z)cjVaon*f!{w-jrqi`F<qnnhpRPzo>=6(cAwAJrF8FhVGV!Q}ZIo7KA&}a1Yhidn+
zoSbxNtHvB{rbSDjbF`6gyzGDXU*W9VM*6?5%0)uI*IX+?hCY*R`vIaVJ0GVy{#t^X
zB|6QFw&*h7pE$*qU)a1hvQ;-ehlf^$%X~Z54uZ@W+*)eXH>2fG>CyrL@bN+BD_5s6
zM#1S!8jGaDe#a(MKNgVa#FPfq!A@QmDZ8$8v+??;sZ8WA>~w$26-MC0RimJ!CaD~O
zOVZN`;kW?Aki+?HldfOW1218?>eNkGg{@nK=05@a6>6(BCp)fg5APjBB|g4r9e7he
z`RHry)8F?~a`!DiFW;;|x%S^ZAMIzm!J|DE%bCx9Ny$9&*SWN9z{1aVX5;F*^;kqs
zc--&Q_2~N^R_eM4i5L`i{1!E_Bu)J!1^7XVH+^~4FLrq|TYJ;tu_j*C2iHu9|DWa4
zfJ1a&>ds@K8~Ku$fBs1vvM3;t2m6WNd=>0_4xgwA7(yNUY~$OoTj$9``{a76t2gv?
z-ppgThjzc;6~9P{&=&RD;Wm9WLP`yH!)CB!4*vbaKt)vYRdeM1&(%&$$XG`2aHZh7
zUGi-KD<3iW#ak~l(XV_Ru}TKhaQVn*lZ$E6HNLJ&T@U;U{Ap3u>XYdOI8|Ao(?tKR
z#!IpDu<C>FrG2v2?;qx{!e{9do>X&Q$XHl>penmd$jJ%&=s4uXBwF`ST&Mm!$*rce
z=sL6VDE2e4BLNrmiB)4y5xL?89LR^~FIWwYZ>;n1>YTp}xIf-o$t@C5TbqErTfAxF
zo;lkk#5bOn-G?fe^+<EL|4>H0G{>iDKs30`Ta6NcF+!b->Pq@)ONMEd2;pLdoV|i^
zS!dY{?vAP1bQ^XbHMXf;M~SV|K6F>kEAlFwDu*~eLmwNJ53gyLQb9XT>pAknOS%Ix
zcmeCkT=kvg2SHcIm~(s&dq9>VG3v_UaqBOY%vGhxkO(hZ-9KsFqqD~uj#ZsfJI~oe
zj5N?%^;wq7jK}}8+c)Hk3KDry_bG5x`b(zSBBUQ;u(y8{IPdQAQX<+YOpUbvcL@ZN
z5M*7->}9a^^f<TN35B9^*k9Y=LoMGQd@lc=jroDFxZ?v}a}Vg4ed-Mt=B?;hyuvbm
zKVpxHP?D`<x1*6n=sNCh4Agl(OD<*f|DDi$TfZM%niJGxJwlJV@w*qfC}Uynzt4t+
n&i%sW&Mt)d?@eqnv-U-L%hbMlbC+bZRw})_#&=3@*@gWdqe8`a

literal 0
HcmV?d00001

diff --git a/docker/setup_env.sh b/docker/setup_env.sh
new file mode 100644
index 000000000..e7b1ecc77
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/docs/Contribution_Guidelines.md b/docs/Contribution_Guidelines.md
new file mode 100644
index 000000000..d7f163a3d
--- /dev/null
+++ b/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. 
diff --git a/docs/Report-v1.0.0-cn.md b/docs/Report-v1.0.0-cn.md
new file mode 100644
index 000000000..3ddc23e0b
--- /dev/null
+++ b/docs/Report-v1.0.0-cn.md
@@ -0,0 +1,135 @@
+# 技术报告 v1.0.0
+
+在2024年3月，我们推出了Open-Sora-Plan，一个旨在复现OpenAI [Sora](https://openai.com/sora)的开源计划。它作为一个基础的开源框架，能够训练视频生成模型包括无条件视频生成，类别引导视频生成，文生视频。
+
+**今天，我们兴奋地展示Open-Sora-Plan v1.0.0，极大地改进视频生成质量、文本控制能力。**
+
+相比于之前的视频生成模型，Open-Sora-Plan v1.0.0 有以下的改进：
+
+1. **CausalVideoVAE高效的训练与推理**。 我们用4×8×8的对视频进行时间和空间的压缩。
+2. **图片视频联合训练提升视觉质量**。 CasualVideoVAE 将首帧看作图片，天然支持同时编码图片和视频。这允许扩散模型提取更多时空细节来改善质量。
+
+
+### Open-Source Release
+我们开源了Open-Sora-Plan去促进视频生成社区的进一步发展。公开代码、数据、模型。
+- 在线演示：Hugging Face [![hf_space](https://img.shields.io/badge/🤗-Open%20In%20Spaces-blue.svg)](https://huggingface.co/spaces/LanguageBind/Open-Sora-Plan-v1.0.0)， [![Replicate demo and cloud API](https://replicate.com/camenduru/open-sora-plan-512x512/badge)](https://replicate.com/camenduru/open-sora-plan-512x512) 和 [![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), 感谢[@camenduru](https://github.com/camenduru)大力支持我们的工作！🤝
+- 代码：所有训练脚本和采样代码。
+- 模型：包括扩散模型和CausalVideoVAE [这里](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.0.0)。
+- 数据：所有原视频和对应描述 [这里](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.0.0)。
+  
+## 效果
+
+Open-Sora-Plan v1.0.0支持图片视频联合训练。我们在此展示视频和图片的重建以及生成：
+
+720×1280**视频重建**。 因为github的限制，原视频放在: [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
+
+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%"/>
+
+65×1024×1024**文生视频**
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/2641a8aa-66ac-4cda-8279-86b2e6a6e011
+
+65×512×512**文生视频** 
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/37e3107e-56b3-4b09-8920-fa1d8d144b9e
+
+
+512×512**文生视频** 
+
+![download](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/491d72bc-e762-48ff-bdcc-cc69350f56d6)
+
+## 详细技术报告
+
+### CausalVideoVAE
+
+#### 模型结构
+
+![image](https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/e3c8b35d-a217-4d96-b2e9-5c248a2859c8)
+
+因果VAE架构继承了[Stable-Diffusion Image VAE](https://github.com/CompVis/stable-diffusion/tree/main)。 为了保证图片VAE的预训练权重可以无缝运用到视频VAE中，模型结构采取如下设计:
+
+1. **CausalConv3D**: 将Conv2D 转变成CausalConv3D可以实现图片和视频的联合训练. CausalConv3D 对第一帧进行特殊处理，因为它无法访问后续帧。对于更多细节，请参考https://github.com/PKU-YuanGroup/Open-Sora-Plan/pull/145
+
+2. **初始化**：将Conv2D扩展到Conv3D常用的[方法](https://github.com/hassony2/inflated_convnets_pytorch/blob/master/src/inflate.py#L5)有两种：平均初始化和中心初始化。 但我们采用了特定的初始化方法（尾部初始化）。 这种初始化方法确保模型无需任何训练就能够直接重建图像，甚至视频。
+   
+#### 训练细节
+
+<img width="833" alt="image" src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/9ffb6dc4-23f6-4274-a066-bbebc7522a14">
+
+我们展示了 17×256×256 下两种不同初始化方法的损失曲线。黄色曲线代表使用尾部初始化的损失，而蓝色曲线对应中心初始化的损失。 如图所示，尾部初始化在损失曲线上表现出更好的性能。 此外，我们发现中心初始化会导致错误累积，导致长时间内崩溃。
+
+#### 推理技巧
+尽管训练Diffusion中VAE始终是冻住的，我们仍然无法负担CasualVideoVAE的花销。在我们的实验中, 80G的显存只能够在半精度下推理一个256×512×512或32×1024×1024的视频 ，这限制了我们扩展到更长更高清的视频。因此我们采用tile convolution，能够以几乎恒定的内存推理任意时长或任意分辨率的视频。
+
+### 数据构建
+我们定义高质量的视频数据集包括两个核心法则：(1) 没有与内容无关的水印。(2) 高质量的文本注释。
+
+**对于法则1**，我们从开源网站（CC0协议）爬取了大约40k videos：1234个来自[mixkit](https://mixkit.co/)，7408个来自[pexels](https://www.pexels.com/)，31616个来自[pixabay](https://pixabay.com/)。我们根据[Panda70M](https://github.com/snap-research/Panda-70M/blob/main/splitting/README.md)提供的场景变换剪切script将这些视频切成大约434k video clips。事实上，根据我们的剪切结果，从这些网上上爬取的99%的视频都是单一的场景。另外，我们发现爬取的数据中超过60%为风景相关视频。更多细节可以在[这](https://github.com/PKU-YuanGroup/Open-Sora-Dataset)找到。
+
+**对于法则2**，很难有大量的高质量的文本注释能够从网上直接爬取。因此我们用成熟的图片标注模型来获取高质量的稠密描述。我们对2个多模态大模型进行消融实验：[ShareGPT4V-Captioner-7B](https://github.com/InternLM/InternLM-XComposer/blob/main/projects/ShareGPT4V/README.md) 和 [LLaVA-1.6-34B](https://github.com/haotian-liu/LLaVA)。前者是专门用来制作文本注释的模型，而后者是一个通用的多模态大模型。经过我们的消融实验，他们在caption的表现差不多。然而他们的推理速度在A800上差距很大：40s/it of batch size of 12 for ShareGPT4V-Captioner-7B，15s/it of batch size of 1 for ShareGPT4V-Captioner-7B。我们开源所有的[文本注释和原视频](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.0.0)。
+
+| 模型名字 | 平均长度 | 最大值 | 标准差 |
+|---|---|---|---|
+| ShareGPT4V-Captioner-7B | 170.0827524529121 |  467 | 53.689967539537776 | 
+| LLaVA-1.6-34B | 141.75851073472666 |  472 | 48.52492072346965 | 
+
+### 训练扩散模型
+与之前的工作类似，我们采用多阶段的级联的训练方法，总共消耗了2048个A800 GPU 小时。我们发现联合图片训练能够显著加速模型的收敛并且增强视觉观感，这与[Latte](https://github.com/Vchitect/Latte)一致。以下是我们的训练花销。
+
+| 名字 | Stage 1 | Stage 2 | Stage 3 | Stage 4 |
+|---|---|---|---|---|
+| 训练视频尺寸 | 17×256×256 |  65×256×256 | 65×512×512 |  65×1024×1024 | 
+| 计算资源 (#A800 GPU x #小时) | 32 × 40 |  32 × 18 |  32 × 6 |  训练中 | 
+| 权重 | [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) |  训练中 | 
+| 日志 | [wandb](https://api.wandb.ai/links/linbin/p6n3evym) |  [wandb](https://api.wandb.ai/links/linbin/t2g53sew) |  [wandb](https://api.wandb.ai/links/linbin/uomr0xzb) | 训练中 | 
+| 训练数据 | ~40k videos |  ~40k videos |  ~40k videos |  ~40k videos | 
+
+## 下版本预览
+### CausalVideoVAE
+目前我们发布的CausalVideoVAE v1.0.0版本存在2个主要的缺陷：**运动模糊**以及**网格效应**。我们对CasualVideoVAE做了一系列的改进使它推理成本更低且性能更强大，我们暂时叫它为预览版本，将在下个版本发布。
+
+**1分钟720×1280视频重建**。 受限于GitHub，我们将原视频放在这：[原视频](https://streamable.com/u4onbb)，[重建视频](https://streamable.com/qt8ncc)。
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/88202804/cdcfa9a3-4de0-42d4-94c0-0669710e407b
+
+我们从kinetic 400的验证集中随机选取100个样本进行评估，结果表如下所示：
+
+|  | SSIM↑ | LPIPS↓ | PSNR↑ | FLOLPIPS↓ |
+|---|---|---|---|---|
+| v1.0.0 | 0.829 |  0.106 |  27.171 |  0.119 | 
+| Preview | 0.877 |  0.064 |  29.695 |  0.070 | 
+
+#### 运动模糊
+
+| **v1.0.0** | **预览版本** |
+| --- | --- |
+| ![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)  |
+
+#### 网格效应
+
+| **v1.0.0** | **预览版本** |
+| --- | --- |
+| ![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)  |
+
+### 数据构建
+
+**数据源**：正如上文提到，我们的数据集中超过60%为风景视频。这意味着我们的开域视频生成能力有限。然而当前的大规模开源数据集大多从YouTube爬取，尽管视频的数量多，但我们担忧视频本身的质量是否达标。因此，我们将继续收集高质量的数据集，同时也欢迎开源社区的推荐。 
+
+**Caption生成流程**：当我们训练时长增加时，我们不得不考虑更有效的视频caption生成方法，而不是多模态图片大模型。我们正在开发一个新的视频注释生成管线，它能够很好的支持长视频，敬请期待。
+
+### 训练扩散模型
+尽管目前v1.0.0展现了可喜的结果，但我们仍然离Sora有一段距离。在接下来的工作中，我们主要围绕这三个方面:
+
+1. **动态分辨率与时长的训练**: 我们的目标是开发出能够以不同分辨率和持续时间训练模型的技术，使训练过程更加灵活、适应性更强。
+
+2. **更长的视频生成**: 我们将探索扩展模型生成能力的方法，使其能够制作更长的视频，超越目前的限制。
+
+3. **更多条件控制**: 我们力求增强模型的条件控制能力，为用户提供更多的选项和对生成视频的控制能力。
+
+另外，通过仔细观察生成的视频，我们发现存在一些不符合常理的斑点或异常的流动，这是由于CasualVideoVAE的性能不足导致的 如上面提到。在未来的实验中，我们将使用更强的VAE，重新训练一个扩散模型。
diff --git a/docs/Report-v1.0.0.md b/docs/Report-v1.0.0.md
new file mode 100644
index 000000000..8b8419976
--- /dev/null
+++ b/docs/Report-v1.0.0.md
@@ -0,0 +1,136 @@
+# 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 are made publicly available.
+- Demo: Hugging Face demo [![hf_space](https://img.shields.io/badge/🤗-Open%20In%20Spaces-blue.svg)](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,234 videos from [mixkit](https://mixkit.co/), 7,408 videos from [pexels](https://www.pexels.com/), and 31,616 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. More details can be found [here](https://github.com/PKU-YuanGroup/Open-Sora-Dataset).
+
+**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). We show some statistics here, and we set the maximum length of the model to 300, which covers almost 99% of the samples.
+
+| Name | Avg length | Max | Std |
+|---|---|---|---|
+| ShareGPT4V-Captioner-7B | 170.0827524529121 |  467 | 53.689967539537776 | 
+| LLaVA-1.6-34B | 141.75851073472666 |  472 | 48.52492072346965 | 
+
+### Training Diffusion Model
+Similar to previous work, we employ a multi-stage cascaded training approach, which consumes a total of 2,528 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 × 22 |  32 × 17 |  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. We are launching an Open-Sora-Dataset project, check out the details at [Open-Sora-Dataset](https://github.com/PKU-YuanGroup/Open-Sora-Dataset)
+
+**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.
diff --git a/docs/Report-v1.1.0.md b/docs/Report-v1.1.0.md
new file mode 100644
index 000000000..ca3ec2ba5
--- /dev/null
+++ b/docs/Report-v1.1.0.md
@@ -0,0 +1,256 @@
+# Report v1.1.0
+
+In April 2024, we launched Open-Sora-Plan v1.0.0, featuring a simple and efficient design along with remarkable performance in text-to-video generation. It has already been adopted as a foundational model in numerous research projects, including its data and model.
+
+**Today, we are excited to present Open-Sora-Plan v1.1.0, which significantly improves video generation quality and duration.**
+
+Compared to the previous version, Open-Sora-Plan v1.1.0, the improvements include:
+
+1. **Better compressed visual representations**. We optimized the CausalVideoVAE architecture, which now has stronger performance and higher inference efficiency.
+2. **Generate higher quality, longer videos**. We used higher quality visual data and captions by [ShareGPT4Video](https://sharegpt4video.github.io/), enabling the model to better understand the workings of the world.
+
+Along with performance improvements, Open-Sora-Plan v1.1.0 maintains the minimalist design and data efficiency of v1.0.0. Remarkably, we found that v1.1.0 exhibits similar performance to the Sora base model, indicating that our version's evolution aligns with the scaling law demonstrated by Sora.
+
+### 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.1.0).
+- Code: All training scripts and sample scripts.
+- Model: Both Diffusion Model and CasualVideoVAE [here](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.1.0).
+- Data: Both raw videos and captions [here](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0).
+
+## Gallery
+
+
+### 221×512×512 Text-to-Video Generation
+
+| 221×512×512 (9.2s) | 221×512×512 (9.2s) | 221×512×512 (9.2s) | 221×512×512 (9.2s) |
+| --- | --- | --- | --- |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/6d18f344-f7da-44eb-9e07-77813f6b5e90" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/71f75e72-e9ee-4ce7-b8ea-d2d45a6f367e" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/80430eae-a3b4-4f24-b448-0db2919327d6" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/dc217894-a8c8-4174-a42c-acd2811f61f5" width=224> |
+| This close-up shot of a Victoria crowned pigeon showcases its striking blue plumage ... |  a cat wearing sunglasses and working as a lifeguard at pool. |  Photorealistic closeup video of two pirate ships battling each other as they sail ... | A movie trailer featuring the adventures ofthe 30 year old spacemanwearing a redwool ...  |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/f9c5823f-aa03-40ee-8335-684684f5c842" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/f207d798-8988-45b0-b836-347e499ee000" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/d40c38dc-9f26-4591-8163-c7089e1553e3" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/67b40ce1-135d-4ae2-a1ee-5a0866461eb2" width=224> |
+| A snowy forest landscape with a dirt road running through it. The road is flanked by ... |  Drone shot along the Hawaii jungle coastline, sunny day. Kayaks in the water. | Alpacas wearing knit wool sweaters, graffiti background, sunglasses.  | The camera rotates around a large stack of vintage televisions all showing different ...  |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/42542b4e-b1b8-49b8-ada4-7bcfde8e1453" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/4eee2619-a5ca-4a32-b350-e65d6220c8f7" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/e6152c93-4edf-4569-8d17-1effd87a7780" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/ec78d855-07c7-4421-896a-3c46a83ec129" width=224> |
+| A drone camera circles around a beautiful historic church built on a rocky outcropping ... | Aerial view of Santorini during the blue hour, showcasing the stunning architecture ...  |  A robot dog explores the surface of Mars, kicking up red dust as it investigates  ... | An aerial shot of a lighthouse standing tall on a rocky cliff, its beacon cutting ...  |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/98897215-eae9-49f3-8fdb-df1d4b74d435" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/10e0f93d-925f-4b38-8205-7d89f49195f1" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/5453cf37-29ac-423d-9fb2-05f23416ca3e" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/77928824-6705-4d83-a7b4-43fdb74790bf" width=224> |
+| 3D animation of a small, round, fluffy creature with big, expressive eyes explores ... |  A corgi vlogging itself in tropical Maui. |  A single drop of liquid metal falls from a floating orb, landing on a mirror-like ... | The video presents an abstract composition centered around a hexagonal shape adorned ...  |
+
+### 65×512×512 Text-to-Video Generation
+
+| 65×512×512 (2.7s) | 65×512×512 (2.7s) | 65×512×512 (2.7s) | 65×512×512 (2.7s) |
+| --- | --- | --- | --- |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/a0601f20-579c-4e2e-832c-5763546718cc" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/55229eca-de3a-476b-930b-13a35eb5db30" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/74b9e7a1-0fa4-4f0d-8faf-0c84d97b11b5" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/66b06822-2652-453e-80fb-dd2988b730ce" width=224> |
+| Extreme close-up of chicken and green pepper kebabs grilling on a barbeque with flames. | 3D animation of a small, round, fluffy creature with big, expressive eyes explores a ...   |  A corgi vlogging itself in tropical Maui. |  In a studio, there is a painting depicting a ship sailing through the rough sea. |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/4d27ff13-e725-4602-bf17-90df2c0d8005" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/049fd4db-f2fe-4633-ab62-1dda8268e090" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/aba01259-60f2-49ef-aa33-e738dc8c9a49" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/98178397-61b0-4200-9ab9-9bef2728ed98" width=224> |
+| A robot dog trots down a deserted alley at night, its metallic paws clinking softly ... | A solitary spider weaves its web in a quiet corner. The web shimmers and glows with ...  |  A lone surfer rides a massive wave, skillfully maneuvering through the surf. The water ... |  A solitary cheetah sprints across the savannah, its powerful muscles propelling it ... |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/3964dfcd-d1b4-406b-916c-d4a702184a27" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/2a5fd1c0-9304-46e4-af35-5ffcc718bf08" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/73e1304a-6241-4e19-9dde-f2b1032edefd" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/c82b822b-cadd-4f11-bcc3-29e3651c02c0" width=224> |
+| A solitary astronaut plants a flag on an alien planet covered in crystal formations ... |  At dawn's first light, a spaceship slowly exits the edge of the galaxy against a ...|  A dapper puppy in a miniature suit, basking in the afternoon sun, adjusting his tie ... |  A wise old elephant painting abstract art with its trunk, each stroke a burst of color ... |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/8edfa249-272c-4773-8728-12686527771e" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/33338d8e-5ea7-4b57-9e94-97f3ee404033" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/227fa6b6-801b-438b-9b30-cd3a4e0a7f2f" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/b5d5ebee-ab13-4e32-8b79-e6cb8f08d4a0" width=224> |
+| In an ornate, historical hall, a massive tidal wave peaks and begins to crash. Two ... | A Shiba Inu dog wearing a beret and black turtleneck.  | A painting of a boat on water comes to life, with waves crashing and the boat becoming ...  | Many spotted jellyfish pulsating under water. Their bodies are transparent and glowing ...  |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/fe8e9450-5a80-435d-b050-6b2fe11cdf53" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/7d9335df-817d-479d-9ea7-615cb50c66b8" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/34c1dda6-e420-4edd-bbcd-38ff8e542ec6" width=224> |  |
+| An animated hedgehog with distinctive spiky hair and large eyes is seen exploring a ... | An animated rabbit in a playful pink snowboarding outfit is carving its way down a ...  | A person clad in a space suit with a helmet and equipped with a chest light and arm ...  |   |
+
+### 65×512×512 Video Editing
+
+| generated 65×512×512 (2.7s) | edited 65×512×512 (2.7s) |
+| --- | --- |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/edb8e8c2-5eef-4c90-85fb-6adb035067c3" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/32d93845-8904-4f8f-832e-37eba2ceb542" width=224>  |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/a70d0d91-0d61-4aa4-9520-6e4a6c477f12" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/7913e06f-1e0b-4d06-8233-72c882c6abfe" width=224>  |
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/0e614ec0-fba0-4f42-a343-d4607966dd40" width=224> | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/c531d930-410c-4614-890d-bae8013f33c2" width=224>  |
+
+### 512×512 Text-to-Image Generation
+
+ <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/e44b7f8a-5da2-49c2-87c4-52ea680ad43b" width=512> 
+
+## Detailed Technical Report
+
+### CasualVideoVAE
+
+#### Model Structure
+
+As the number of frames increases, the encoder overhead of CausalVideoVAE gradually rises. When training with 257 frames, 80GB of VRAM is insufficient for the VAE to encode the video. Therefore, we reduced the number of CausalConv3D layers, retaining only the last two stages of CausalConv3D in the encoder. This change significantly lowers the overhead while maintaining nearly the same performance. Note that we only modified the encoder; the decoder still retains all CausalConv3D layers, as training the Diffusion Model does not require the decoder.
+
+<img width="722" alt="vaemodel" src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/b6f8a638-cfb2-40f3-94be-45af8dbad18e">
+
+We compare the computational overhead of the two versions by testing the forward inference of the encoder on the H100.
+
+| Version | 129×256×256 |   | 257×256×256 | | 513×256×256 | |
+|---|---|---|---|---|---|---|
+|  |  Peak Mem. |  Speed  | Peak Mem. |  Speed  |Peak Mem. |  Speed  |
+| v1.0.0 |  22G |   2.9 it/s  | OOM |  -   | OOM |   -  |
+| v1.1.0 |  18G |  4.9 it/s   | 34G  |  2.5 it/s   | 61G |   1.2 it/s   |
+
+
+#### Temporal Module
+
+<img width="480" alt="vaemodel" src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/7c8a4263-b7d1-4edc-a60c-801ae9b4344f">
+
+In v1.0.0, our temporal module had only TemporalAvgPool. TemporalAvgPool leads to the loss of high-frequency information in the video, such as details and edges. To address this issue, we improved this module in v1.1.0. As shown in the figure below, we introduced convolution and added learnable weights, allowing different branches to decouple different features. When we omit CausalConv3D, the video is reconstructed very blurry. Similarly, when we omit TemporalAvgPool, the video becomes very sharp.
+
+|  | SSIM↑ | LPIPS↓ | PSNR↑ |
+|---|---|---|---|
+| Base | 0.850 |  0.091 |  28.047 |
+| + Frames | 0.868 |  0.070 |  28.829 | 
+| + Reset mixed factor | 0.873 |  0.070 |  29.140 | 
+
+
+
+
+
+#### Training Details
+
+Similar to v1.0.0, we initialized from the Latent Diffusion's VAE and used tail initialization. For CausalVideoVAE, we trained for 100k steps in the first stage with a video shape of 9×256×256. Subsequently, we increased the frame count from 9 to 25 and found that this significantly improved the model's performance. It is important to clarify that we enabled the mixed factor during both the first and second stages, with a value of a (sigmoid(mixed factor)) reaching 0.88 at the end of training, indicating the model's tendency to retain low-frequency information. In the third stage, we reinitialized the mixed factor to 0.5 (sigmoid(0.5) = 0.6225), which further enhanced the model's capabilities.
+
+#### Loss Function
+
+We found that using GAN loss helps retain high-frequency information and alleviates grid artifacts. Additionally, we observed that switching from 2D GAN to 3D GAN provides further improvements.
+
+| GAN Loss/Step | SSIM↑ | LPIPS↓ | PSNR↑ |
+|---|---|---|---|
+| 2D/80k | 0.879 |  0.068 |  29.480 |
+| 3D/80k | 0.882 |  0.067 |  29.890 | 
+
+#### Inference Tricks
+Therefore, we introduced a method called **temporal rollback tiled convolution**, a tiling approach specifically designed for CausalVideoVAE. Specifically, all windows except the first one discard the first frame because the first frame in a window is treated as an image, while the remaining frames should be treated as video frames.
+
+<img width="633" alt="tiled_temp" src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/0a06011e-1d6c-410a-9f1c-82c4122a018a">
+
+We tested the speed on the H100 with a window size of 65×256×256.
+
+| Version | 129×256×256 |   | 257×256×256 | | 513×256×256 | |
+|---|---|---|---|---|---|---|
+|  |  Peak Mem. |  Speed  | Peak Mem. |  Speed  |Peak Mem. |  Speed  |
+| 4×8×8 |  10G |   1.3 s/it  | 10G | 2.6 s/it   | 10G |   5.3 s/it  |
+
+### Data Construction
+Since Open-Sora-Plan supports joint training of images and videos, our data collection is divided into two parts: images and videos. Images do not need to originate from videos; they are independent datasets. We spent approximately 32×240 H100 hours generating image and video captions, and all of this is **open source**!
+
+#### Image-Text Collection Pipeline
+We obtained 11 million image-text pairs from [Pixart-Alpha](https://huggingface.co/datasets/PixArt-alpha/SAM-LLaVA-Captions10M), with captions generated by [LLaVA](https://github.com/haotian-liu/LLaVA). Additionally, we utilized the high-quality OCR dataset [Anytext-3M](https://github.com/tyxsspa/AnyText), which pairs each image with corresponding OCR characters. However, these captions were insufficient to describe the entire image, so we used [InternVL-1.5](https://github.com/OpenGVLab/InternVL) for supplementary descriptions. Since T5 only supports English, we filtered for English data, which constitutes about half of the complete dataset. Furthermore, we selected high-quality images from [Laion-5B](https://laion.ai/blog/laion-5b/) to enhance human-like generation quality. The selection criteria included high resolution, high aesthetic scores, and watermark-free images containing people.
+
+Here, we are open-sourcing the prompt used for InternVL-1.5:
+```
+# for anytext-3m
+Combine this rough caption: "{}", analyze the image in a comprehensive and detailed manner. "{}" can be recognized in the image.
+# for human-160k
+Analyze the image in a comprehensive and detailed manner.
+```
+
+| Name | Image Source | Text Captioner | Num pair |
+|---|---|---|---|
+| SAM-11M | [SAM](https://ai.meta.com/datasets/segment-anything/) |  [LLaVA](https://github.com/haotian-liu/LLaVA) |  11,185,255 |
+| Anytext-3M-en | [Anytext](https://github.com/tyxsspa/AnyText) |  [InternVL-1.5](https://github.com/OpenGVLab/InternVL) |  1,886,137 | 
+| Human-160k | [Laion](https://laion.ai/blog/laion-5b/) |  [InternVL-1.5](https://github.com/OpenGVLab/InternVL) |  162,094 | 
+
+
+#### Video-Text Collection Pipeline
+In v1.0.0, we sampled one frame from each video to generate captions. However, as video length increased, a single frame could not adequately describe the entire video's content or temporal movements. Therefore, we used a video captioner to generate captions for the entire video clip. Specifically, we used [ShareGPT4Video](https://sharegpt4video.github.io/), which effectively covers temporal information and describes the entire video content. The v1.1.0 video dataset comprises approximately 3k hours, compared to only 300 hours in v1.0.0. As before, we have open-sourced all text annotations and videos (both under the CC0 license), which can be found [here](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main).
+
+
+
+| Name | Hours | Num frames | Num pair |
+|---|---|---|---|
+| [Mixkit](https://mixkit.co/) | 42.0h |  65 |  54,735 |
+|   |  |  513 |  1,997 | 
+| [Pixabay](https://pixabay.com/) | 353.3h |  65 | 601,513 |
+|   |  |  513 |  51,483 | 
+| [Pexel](https://www.pexels.com/) | 2561.9h |  65 |  3,832,666 |
+|   |  |  513 |  271,782 | 
+
+### Training Diffusion Model
+Similar to our previous work, we employed a multi-stage cascaded training method. Below is our training card:
+
+#### Stage 1
+
+Surprisingly, we initially believed that the performance of the diffusion model would improve with longer training. However, by observing the [logs](https://api.wandb.ai/links/linbin/o76j03j4), we found that videos generated at 50k steps were of higher quality than those at 70-100k steps. In fact, extensive sampling revealed that checkpoints at 40-60k steps outperformed those at 80-100k steps. Quantitatively, 50k steps correspond to approximately 2 epochs of training. It is currently unclear whether this is due to overfitting from a small dataset or the limited capacity of the 2+1D model.
+
+#### Stage 2
+
+In the second stage, we used Huawei Ascend computing power for training. This stage's training and inference were fully supported by Huawei. We conducted sequence parallel training and inference on a large-scale cluster, distributing one sample across eight ranks. Models trained on Huawei Ascend can also be loaded into GPUs and generate videos of the same quality.
+
+
+#### Stage 3
+
+In the third stage, we further increased the frame count to 513 frames, approximately 21 seconds at 24 FPS. However, this stage presents several challenges, such as ensuring temporal consistency in the 2+1D model over long durations and whether the current amount of data is sufficient. We are still training the model for this stage and continuously monitoring its progress.
+
+| Name | Stage 1 | Stage 2 | Stage 3 |
+|---|---|---|---|
+| Training Video Size | 65×512×512 |  221×512×512 | 513×512×512 |
+| Compute (#Num x #Hours) | 80 H100 × 72 | 512 Ascend × 72 |  Under Training |
+| Checkpoint | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/65x512x512) | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/221x512x512) | Under Training |
+| Log | [wandb](https://api.wandb.ai/links/linbin/o76j03j4) | - |  - |
+| Training Data | ~3k hours videos + 13M images |  |  |
+
+### Video Editing
+
+The recently proposed [ReVideo](https://mc-e.github.io/project/ReVideo/) achieves accurate video editing by modifying the first frame and applying motion control within the edited area. Although it achieves excellent video editing performance, the editing length is limited by the base model [SVD](https://github.com/Stability-AI/generative-models). Open-Sora, as a fundamental model for long-video generation, can compensate for this issue. Currently, we are collaborating with the ReVideo team to use Open-Sora as the base model for long video editing. Some preliminary results are shown [here](). 
+
+The initial version still needs improvement in several aspects. In the future, we will continue to explore integration with ReVideo to develop improved long-video editing models.
+
+## Failed Case and Discussion
+
+Despite the promising results of v1.1.0, there remains a gap between our model and Sora. Here, we present some failure cases and discuss them.
+
+### CasualVideoVAE
+
+Despite the significant performance improvement of VAE in v1.1.0 over v1.0.0, we still encounter failures in challenging cases, such as sand dunes and leaves. The video on the left shows the reconstructed video downsampled by a factor of 4 in time, while the video on the right is downsampled by a factor of 2. Both exhibit jitter when reconstructing fine-grained features. This indicates that reducing temporal downsampling alone cannot fully resolve the jitter issue.
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/1a87d6d8-4bf1-4b4e-83bb-84870c5c3a11
+
+https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/1a87d6d8-4bf1-4b4e-83bb-84870c5c3a11
+
+### Diffusion Model
+
+#### Semantic distortion
+
+On the left is a video generated by v1.1.0 showing a puppy in the snow. In this video, the puppy's head exhibits semantic distortion, indicating that the model struggles to correctly identify which head belongs to which dog. On the right is a video generated by Sora's [base model](https://openai.com/index/video-generation-models-as-world-simulators/). We observe that Sora's early base model also experienced semantic distortion issues. This suggests that we may achieve better results by scaling up the model and increasing the amount of training data.
+
+Prompt：A litter of golden retriever puppies playing in the snow.Their heads pop out of the snow, covered in.
+
+| Our | Sora Base×1 | Sora Base×4 | Sora Base×32 |
+|---|---|---|---|
+| <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/1d456168-afad-4e22-ae3b-fc28eca935e8" width=224>  |<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/c4ca99d9-9492-45c8-a75e-6efe21c330aa" width=224>  |<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/b7b52894-f58b-4e64-858b-015247108b8b" width=224>  | <img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/dcaf793b-1da4-4cc1-9c17-8a46d55e80e6" width=224> |
+
+#### Limited dynamics
+
+The primary difference between videos and images lies in their dynamic nature, where objects undergo a series of changes across consecutive frames. However, the videos generated by v1.1.0 still contain many instances of limited dynamics. Upon reviewing a large number of training videos, we found that while web-crawled videos have high visual quality, they are often filled with meaningless close-up shots. These close-ups typically show minimal movement or are even static. On the left, we present a generated video of a bird, while on the right is a training video we found, which is almost static. There are many similar videos in the dataset from stock footage sites.
+
+Prompt：This close-up shot of a Victoria crowned pigeon showcases its striking blue plumage and red chest. Its crest is made of delicate, lacy feathers, while its eye is a striking red color. The bird's head is tilted slightly to the side,giving the impression of it looking regal and majestic. The background is blurred,drawing attention to the bird's striking appearance.
+
+
+| Our | Raw video |
+|---|---|
+|<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/7ffb6bc6-b52c-488e-9f29-d7d90bda44d6" width=224>  |<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/95fbbca4-e206-42f6-8c6b-af1063c442c6" width=224>  | 
+
+#### Negative prompt
+
+We found that using negative prompts can significantly improve video quality, even though we did not explicitly tag the training data with different labels. On the left is a video sampled using a negative prompt, while on the right is a video generated without a negative prompt. This suggests that we may need to incorporate more prior knowledge into the training data. For example, when a video has a watermark, we should note "watermark" in the corresponding caption. When a video's bitrate is too low, we should add more tags to distinguish it from high-quality videos, such as "low quality" or "blurry." We believe that explicitly injecting these priors can help the model differentiate between the vast amounts of pretraining data (low quality) and the smaller amounts of fine-tuning data (high quality), thereby generating higher quality videos.
+
+Prompt：A litter of golden retriever puppies playing in the snow.Their heads pop out of the snow, covered in.
+Negative Prompt：distorted, discontinuous, ugly, blurry, low resolution, motionless, static, low quality
+
+
+| With Negative Prompt | Without Negative Prompt |
+|---|---|
+|<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/1d456168-afad-4e22-ae3b-fc28eca935e8" width=224>  |<img src="https://github.com/PKU-YuanGroup/Open-Sora-Plan/assets/62638829/7ad17f96-bfab-455a-830d-0daebccaf6fb" width=224>  | 
+
+## Future Work
+
+In our future work, we will focus on two main areas: (1) data scaling and (2) model design. Once we have a robust baseline model, we will extend it to handle variable durations and conditional control models.
+
+### Data Scaling
+
+#### Data source
+
+As mentioned earlier, our dataset is entirely sourced from stock footage websites. Although these videos are of high quality, many consist of close-up shots of specific areas, resulting in slow motion in the videos. We believe this is one of the main reasons for the limited dynamics observed. Therefore, we will continue to collect datasets from diverse sources to address this issue.
+
+#### Data volume
+
+In v1.1.0, our dataset comprises only ~3k hours of video. We are actively collecting more data and anticipate that the video dataset for the next version will reach ~100k hours. We welcome recommendations from the open-source community for additional datasets.
+
+### Model Design
+
+#### CasualVideoVAE
+In our internal testing, even without downsampling in time, we found that it is not possible to completely resolve the jitter issue in reconstructing fine-grained features. Therefore, we need to reconsider how to mitigate video jitter to the greatest extent possible while simultaneously supporting both images and videos. We will introduce a more powerful CasualVideoVAE in the next version.
+
+#### Diffusion Model
+In v1.1.0, we found that 2+1D models can generate higher-quality videos in short durations. However, for long videos, they tend to exhibit discontinuities and inconsistencies. Therefore, we will explore more possibilities in model architecture to address this issue.
diff --git a/docs/Report-v1.2.0.md b/docs/Report-v1.2.0.md
new file mode 100644
index 000000000..8a122211a
--- /dev/null
+++ b/docs/Report-v1.2.0.md
@@ -0,0 +1,193 @@
+# Report v1.2.0
+
+In May 2024, we launched Open-Sora-Plan v1.1.0, featuring a 2+1D model architecture that could be quickly utilized for exploratory training in text-to-video generation tasks. However, when handling dense visual tokens, the 2+1D architecture could not simultaneously process spatial and temporal dimensions. Therefore, we transitioned to **a 3D full attention architecture**, which better captures the joint spatial-temporal features. Although this version is experimental, it advances video generation architecture to a new realm, leading us to release it as v1.2.0.
+
+Compared to previous video generation models, Open-Sora-Plan v1.2.0 offers the following improvements:
+
+1. **Better compressed visual representations**. We optimized the structure of CausalVideoVAE, which now delivers enhanced performance and higher inference efficiency.
+2. **Better video generation architecture**. Instead of 2+1D, we use a diffusion model with a 3D full attention architecture, which provides a better understanding of the world.
+
+
+### Open-Source Release
+We open-source the Open-Sora-Plan to facilitate future development of Video Generation in the community. Code, data, model are made publicly available.
+- Code: All training scripts and sample scripts.
+- Model: Both Diffusion Model and CasualVideoVAE [here](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0).
+- Data: Filtered data [here](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0).
+
+## Gallery
+
+93×1280×720 Text-to-Video Generation. The video quality has been compressed for playback on GitHub.
+
+<table class="center">
+<tr>
+  <td><video src="https://github.com/user-attachments/assets/1c84bc92-d585-46c9-ae7c-e5f79cefea88" autoplay></td>
+</tr>
+</table>
+
+## Detailed Technical Report
+
+### CausalVideoVAE
+
+#### Model Structure
+
+The VAE in version 1.2.0 maintains the overall architecture of the previous version but merges the temporal and spatial downsampling layers. In version 1.1.0, we performed spatial downsampling (stride=1,2,2) followed by temporal downsampling (stride=2,1,1). In version 1.2.0, we conduct both spatial and temporal downsampling simultaneously (stride=2,2,2) and perform spatial-temporal upsampling in the decoder (interpolate_factor=2,2,2).
+
+Due to the absence of additional convolutions during downsampling and upsampling, this method more seamlessly inherits the weights from the SD2.1 VAE, leading to improved initialization of our VAE.
+
+<img src="https://s21.ax1x.com/2024/07/24/pkHrHx0.png" width=768>
+
+
+#### Training Details
+
+
+As with v1.1.0, we initialize from the [SD2.1 VAE](https://huggingface.co/stabilityai/sd-vae-ft-mse) using tail initialization. We perform the first phase of training on the Kinetic400 video dataset, then use the EMA weights from this phase to initialize the second phase, which is fine-tuned on high-quality data (collected in v1.1.0). All training is conducted on 25-frame 256×256 videos using **one A100 node**.
+
+
+| Training stage | Dataset | Training steps  |
+|---|---|---|
+| 1 |  K400 |  200,000 |
+| 2 |  collected in v1.1.0 |  450,000   |
+
+
+#### Evaluation
+
+We evaluated our VAE on the validation sets of two video datasets: [Webvid](https://github.com/m-bain/webvid) and [Panda70m](https://github.com/snap-research/Panda-70M/), and compared it with our [v1.1.0](https://github.com/PKU-YuanGroup/Open-Sora-Plan/blob/main/docs/Report-v1.1.0.md), [SD2.1 VAE](https://huggingface.co/stabilityai/sd-vae-ft-mse), [CV-VAE](https://github.com/AILab-CVC/CV-VAE), and [Open-Sora's VAE](https://github.com/hpcaitech/Open-Sora). The Webvid validation set contains 5k videos, while the Panda70m validation set has 6k videos. The videos were resized to 256 pixels on the short side, center-cropped to 256x256, and then 33 consecutive frames were extracted. We used PSNR, SSIM, and LPIPS metrics, and measured the encoding speed on an A100 GPU. The specific results are as follows:
+
+
+**WebVid**
+
+| Model | Compress Ratio |PNSR↑ | SSIM↑ |LPIPS↓ |
+|---|---|---|---|---|
+| SD2-1 VAE | 1x8x8 | 30.19 | 0.8379 | 0.0568 |
+| SVD VAE | 1x8x8 |<ins>31.15</ins> |<ins>0.8686</ins> | **0.0547** | 
+| CV-VAE | 4x8x8 | 30.76 | 0.8566 | 0.0803 |
+| Open-Sora VAE | 4x8x8 | 31.12 | 0.8569 | 0.1003 |
+|  Open-Sora Plan v1.1 | 4x8x8 | 30.26 | 0.8597 |<ins>0.0551</ins> |
+|  Open-Sora Plan v1.2  | 4x8x8| **31.16** | **0.8694** | 0.0586 |
+
+**Panda70M**
+
+| Model | Compress Ratio| PNSR↑ | SSIM↑ |LPIPS↓ |
+|---|---|---|---|---|
+| SD2-1 VAE | 1x8x8 |30.40 | 0.8894 | 0.0396 |
+| SVD VAE | 1x8x8 |<ins>31.00</ins> | **0.9058** | **0.0379** | 
+| CV-VAE  | 4x8x8| 29.57 | 0.8795 | 0.0673 |
+| Open-Sora VAE | 4x8x8 | **31.06** | 0.8969 | 0.0666 |
+| Open-Sora Plan v1.1 | 4x8x8 | 29.16 | 0.8844 | 0.0481 |
+|  Open-Sora Plan v1.2  | 4x8x8| 30.49 |<ins>0.8970</ins> |<ins>0.0454</ins>|
+
+**Encode Time on A100**
+
+|Input Size| CV-VAE | Open-Sora | Open-Sora Plan v1.1 | Open-Sora Plan v1.2 |
+|---|---|---|---|---|
+| 33x256x256 | 0.186 | 0.147 |<ins>0.104</ins> | **0.102** |
+| 81x256x256 | 0.465 | 0.357 |<ins>0.243</ins> | **0.242** |
+
+### Training Text-to-Video Diffusion Model
+
+#### Model Structure
+
+The most significant change is that we **replaced all 2+1D Transformer blocks with 3D full attention blocks**. Each video is first processed by a patch embedding layer, which downsamples the spatial dimensions by a factor of 2. The video is then flattened into a one-dimensional sequence across the frame, width, and height dimensions. We replaced [T5-XXL](https://huggingface.co/DeepFloyd/t5-v1_1-xxl) with [mT5-XXL](https://huggingface.co/google/mt5-xxl) to enhance multilingual adaptation. Additionally, we incorporated RoPE.
+
+
+### Sequence Parallelism
+
+Due to the high computational complexity of 3D full attention, we must allocate a video across 2 GPUs for parallel processing when training with long-duration and high-resolution videos. We can control the number of GPUs used for a video sample by adjusting the batch size on a node. For example, with `sp_size=8` and `train_sp_batch_size=4`, 2 GPUs are used for a single sample. **We support sequence parallelism for both training and inference**.
+
+**Training on 93×720p**, we report speed on H100.
+
+| GPU （sp_size） | batch size | Enable sp | Train_sp_batch_size | Speed | Step per day |
+|---|---|---|---|---|---|
+|8|8|×|-|100s/step|~850|
+|8|-|√|4|53s/step|~1600|
+|8|-|√|2|27s/step|~3200|
+
+**Inference on 93×720p**, we report speed on H100.
+
+| Size | 1 GPU | 8 GPUs | 
+|---|---|---|
+|29×720p|420s/100step|80s/100step|
+|93×720p|3400s/100step|450s/100step|
+
+#### Dynamic training
+
+Deep neural networks are typically trained using batched inputs. For efficient hardware processing, batch shapes are fixed, leading to a fixed data size. This requires either cropping or padding images to a uniform size, both of which have drawbacks: cropping degrades performance, while padding is inefficient and results in significant information loss. Generally, there are three methods for training with arbitrary token counts: Patch n' Pack, bucket, and pad-mask.
+
+
+
+**Patch n' Pack** ([NaViT](https://arxiv.org/abs/2307.06304)): bypasses the fixed sequence length limitation by combining tokens from multiple samples into a new sample. This approach allows variable-resolution images while maintaining aspect ratios by packaging multiple samples together, thereby reducing training time and enhancing performance and flexibility. However, this method involves significant code modifications and requires re-adaptation when exploring different model architectures in fields with unstable model designs.
+
+
+**Bucket** ([Pixart-alpha](https://arxiv.org/abs/2310.00426), [Open-Sora](https://github.com/hpcaitech/Open-Sora)): This method packages data of different resolutions into buckets, sampling batches from each bucket to ensure same resolution within each batch. It requires minimal code modifications to the model, mainly adjusting the data sampling strategy.
+
+**Pad-mask** ([FiT](https://arxiv.org/abs/2402.12376), our v1.0/v1.1): This method sets a maximum resolution and pads all data to this resolution, generating a corresponding mask. Although the approach is straightforward, it is computationally inefficient.
+
+We believe that current video generation models are still in an exploratory phase. Extensive modifications to model code during this period can incur unnecessary development costs. The pad-mask method, while straightforward, is computationally inefficient and can waste resources in video, which involves dense computations. Ultimately, we chose the bucket strategy, which requires no modifications to the model code. Next, we will explain how our bucket strategy supports arbitrary lengths and resolutions. For simplicity, we will use video duration as an example:
+
+<img src="https://s21.ax1x.com/2024/07/24/pkHr4aQ.png" width=768>
+
+
+We define a megabatch as the total data processed in a single step across all GPUs. A megabatch can be divided into multiple batches, with each batch corresponding to the data processed by a single GPU.
+
+**Sort by frame**: The first step is to count the number of frames in all video data and sort them. This step aims to group similar data together, with sorting being one method to achieve this.
+
+**Group megabatch**: Next, all data is divided into groups, each forming a megabatch. Since all data is pre-sorted, most videos within a megabatch have the same number of frames. However, there will always be boundary cases, such as having both 61-frame and 1-frame videos in a single megabatch.
+
+**Re-organize megabatch**: We re-organize these special megabatches, which actually constitute a small proportion. We randomly replace the minority data in the megabatch with the majority data, thus re-organizing it into a megabatch with same frame counts.
+
+**Shuffle megabatch**: To ensure data randomness, we shuffle both within each megabatch and between different megabatches.
+
+When supporting dynamic resolutions, we simply replace each sample's frame sequence with (frame × height × width). This method ensures that the data dimension processed by each GPU in every step is the same, preventing situations where GPU1 waits for GPU0 to finish processing a longer video. Moreover, it is entirely decoupled from the model code, serving as a plug-and-play video sampling strategy.
+
+
+#### Training stage
+
+Similar to previous work, we use a multi-stage training approach. With the 3D DiT architecture, all parameters can be transferred from images to videos without loss. To explore training costs, all parameters of the diffusion model are trained from scratch. Therefore, we first train an text-to-image model, using the training strategy from [Pixart-alpha](https://arxiv.org/abs/2310.00426).
+
+The video model is initialized with weights from a 480p image model. We first train 480p videos with 29 frames. Next, we adapt the weights to 720p resolution, training on approximately 6 million higher-quality (HQ) samples from Panda70M, filtered for aesthetic quality and motion. Finally, we refine the model with a more higher-quality (HQ) subset of 1 million samples. After that, we use a filtered data (collected in v1.1.0) for fine-tuning 93-frame 720p videos. Below is our training card. We release the annotation file [here](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/anno_json).
+
+| Name | Stage 1 | Stage 2 | Stage 3 | Stage 4 |Stage 5 |
+|---|---|---|---|---|---|
+| Training Video Size | 1×320×240 |  1×640×480 | 29×640×480 |  29×1280×720 | 93×1280×720 |
+| Training Step| 146k |  200k | 30k | 21k | 3k |
+| Compute (#Num x #Hours) | 32 Ascend × 81 | 32 Ascend × 142 |  128 Ascend × 38 | 256 H100 × 64 | 256 H100 × 84 |
+| Checkpoint | - | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/1x480p) | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/29x480p) | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/29x720p) | [HF](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/93x720p) |
+| Log | - | - | [wandb](https://api.wandb.ai/links/1471742727-Huawei/trdu2kba) | [wandb](https://api.wandb.ai/links/linbin/vvxvcd7s) | [wandb](https://api.wandb.ai/links/linbin/easg3qkl)
+| Training Data | [10M SAM](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0/blob/main/anno_json/sam_image_11185255_resolution.json) | 5M internal image data | [6M HQ Panda70M](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0/blob/main/anno_json/Panda70M_HQ6M.json) | [6M HQ Panda70M](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0/blob/main/anno_json/Panda70M_HQ6M.json) | [1M HQ Panda70M](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0/blob/main/anno_json/Panda70M_HQ1M.json) and [100k HQ data](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/anno_json) (collected in v1.1.0) |
+
+Additionally, we fine-tuned 3.5k steps from the final 93×720p to get [93×480p](https://huggingface.co/LanguageBind/Open-Sora-Plan-v1.2.0/tree/main/93x480p) for community research use.
+
+### Training Image-to-Video Diffusion Model
+
+#### Model Structure
+
+<img src="https://s21.ax1x.com/2024/08/12/pApZZJf.png">
+
+To reuse the weights of the Text-to-Video model, our Image-to-Video model is inspired by the Stable Diffusion Inpainting Model and adopts a strategy based on frame-level inpainting. By incorporating three types of information—original noise, masked video, and mask—under different control frame conditions, our model can generate coherent videos while ensuring flexibility in its usage.
+
+Compared to the denoiser structure of the Text-to-Video model, the Inpainting Model's denoiser has only changed the number of channels in the `conv in` layer. To ensure the model has a good prior knowledge, we introduce the masked video and mask information through zero initialization. We believe this is due to the 2+1D structure's lack of ability to establish long-range information dependencies, and relying solely on attention in the temporal dimension makes it difficult to capture information changes under frame control. In Text-to-Video tasks, this phenomenon is not as evident because all frames share the same text prompt embedding. However, in Image-to-Video tasks, simply concatenating images in the channel dimension does not ensure the model can accurately capture changes between frames. This is because the model cannot directly replicate image information from the channels to reduce the loss, and the 2+1D structure's interaction solely on the temporal axis fails to allow the model to discern which information from the control frames can be utilized, especially there are significant differences between frames. Therefore, without a shared image-semantic information, the control frame information might not be effectively conveyed to each frame.
+
+##### About Semantic Adapter
+
+In previous models based on the Unet 2+1D architecture, it is necessary to input the control frames into the CLIP model to obtain semantic embeddings. These semantic embeddings are then injected into the denoiser through cross-attention. The structure that extracts CLIP embeddings and injects them into the denoiser is commonly referred to as a semantic adapter.
+
+In the 2+1D architecture, the semantic adapter is commonly present. Additionally, papers like [DynamiCrafter](https://arxiv.org/abs/2310.12190) have pointed out that incorporating the semantic adapter helps maintain stability in the generated videos. We believe this is because the 2+1D structure lacks the ability to establish long-range information dependencies, and relying solely on attention in the temporal dimension makes it difficult to capture information changes under frame control. In the Text-to-Video task, this phenomenon is not as evident because all frames share the same text prompt embedding. However, in the Image-to-Video task, without shared semantic information, it may lead to the inability to effectively transfer control frame information to each individual frame.
+
+<center>
+<figure>
+    <img src="https://github.com/user-attachments/assets/06df193a-fe89-42c1-8c01-fb3b7c2be0e3" height=400 />
+	<img src="https://github.com/user-attachments/assets/09906df4-ab9f-443d-8d38-512e16075b0c" height=400 />
+</figure>
+</center>
+
+We conducted a simple comparison of the performance of using the Inpainting Model under the 2+1D structure (Open-Sora Plan v1.1, left in the figure) versus the 3D structure (Open-Sora Plan v1.2, right in the figure). With the same number of optimization steps, the probability of unstable visual performance in the 2+1D structure was significantly higher than in the 3D structure. Even at convergence, the 2+1D structure's visual stability was still inferior to that of the 3D structure, and it was even worse than the early training stages of the 3D structure.
+
+## Future Work and Discussion
+
+#### CausalVideoVAE
+We observed that high-frequency motion information in videos tends to exhibit jitter, and increasing training duration and data volume does not significantly alleviate this issue. In videos, compressing the duration while maintaining the original latent dimension can lead to significant information loss. A more robust VAE will be released in the next version.
+
+#### Diffusion Model
+We replaced T5 with mT5 to enhance multilingual capabilities, but this capability is limited as our training data is currently only in English. The multilingual ability primarily comes from the mT5 mapping space. We will explore additional text encoders and expand the data in the next steps.
+
+Our model performs well in generating character consistency, likely due to panda70m being a character-centric dataset. However, it still shows poor performance in text consistency and object generalization. We suspect this may be due to the limited amount of data the model has seen, as evidenced by the non-convergence of the loss in the final stage. **We hope to collaborate with the open-source community to optimize the 3D DiT architecture.**
diff --git a/examples/conditional_images_path.txt b/examples/conditional_images_path.txt
new file mode 100644
index 000000000..e8159a9aa
--- /dev/null
+++ b/examples/conditional_images_path.txt
@@ -0,0 +1,2 @@
+/path/to/image_0.png
+/path/to/image_1_0.png,/path/to/image_1_1.png
diff --git a/examples/get_latents_std.py b/examples/get_latents_std.py
new file mode 100644
index 000000000..5f9937024
--- /dev/null
+++ b/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 = }')
\ No newline at end of file
diff --git a/examples/prompt_list_0.txt b/examples/prompt_list_0.txt
new file mode 100644
index 000000000..92fb551c6
--- /dev/null
+++ b/examples/prompt_list_0.txt
@@ -0,0 +1,20 @@
+A young man at his 20s is sitting on a piece of cloud in the sky, reading a book.
+An extreme close-up of an gray-haired man with a beard in his 60s, he is deep in thought pondering the history of the universe as he sits at a cafe in Paris, his eyes focus on people offscreen as they walk as he sits mostly motionless, he is dressed in a wool coat suit coat with a button-down shirt, he wears a brown beret and glasses and has a very professorial appearance, and the end he offers a subtle closed-mouth smile as if he found the answer to the mystery of life, the lighting is very cinematic with the golden light and the Parisian streets and city in the background, depth of field, cinematic 35mm film.
+A close-up of a woman’s face, illuminated by the soft light of dawn, her expression serene and content as she wakes up in a cozy bedroom.
+An intense close-up of a detective’s face, lit by a single desk lamp, his eyes scanning a wall covered in photos and notes, deep in thought.
+Audience members in a theater are captured in a series of medium shots, with a young man and woman in formal attire centrally positioned and illuminated by a spotlight effect.
+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.
+a realistic 3d rendering of a female character with curly blonde hair and blue eyes. she is wearing a black tank top and has a neutral expression while facing the camera directly. the background is a plain blue sky, and the scene is devoid of any other objects or text. the character is detailed, with realistic textures and lighting, suitable for a video game or high-quality animation. there is no movement or additional action in the video. the focus is entirely on the character's appearance and realistic rendering.
+A panda strumming a guitar under a bamboo grove, its paws gently plucking the strings as a group of mesmerized rabbits watch, the music blending with the rustle of bamboo leaves. HD.
+A close-up of a woman with a vintage hairstyle and bright red lipstick, gazing seductively into the camera, the background blurred to keep the focus solely on her.
+In the jungle, a hidden temple stands guarded by statues of lions, their eyes glowing with emerald light, protecting secrets untold for millennia. 8K.
+A close-up of an old man’s weathered face, with deep wrinkles and a thick white mustache, looking out to sea, the wind gently blowing through his hair.
+An intense close-up of a soldier’s face, covered in dirt and sweat, his eyes filled with determination as he surveys the battlefield.
+A river that flows uphill, defying gravity as it returns lost treasures from the sea to the mountain top, each item telling a story of a voyage gone by. HD.
+A close-up of a man’s face, lit only by the glow of his computer screen, his eyes wide and unblinking as he discovers something shocking online.
+On a deserted island, palm trees sway to summon a rainstorm, their leaves conducting the wind like maestros, orchestrating a symphony of thunder and lightning. High Resolution.
+An extreme close-up of a middle-aged man’s face, with a five o’clock shadow, staring pensively into the distance as rain softly taps against the window beside him, his thoughts deep and contemplative.
+A close-up of a man’s face, his expression one of deep concentration as he works on a complex task.
+Drone view of waves crashing against the rugged cliffs along Big Sur's garay point beach.The crashing blue waters create white-tipped waves,while the golden light of the setting sun illuminates the rocky shore. A small island with a lighthouse sits in the distance, and green
+shrubbery covers the cliffs edge. The steep drop from the road down to the beach is adramatic feat, with the cliff's edges jutting out over the sea. This is a view that captures the raw beauty of the coast and the rugged landscape of the Pacific Coast Highway.
+a close-up shot of a woman standing in a dimly lit room. she is wearing a traditional chinese outfit, which includes a red and gold dress with intricate designs and a matching headpiece. the woman has her hair styled in an updo, adorned with a gold accessory. her makeup is done in a way that accentuates her features, with red lipstick and dark eyeshadow. she is looking directly at the camera with a neutral expression. the room has a rustic feel, with wooden beams and a stone wall visible in the background. the lighting in the room is soft and warm, creating a contrast with the woman's vibrant attire. there are no texts or other objects in the video. the style of the video is a portrait, focusing on the woman and her attire.
\ No newline at end of file
diff --git a/examples/rec_image.py b/examples/rec_image.py
new file mode 100644
index 000000000..ad869eb04
--- /dev/null
+++ b/examples/rec_image.py
@@ -0,0 +1,67 @@
+import sys
+from PIL import Image
+import torch
+from torchvision.transforms import ToTensor, Compose, Resize, Normalize, Lambda
+from torch.nn import functional as F
+import argparse
+import numpy as np
+from opensora.models import CausalVAEModelWrapper
+
+def preprocess(video_data: torch.Tensor, short_size: int = 128) -> torch.Tensor:
+    transform = Compose(
+        [
+            ToTensor(),
+            Lambda(lambda x: 2. * x - 1.), 
+            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
+    short_size = args.short_size
+    device = args.device
+    kwarg = {}
+    
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir', **kwarg).to(device)
+    vae = CausalVAEModelWrapper(args.ae_path, **kwarg).to(device)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+    vae.eval()
+    vae = vae.to(device)
+    vae = vae.half()
+    
+    with torch.no_grad():
+        x_vae = preprocess(Image.open(image_path), short_size)
+        x_vae = x_vae.to(device, dtype=torch.float16)  # b c t h w
+        latents = vae.encode(x_vae)
+        latents = latents.to(torch.float16)
+        image_recon = vae.decode(latents)  # b t c h w
+    x = image_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('--ae', type=str, default='')
+    parser.add_argument('--ae_path', type=str, default='')
+    parser.add_argument('--model_path', type=str, default='results/pretrained')
+    parser.add_argument('--short_size', type=int, default=336)
+    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)
diff --git a/examples/rec_video.py b/examples/rec_video.py
new file mode 100644
index 000000000..06ee2b0d2
--- /dev/null
+++ b/examples/rec_video.py
@@ -0,0 +1,131 @@
+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
+from opensora.models import CausalVAEModelWrapper
+from opensora.dataset.transform import ToTensorVideo, CenterCropResizeVideo
+
+
+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(0, 2, 3, 1).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
+
+
+def preprocess(video_data: torch.Tensor, height: int = 128, width: int = 128) -> torch.Tensor:
+    transform = Compose(
+        [
+            ToTensorVideo(),
+            CenterCropResizeVideo((height, width)),
+            Lambda(lambda x: 2. * x - 1.)
+        ]
+    )
+
+    video_outputs = transform(video_data)
+    video_outputs = torch.unsqueeze(video_outputs, 0)
+
+    return video_outputs
+
+
+def main(args: argparse.Namespace):
+    device = args.device
+    kwarg = {}
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir', **kwarg).to(device)
+    vae = CausalVAEModelWrapper(args.ae_path, **kwarg).to(device)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.eval()
+    vae = vae.to(device)
+    vae = vae.half()
+
+    with torch.no_grad():
+        x_vae = preprocess(read_video(args.video_path, args.num_frames, args.sample_rate), args.height,
+                           args.width)
+        print(x_vae.shape)
+        x_vae = x_vae.to(device, dtype=torch.float16)  # b c t h w
+        latents = vae.encode(x_vae)
+        print(latents.shape)
+        latents = latents.to(torch.float16)
+        video_recon = vae.decode(latents)  # b t c h w
+        print(video_recon.shape)
+
+    custom_to_video(video_recon[0], fps=args.fps, 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('--ae', type=str, default='')
+    parser.add_argument('--ae_path', type=str, default='')
+    parser.add_argument('--model_path', type=str, default='results/pretrained')
+    parser.add_argument('--fps', type=int, default=30)
+    parser.add_argument('--height', type=int, default=336)
+    parser.add_argument('--width', type=int, default=336)
+    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('--tile_sample_min_size', type=int, default=256)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--save_memory', action='store_true')
+
+    args = parser.parse_args()
+    main(args)
diff --git a/opensora/__init__.py b/opensora/__init__.py
new file mode 100644
index 000000000..4e768b56d
--- /dev/null
+++ b/opensora/__init__.py
@@ -0,0 +1 @@
+# 
\ No newline at end of file
diff --git a/opensora/acceleration/__init__.py b/opensora/acceleration/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/acceleration/communications.py b/opensora/acceleration/communications.py
new file mode 100644
index 000000000..6f69a73a0
--- /dev/null
+++ b/opensora/acceleration/communications.py
@@ -0,0 +1,133 @@
+import torch
+import torch.distributed as dist
+from einops import rearrange
+from opensora.acceleration.parallel_states import hccl_info, lccl_info, enable_LCCL
+try:
+    from lcalib.functional import lcal_all2allvc
+except:
+    lcal_all2allvc = None
+
+def broadcast(input_: torch.Tensor):
+    sp_size = hccl_info.world_size
+    src = hccl_info.rank // sp_size * sp_size
+    dist.broadcast(input_, src=src, group=hccl_info.group)
+
+_COUNT = 0
+def _all_to_all(
+    input_: torch.Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+):
+    group = hccl_info.group
+    sp_size = hccl_info.world_size
+    input_list = [t.contiguous() for t in torch.tensor_split(input_, sp_size, scatter_dim)]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+def _single_all_to_all(
+    input_: torch.Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    enable_HCCL=False,
+):
+    if enable_LCCL:
+        sp_size = lccl_info.world_size
+    else:
+        sp_size = hccl_info.world_size
+    inp_shape = list(input_.shape)
+    inp_shape[scatter_dim] = inp_shape[scatter_dim] // sp_size
+    if scatter_dim < 1:
+        input_t = input_.reshape(
+            [sp_size, inp_shape[scatter_dim]] + \
+            inp_shape[scatter_dim + 1:]
+        )
+    else:
+        # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+        input_t = input_.reshape(
+            [-1, sp_size, inp_shape[scatter_dim]] + \
+            inp_shape[scatter_dim + 1:]
+        ).transpose(0, 1).contiguous()
+
+    output = torch.empty_like(input_t)
+    if enable_LCCL and not enable_HCCL:
+        matrix_count = torch.ones([sp_size, sp_size], dtype=torch.int64, device=input_t.device) * (
+                    input_t.numel() // sp_size)
+        lcal_all2allvc(input_t, output, matrix_count, lccl_info.group)
+    else:
+        dist.all_to_all_single(output, input_t, group=hccl_info.group)
+    # if scattering the seq-dim, transpose the heads back to the original dimension
+    if scatter_dim < 1:
+        output = output.transpose(0, 1).contiguous()
+
+    return output.reshape(
+        inp_shape[: gather_dim] + [inp_shape[gather_dim] * sp_size, ] + inp_shape[gather_dim + 1:])
+
+
+class _AllToAll(torch.autograd.Function):
+    """All-to-all communication.
+
+    Args:
+        input_: input matrix
+        process_group: communication group
+        scatter_dim: scatter dimension
+        gather_dim: gather dimension
+    """
+
+    @staticmethod
+    def forward(ctx, input_, scatter_dim, gather_dim, all_to_all_func):
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.all_to_all = all_to_all_func
+        output = ctx.all_to_all(input_, scatter_dim, gather_dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output = ctx.all_to_all(
+            grad_output,
+            ctx.gather_dim,
+            ctx.scatter_dim,
+        )
+        return (
+            grad_output,
+            None,
+            None,
+            None,
+        )
+
+def all_to_all_SBH(
+    input_: torch.Tensor,
+    scatter_dim: int = 1,
+    gather_dim: int = 0,
+):
+    return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all)
+
+def all_to_all_BSND(
+    input_: torch.Tensor,
+    scatter_dim: int = 2,
+    gather_dim: int = 1,
+):
+    return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all)
+
+
+def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num):
+    def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
+        hidden_states = _single_all_to_all(hidden_states, scatter_dim=2, gather_dim=0, enable_HCCL=True)
+        encoder_hidden_states = _single_all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0, enable_HCCL=True)
+        attention_mask = _single_all_to_all(attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True)
+        encoder_attention_mask = _single_all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True)
+        return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
+
+    sp_size = hccl_info.world_size
+    frame = hidden_states.shape[2]
+    assert frame % sp_size == 0, "frame should be a multiple of sp_size"
+
+    encoder_hidden_states = rearrange(encoder_hidden_states, 'b 1 (n x) h -> b n x h',
+                                     n=sp_size, x=encoder_hidden_states.shape[2]//sp_size).contiguous()
+    hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask = all_to_all(hidden_states,
+                                                                                            encoder_hidden_states,
+                                                                                            attention_mask.repeat(1, sp_size, 1, 1),
+                                                                                            encoder_attention_mask.repeat(1, sp_size, 1))
+
+    return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num
\ No newline at end of file
diff --git a/opensora/acceleration/parallel_states.py b/opensora/acceleration/parallel_states.py
new file mode 100644
index 000000000..e3f96e121
--- /dev/null
+++ b/opensora/acceleration/parallel_states.py
@@ -0,0 +1,57 @@
+import torch
+import torch_npu
+import torch.distributed as dist
+import os
+try:
+    from lcalib.functional import lcal_initialize
+    enable_LCCL = True
+except:
+    lcal_initialize = None
+    enable_LCCL = False
+class COMM_INFO:
+    def __init__(self):
+        self.group = None
+        self.world_size = 0
+        self.rank = -1
+
+lccl_info = COMM_INFO()
+hccl_info = COMM_INFO()
+_SEQUENCE_PARALLEL_STATE = False
+def initialize_sequence_parallel_state(sequence_parallel_size):
+    global _SEQUENCE_PARALLEL_STATE
+    if sequence_parallel_size > 1:
+        _SEQUENCE_PARALLEL_STATE = True
+        initialize_sequence_parallel_group(sequence_parallel_size)
+
+def set_sequence_parallel_state(state):
+    global _SEQUENCE_PARALLEL_STATE
+    _SEQUENCE_PARALLEL_STATE = state
+
+def get_sequence_parallel_state():
+    return _SEQUENCE_PARALLEL_STATE
+
+def initialize_sequence_parallel_group(sequence_parallel_size):
+    """Initialize the sequence parallel group."""
+    rank = int(os.getenv('RANK', '0'))
+    world_size = int(os.getenv("WORLD_SIZE", '1'))
+    assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size"
+    # hccl
+    hccl_info.world_size = sequence_parallel_size
+    hccl_info.rank = rank
+    num_sequence_parallel_groups: int = world_size // sequence_parallel_size
+    for i in range(num_sequence_parallel_groups):
+        ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size)
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            hccl_info.group = group
+
+    if enable_LCCL:
+        assert sequence_parallel_size == 8, "sequence_parallel_size should be 8 when enable_LCCL is True"
+        rank %= sequence_parallel_size
+        lccl_info.world_size = sequence_parallel_size
+        lccl_info.group = lcal_initialize(rank, sequence_parallel_size)
+        lccl_info.rank = rank
+
+def destroy_sequence_parallel_group():
+    """Destroy the sequence parallel group."""
+    dist.destroy_process_group()
diff --git a/opensora/adaptor/__init__.py b/opensora/adaptor/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/adaptor/bf16_optimizer.py b/opensora/adaptor/bf16_optimizer.py
new file mode 100644
index 000000000..e0401c0b3
--- /dev/null
+++ b/opensora/adaptor/bf16_optimizer.py
@@ -0,0 +1,458 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from collections import OrderedDict
+import torch
+import sys
+import os
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+from deepspeed import comm as dist
+from deepspeed.runtime.constants import PIPE_REPLICATED
+from deepspeed.runtime import ZeROOptimizer
+from packaging import version as pkg_version
+
+from deepspeed.git_version_info import version
+from deepspeed.runtime.utils import (get_global_norm_of_tensors, clip_tensors_by_global_norm, DummyOptim,
+                                     align_dense_tensors, all_gather_dp_groups, bwc_tensor_model_parallel_rank,
+                                     is_model_parallel_parameter, see_memory_usage, graph_process)
+
+from deepspeed.utils import link_hp_params, fragment_address
+from deepspeed.checkpoint import enable_universal_checkpoint
+from deepspeed.checkpoint.constants import (DS_VERSION, PARTITION_COUNT, BASE_OPTIMIZER_STATE,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS, CLIP_GRAD, GROUP_PADDINGS,
+                                            PARAM_SLICE_MAPPINGS)
+
+setattr(sys.modules[__name__], 'fragment_address', fragment_address)
+
+
+def contigous_flatten(tensors):
+    return _flatten_dense_tensors([tensor.contiguous() for tensor in tensors])
+
+
+class BF16_Optimizer(ZeROOptimizer):
+
+    def __init__(self,
+                 init_optimizer,
+                 param_names,
+                 mpu=None,
+                 clip_grad=0.0,
+                 norm_type=2,
+                 allgather_bucket_size=5000000000,
+                 dp_process_group=None,
+                 timers=None,
+                 grad_acc_dtype=None,
+                 graph_harvesting=False):
+        # super().__init__()
+        # base_class = ZeROOptimizer.__bases__[0]
+        # # 直接调用基类的 __init__ 方法
+        # base_class.__init__()
+        see_memory_usage('begin bf16_optimizer', force=True)
+        self.timers = timers
+        self.optimizer = init_optimizer
+        self.param_names = param_names
+        self.using_real_optimizer = not isinstance(self.optimizer, DummyOptim)
+
+        assert grad_acc_dtype in [torch.float32, torch.bfloat16
+                                  ], f"BF16Optimizer: Unsupported gradient accumulation data type: {grad_acc_dtype}"
+        self.grad_acc_dtype = grad_acc_dtype
+
+        self.clip_grad = clip_grad
+        self.norm_type = norm_type
+        self.mpu = mpu
+        self.allgather_bucket_size = int(allgather_bucket_size)
+        self.dp_process_group = dp_process_group
+        self.dp_rank = dist.get_rank(group=self.dp_process_group)
+        self.real_dp_process_group = [dp_process_group for i in range(len(self.optimizer.param_groups))]
+
+        # Use torch (un)flatten ops
+        self.flatten = contigous_flatten
+        self.unflatten = _unflatten_dense_tensors
+
+        #align nccl all-gather send buffers to 4-bye boundary
+        self.nccl_start_alignment_factor = 16
+
+        # Build BF16/FP32 groups
+        self.bf16_groups = []
+        self.bf16_groups_flat = []
+        self.bf16_partitioned_groups = []
+
+        self.fp32_groups_flat_partition = []
+
+        # Maintain different fp32 gradients views for convenience
+        self.fp32_groups_gradients = []
+        self.fp32_groups_gradient_dict = {}
+        self.fp32_groups_gradients_flat = []
+        self.fp32_groups_actual_gradients_flat = []
+        self.fp32_groups_gradient_flat_partition = []
+        self.fp32_groups_has_gradients = []
+
+        self.group_paddings = []
+        self.graph_harvesting = graph_harvesting
+        if self.using_real_optimizer:
+            self._setup_for_real_optimizer()
+
+        see_memory_usage('end bf16_optimizer', force=True)
+
+    def _setup_for_real_optimizer(self):
+        dp_world_size = dist.get_world_size(group=self.dp_process_group)
+        self.partition_count = [dp_world_size for i in range(len(self.optimizer.param_groups))]
+
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            see_memory_usage(f'before initializing group {i}', force=True)
+
+            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
+
+            # grab the original list
+            trainable_parameters = [param for param in param_group['params'] if param.requires_grad]
+            self.bf16_groups.append(trainable_parameters)
+
+            # create flat bf16 params
+            self.bf16_groups_flat.append(
+                self._flatten_dense_tensors_aligned(self.bf16_groups[i],
+                                                    self.nccl_start_alignment_factor * dp_world_size))
+
+            # Make bf16 params point to flat tensor storage
+            self._update_storage_to_flattened_tensor(tensor_list=self.bf16_groups[i],
+                                                     flat_tensor=self.bf16_groups_flat[i])
+
+            # divide flat weights into equal sized partitions
+            partition_size = self.bf16_groups_flat[i].numel() // dp_world_size
+            bf16_dp_partitions = [
+                self.bf16_groups_flat[i].narrow(0, dp_index * partition_size, partition_size)
+                for dp_index in range(dp_world_size)
+            ]
+            self.bf16_partitioned_groups.append(bf16_dp_partitions)
+
+            # create fp32 params partition
+            self.fp32_groups_flat_partition.append(bf16_dp_partitions[partition_id].clone().float().detach())
+            self.fp32_groups_flat_partition[i].requires_grad = True
+
+            num_elem_list = [t.numel() for t in self.bf16_groups[i]]
+
+            # create fp32 gradients
+            self.fp32_groups_gradients_flat.append(
+                torch.zeros_like(self.bf16_groups_flat[i], dtype=self.grad_acc_dtype))
+
+            # track individual fp32 gradients for entire model
+            fp32_gradients = self._split_flat_tensor(flat_tensor=self.fp32_groups_gradients_flat[i],
+                                                     num_elem_list=num_elem_list)
+            self.fp32_groups_gradients.append(fp32_gradients)
+            self.fp32_groups_gradient_dict[i] = fp32_gradients
+
+            # flat tensor corresponding to actual fp32 gradients (i.e., minus alignment padding)
+            length_without_padding = sum(num_elem_list)
+            self.fp32_groups_actual_gradients_flat.append(
+                torch.narrow(self.fp32_groups_gradients_flat[i], 0, 0, length_without_padding))
+
+            # flat tensor corresponding to gradient partition
+            self.fp32_groups_gradient_flat_partition.append(
+                torch.narrow(self.fp32_groups_gradients_flat[i], 0, partition_id * partition_size, partition_size))
+
+            # track fp32 gradient updates
+            self.fp32_groups_has_gradients.append([False] * len(self.bf16_groups[i]))
+
+            # Record padding required for alignment
+            if partition_id == dist.get_world_size(group=self.real_dp_process_group[i]) - 1:
+                padding = self.bf16_groups_flat[i].numel() - length_without_padding
+            else:
+                padding = 0
+
+            self.group_paddings.append(padding)
+
+            # update optimizer param groups to reference fp32 params partition
+            param_group['params'] = [self.fp32_groups_flat_partition[i]]
+
+            see_memory_usage(f'after initializing group {i}', force=True)
+
+        see_memory_usage('before initialize_optimizer', force=True)
+        self.initialize_optimizer_states()
+        see_memory_usage('end initialize_optimizer', force=True)
+
+        # Need optimizer states initialized before linking lp to optimizer state
+        self._link_all_hp_params()
+        self._enable_universal_checkpoint()
+        self._param_slice_mappings = self._create_param_mapping()
+
+    def _enable_universal_checkpoint(self):
+        for lp_param_group in self.bf16_groups:
+            enable_universal_checkpoint(param_list=lp_param_group)
+
+    def _create_param_mapping(self):
+        param_mapping = []
+        for i, _ in enumerate(self.optimizer.param_groups):
+            param_mapping_per_group = OrderedDict()
+            for lp in self.bf16_groups[i]:
+                if lp._hp_mapping is not None:
+                    lp_name = self.param_names[lp]
+                    param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address()
+            param_mapping.append(param_mapping_per_group)
+
+        return param_mapping
+
+    def _link_all_hp_params(self):
+        dp_world_size = dist.get_world_size(group=self.dp_process_group)
+        for i, _ in enumerate(self.optimizer.param_groups):
+            # Link bf16 and fp32 params in partition
+            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
+            partition_size = self.bf16_groups_flat[i].numel() // dp_world_size
+            flat_hp_partition = self.fp32_groups_flat_partition[i]
+            link_hp_params(lp_param_list=self.bf16_groups[i],
+                           flat_hp_partition=flat_hp_partition,
+                           gradient_dict=self.fp32_groups_gradient_dict,
+                           offload_gradient_dict=None,
+                           use_offload=False,
+                           param_group_index=i,
+                           partition_start=partition_id * partition_size,
+                           partition_size=partition_size,
+                           partition_optimizer_state=self.optimizer.state[flat_hp_partition],
+                           dp_group=self.real_dp_process_group[i])
+
+    def initialize_optimizer_states(self):
+        """Take an optimizer step with zero-valued gradients to allocate internal
+        optimizer state.
+
+        This helps prevent memory fragmentation by allocating optimizer state at the
+        beginning of training instead of after activations have been allocated.
+        """
+        for param_partition, grad_partition in zip(self.fp32_groups_flat_partition,
+                                                   self.fp32_groups_gradient_flat_partition):
+            # In case of grad acc dtype different than FP32, need to cast to high precision.
+            param_partition.grad = grad_partition.to(
+                param_partition.dtype) if grad_partition.dtype != param_partition.dtype else grad_partition
+
+        self.optimizer.step()
+
+        if self.grad_acc_dtype is not torch.float32:
+            for param_partition in self.fp32_groups_flat_partition:
+                param_partition.grad = None
+
+        self.clear_hp_grads()
+
+    def _split_flat_tensor(self, flat_tensor, num_elem_list):
+        assert sum(num_elem_list) <= flat_tensor.numel()
+        tensor_list = []
+        offset = 0
+        for num_elem in num_elem_list:
+            dense_tensor = torch.narrow(flat_tensor, 0, offset, num_elem)
+            tensor_list.append(dense_tensor)
+            offset += num_elem
+
+        return tensor_list
+
+    def _update_storage_to_flattened_tensor(self, tensor_list, flat_tensor):
+        updated_params = self.unflatten(flat_tensor, tensor_list)
+        for p, q in zip(tensor_list, updated_params):
+            p.data = q.data
+
+    def _flatten_dense_tensors_aligned(self, tensor_list, alignment):
+        return self.flatten(align_dense_tensors(tensor_list, alignment))
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        if closure is not None:
+            raise NotImplementedError(f'{self.__class__} does not support closure.')
+
+        all_groups_norm = get_global_norm_of_tensors(input_tensors=self.get_grads_for_norm(),
+                                                     mpu=self.mpu,
+                                                     norm_type=self.norm_type,
+                                                     use_graph=self.graph_harvesting)
+        self._global_grad_norm = all_groups_norm
+
+        assert all_groups_norm > 0.
+        if self.clip_grad > 0.:
+            clip_tensors_by_global_norm(input_tensors=self.get_grads_for_norm(for_clipping=True),
+                                        max_norm=self.clip_grad,
+                                        global_norm=all_groups_norm,
+                                        mpu=self.mpu,
+                                        use_graph=self.graph_harvesting)
+
+        self.optimizer.step()
+
+        self.update_lp_params()
+
+        self.clear_hp_grads()
+
+    def backward(self, loss, update_hp_grads=True, clear_lp_grads=False, **bwd_kwargs):
+        """Perform a backward pass and copy the low-precision gradients to the
+        high-precision copy.
+
+        We copy/accumulate to the high-precision grads now to prevent accumulating in the
+        bf16 grads after successive backward() calls (i.e., grad accumulation steps > 1)
+
+        The low-precision grads are deallocated during this procedure.
+        """
+        self.clear_lp_grads()
+        loss.backward(**bwd_kwargs)
+
+        if update_hp_grads:
+            self.update_hp_grads(clear_lp_grads=clear_lp_grads)
+
+    @torch.no_grad()
+    def update_hp_grads(self, clear_lp_grads=False):
+
+        def _update_hp_grads_func(clear_lp_grads=False):
+            for i, group in enumerate(self.bf16_groups):
+                for j, lp in enumerate(group):
+                    if lp.grad is None:
+                        continue
+                    hp_grad = self.fp32_groups_gradients[i][j]
+                    assert hp_grad is not None, \
+                        f'high precision param has no gradient, lp param_id = {id(lp)} group_info = [{i}][{j}]'
+                    hp_grad.data.add_(lp.grad.data.to(hp_grad.dtype).view(hp_grad.shape))
+                    lp._hp_grad = hp_grad
+                    self.fp32_groups_has_gradients[i][j] = True
+                    # clear gradients
+                    if clear_lp_grads:
+                        lp.grad._zero()
+
+        if self.graph_harvesting:
+            graph_process(False, _update_hp_grads_func, clear_lp_grads)
+        else:
+            _update_hp_grads_func(clear_lp_grads)
+        #cpu op
+        for i, group in enumerate(self.bf16_groups):
+            for j, lp in enumerate(group):
+                if lp.grad is None:
+                    continue
+                self.fp32_groups_has_gradients[i][j] = True
+
+    @torch.no_grad()
+    def get_grads_for_reduction(self):
+        return self.fp32_groups_gradients_flat
+
+    @torch.no_grad()
+    def get_grads_for_norm(self, for_clipping=False):
+        grads = []
+        tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu)
+        for i, group in enumerate(self.bf16_groups):
+            for j, lp in enumerate(group):
+                if not for_clipping:
+                    if hasattr(lp, PIPE_REPLICATED) and lp.ds_pipe_replicated:
+                        continue
+
+                    if not (tensor_mp_rank == 0 or is_model_parallel_parameter(lp)):
+                        continue
+
+                if not self.fp32_groups_has_gradients[i][j]:
+                    continue
+
+                grads.append(self.fp32_groups_gradients[i][j])
+
+        return grads
+
+    @torch.no_grad()
+    def update_lp_params(self):
+        for i, (bf16_partitions,
+                fp32_partition) in enumerate(zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition)):
+            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
+            bf16_partitions[partition_id].data.copy_(fp32_partition.data)
+            # print_rank_0(f'update_lp_params {i=} {partition_id=}', force=True)
+            # if i == 0:
+            #     print_rank_0(f'{fp32_partition[:10]=}', force=True)
+
+        all_gather_dp_groups(groups_flat=self.bf16_groups_flat,
+                             partitioned_param_groups=self.bf16_partitioned_groups,
+                             dp_process_group=self.real_dp_process_group,
+                             start_alignment_factor=self.nccl_start_alignment_factor,
+                             allgather_bucket_size=self.allgather_bucket_size)
+
+    def clear_hp_grads(self):
+        for flat_gradients in self.fp32_groups_gradients_flat:
+            flat_gradients.zero_()
+
+        for i, group in enumerate(self.fp32_groups_gradients):
+            self.fp32_groups_has_gradients[i] = [False] * len(group)
+
+    def clear_lp_grads(self):
+        for group in self.bf16_groups:
+            for param in group:
+                if param.grad is not None:
+                    # Using zero_() fixed memory address for graph replay
+                    param.grad.zero_()
+
+    def state_dict(self):
+        state_dict = {}
+        state_dict[CLIP_GRAD] = self.clip_grad
+        state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict()
+        state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = self.fp32_groups_flat_partition
+        state_dict[GROUP_PADDINGS] = self.group_paddings
+        state_dict[PARTITION_COUNT] = self.partition_count
+        state_dict[DS_VERSION] = version
+        state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings
+
+        return state_dict
+
+    # Restore base optimizer fp32 weights bfloat16 weights
+    def _restore_from_bit16_weights(self):
+        for i, group in enumerate(self.bf16_groups):
+            partition_id = dist.get_rank(group=self.real_dp_process_group[i])
+            for bf16_partitions, fp32_partition in zip(self.bf16_partitioned_groups, self.fp32_groups_flat_partition):
+                fp32_partition.data.copy_(bf16_partitions[partition_id].data)
+
+    def refresh_fp32_params(self):
+        self._restore_from_bit16_weights()
+
+    def load_state_dict(self,
+                        state_dict_list,
+                        checkpoint_folder,
+                        load_optimizer_states=True,
+                        load_from_fp32_weights=False,
+                        load_serial=None):
+        if checkpoint_folder:
+            self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights)
+        else:
+            self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights)
+
+    def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False):
+
+        dp_rank = dist.get_rank(group=self.dp_process_group)
+        current_rank_sd = state_dict_list[dp_rank]
+
+        ckpt_version = current_rank_sd.get(DS_VERSION, False)
+        assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed"
+        ckpt_version = pkg_version.parse(ckpt_version)
+
+        self.clip_grad = current_rank_sd.get(CLIP_GRAD, self.clip_grad)
+
+        if load_optimizer_states:
+            self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE])
+
+        if load_from_fp32_weights:
+            for current, saved in zip(self.fp32_groups_flat_partition,
+                                      current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]):
+                src_tensor = _get_padded_tensor(saved, current.numel())
+                current.data.copy_(src_tensor.data)
+
+        if load_optimizer_states:
+            self._link_all_hp_params()
+
+    def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights):
+        self._load_hp_checkpoint_state(checkpoint_folder)
+
+    @property
+    def param_groups(self):
+        """Forward the wrapped optimizer's parameters."""
+        return self.optimizer.param_groups
+
+    def _load_hp_checkpoint_state(self, checkpoint_dir):
+        checkpoint_dir = os.path.join(checkpoint_dir, "zero")
+        tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu)
+        tp_world_size = self.mpu.get_slice_parallel_world_size()
+
+        for i, _ in enumerate(self.optimizer.param_groups):
+            for lp in self.bf16_groups[i]:
+                if lp._hp_mapping is not None:
+                    #print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}")
+                    lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank,
+                                                tp_world_size)
+
+
+def _get_padded_tensor(src_tensor, size):
+    if src_tensor.numel() >= size:
+        return src_tensor
+    padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device)
+    slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel())
+    slice_tensor.data.copy_(src_tensor.data)
+    return padded_tensor
diff --git a/opensora/adaptor/engine.py b/opensora/adaptor/engine.py
new file mode 100644
index 000000000..5208a9e41
--- /dev/null
+++ b/opensora/adaptor/engine.py
@@ -0,0 +1,3568 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import re
+import stat
+import torch
+import hashlib
+from collections import defaultdict, OrderedDict, deque
+from shutil import copyfile
+import gc
+
+from torch.nn.modules import Module
+from torch.nn.parameter import Parameter
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from typing import Callable, Dict, Union, Iterable
+
+import deepspeed
+
+from deepspeed import comm as dist
+from deepspeed.runtime.utils import see_memory_usage, DummyOptim
+from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
+from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
+from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+from deepspeed.runtime.zero.utils import is_zero_supported_optimizer, ZeRORuntimeException
+from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload
+from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION
+
+from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer
+from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
+from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
+
+from deepspeed.runtime.config import DEEPSPEED_OPTIMIZERS, \
+    ADAGRAD_OPTIMIZER, ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, \
+    TORCH_ADAM_PARAM, ADAM_W_MODE, ADAM_W_MODE_DEFAULT, ZERO_ONE_ADAM_OPTIMIZER, MUADAM_OPTIMIZER, MUADAMW_OPTIMIZER, \
+    MUSGD_OPTIMIZER, LION_OPTIMIZER
+
+from deepspeed.runtime.dataloader import DeepSpeedDataLoader
+from deepspeed.runtime.constants import \
+    ROUTE_TRAIN, ROUTE_PREDICT, ROUTE_EVAL, \
+    PLD_THETA, PLD_GAMMA, BFLOAT16, FP16, AMP, GRADIENT_ACCUMULATION_STEPS, \
+    DATA_PARALLEL_GROUP, GLOBAL_RANK
+from deepspeed.runtime.zero.config import ZeroStageEnum
+from deepspeed.compression import compression_scheduler
+from deepspeed.compression.constants import \
+    WEIGHT_QUANTIZE_IN_FORWARD_ENABLED, \
+    WEIGHT_QUANTIZATION, SHARED_PARAMETERS, \
+    WEIGHT_QUANTIZE_ENABLED, \
+    WEIGHT_QUANTIZE_GROUPS, \
+    WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE, \
+    WEIGHT_QUANTIZE_CHANGE_RATIO, \
+    WEIGHT_QUANTIZE_TYPE, \
+    WEIGHT_QUANTIZE_ROUNDING, \
+    WEIGHT_QUANTIZE_VERBOSE, \
+    WEIGHT_QUANTIZE_KERNEL
+from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT, FROZEN_PARAM_FRAGMENTS
+from deepspeed.runtime.sparse_tensor import SparseTensor
+
+from deepspeed.runtime import lr_schedules
+from deepspeed.utils import groups
+from deepspeed.utils import logger, log_dist, instrument_w_nvtx
+from deepspeed.utils.timer import NoopTimer, ThroughputTimer, SynchronizedWallClockTimer, \
+    FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER, \
+    STEP_MICRO_TIMER, \
+    FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER, \
+    STEP_GLOBAL_TIMER
+from deepspeed.utils.debug import debug_extract_module_and_param_names
+from deepspeed.monitor.monitor import MonitorMaster
+from deepspeed.runtime.progressive_layer_drop import ProgressiveLayerDrop
+from deepspeed.runtime.utils import clip_grad_norm_
+from deepspeed.runtime.eigenvalue import Eigenvalue
+from deepspeed.runtime.data_pipeline.constants import DATA_SAMPLING, \
+    DATA_ROUTING, DATA_SAMPLING_ENABLED, CURRICULUM_LEARNING, \
+    CURRICULUM_LEARNING_ENABLED, DATA_SAMPLING_NUM_WORKERS, RANDOM_LTD, \
+    RANDOM_LTD_ENABLED, RANDOM_LTD_LAYER_ID, RANDOM_LTD_LAYER_NUM, \
+    RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE, RANDOM_LTD_LAYER_TOKEN_LR_ENABLED, \
+    RANDOM_LTD_GLOBAL_BATCH_SIZE, RANDOM_LTD_MICRO_BATCH_SIZE, DATA_EFFICIENCY
+from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler
+from deepspeed.runtime.data_pipeline.data_routing.scheduler import RandomLTDScheduler
+from deepspeed.runtime.data_pipeline.data_routing.helper import remove_random_ltd_state_dict
+from deepspeed.runtime.data_pipeline.data_routing.basic_layer import RandomLayerTokenDrop
+
+from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
+from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+
+from deepspeed.runtime.pipe.module import PipelineModule
+from deepspeed.runtime.utils import get_ma_status
+from deepspeed.ops.adam import FusedAdam
+from deepspeed.moe.sharded_moe import TopKGate, MOELayer
+from deepspeed.moe.layer import MoE
+from deepspeed.moe.utils import is_moe_param
+from deepspeed.git_version_info import version
+
+from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler
+from deepspeed.utils.logging import print_json_dist, print_configuration
+
+from deepspeed.accelerator import get_accelerator
+
+from deepspeed.runtime.config import DtypeEnum
+
+from opensora.adaptor.zp_manager import zp_manager
+
+MEMORY_OPT_ALLREDUCE_SIZE = 500000000
+
+DeepSpeedOptimizerCallable = \
+    Callable[[Union[Iterable[Parameter], Dict[str, Iterable]]], Optimizer]
+DeepSpeedSchedulerCallable = Callable[[Optimizer], _LRScheduler]
+
+try:
+    import apex
+    from apex import amp
+    APEX_INSTALLED = True
+except ImportError:
+    # Fail silently so we don't spam logs unnecessarily if user isn't using amp
+    APEX_INSTALLED = False
+
+
+def split_half_float_double_sparse(tensors):
+    device_type = get_accelerator().device_name()
+    supported_types = [
+        "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type),
+        "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type),
+        SparseTensor.type()
+    ]
+
+    for t in tensors:
+        assert t.type() in supported_types, f"attempting to reduce an unsupported grad type: {t.type()}"
+
+    buckets = []
+    for i, dtype in enumerate(supported_types):
+        bucket = [t for t in tensors if t.type() == dtype]
+        if bucket:
+            buckets.append((dtype, bucket))
+    return buckets
+
+
+class EngineTimers(object):
+    r"""Wallclock timers for DeepSpeedEngine"""
+
+    def __init__(self, enable_micro_timers, enable_global_timers):
+        self.forward_timers = []
+        self.backward_timers = []
+        self.backward_inner_timers = []
+        self.backward_reduce_timers = []
+        self.step_timers = []
+        self.global_timers = []
+        self.micro_timers = []
+
+        if enable_micro_timers:
+            self.forward_timers += [FORWARD_MICRO_TIMER]
+            self.backward_timers += [BACKWARD_MICRO_TIMER]
+            self.backward_inner_timers += [BACKWARD_INNER_MICRO_TIMER]
+            self.backward_reduce_timers += [BACKWARD_REDUCE_MICRO_TIMER]
+            self.step_timers += [STEP_MICRO_TIMER]
+            self.micro_timers += [
+                FORWARD_MICRO_TIMER, BACKWARD_MICRO_TIMER, BACKWARD_INNER_MICRO_TIMER, BACKWARD_REDUCE_MICRO_TIMER,
+                STEP_MICRO_TIMER
+            ]
+
+        if enable_global_timers:
+            self.forward_timers += [FORWARD_GLOBAL_TIMER]
+            self.backward_timers += [BACKWARD_GLOBAL_TIMER]
+            self.backward_inner_timers += [BACKWARD_INNER_GLOBAL_TIMER]
+            self.backward_reduce_timers += [BACKWARD_REDUCE_GLOBAL_TIMER]
+            self.step_timers += [STEP_GLOBAL_TIMER]
+            self.global_timers += [
+                FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, BACKWARD_INNER_GLOBAL_TIMER, BACKWARD_REDUCE_GLOBAL_TIMER,
+                STEP_GLOBAL_TIMER
+            ]
+
+
+class DeepSpeedEngine(Module):
+    r"""DeepSpeed engine for training."""
+
+    def __init__(
+        self,
+        args,
+        model,
+        optimizer=None,
+        model_parameters=None,
+        training_data=None,
+        lr_scheduler=None,
+        mpu=None,
+        dist_init_required=None,
+        collate_fn=None,
+        config=None,
+        config_class=None,
+        dont_change_device=False,
+    ):
+        super(DeepSpeedEngine, self).__init__()
+        self.dont_change_device = dont_change_device
+        self.client_optimizer = optimizer
+        self.client_lr_scheduler = lr_scheduler
+        self.training_data = training_data
+        self.collate_fn = collate_fn
+        self.mpu = mpu
+        self.all_to_all_group = None
+        self.data_parallel_group = None
+        self.global_steps = 0
+        self.global_samples = 0
+        self.micro_steps = 0
+        self.skipped_steps = 0
+        self.gradient_average = True
+        self.warn_unscaled_loss = True
+        self.config = config
+        self._config = config_class
+        self.loaded_checkpoint_mp_world_size = None
+        self.loaded_checkpoint_dp_world_size = None
+        self.loaded_checkpoint_zp_world_size = None
+        self.enable_backward_allreduce = True
+        self.progressive_layer_drop = None
+        self.eigenvalue = None
+        self.block_eigenvalue = None
+        self.gas_boundary_ctr = 0
+        self.dist_backend = get_accelerator().communication_backend_name()
+        self.has_moe_layers = False
+        self.num_experts = []
+        self.gate_modules = []
+        self.moe_layers = []
+        self._step_applied = False
+        self._global_grad_norm = None
+        self.use_ds_comm = False  # False --> Use torch.dist, True --> Use ds.comm backend.
+
+        self.checkpoint_engine = None
+
+        self._is_gradient_accumulation_boundary = None
+        self.scale_wrt_gas = None
+        self.losses = 0.0
+
+        # for debug purposes - can then debug print: debug_get_module_name(module)
+        debug_extract_module_and_param_names(model)
+
+        self._do_args_sanity_check(args)
+        self._configure_with_arguments(args, mpu)
+        self._do_sanity_check()
+        see_memory_usage(f"DeepSpeed Engine: After args sanity test", force=self.memory_breakdown())
+        if mpu is not None:
+            if self.elasticity_enabled():
+                if not self.is_elastic_model_parallel_supported():
+                    assert not self.elasticity_enabled(), ("Elasticity is not currently supported"
+                                                           " with model parallelism.")
+
+        self._set_distributed_vars(args)
+
+        dist.configure(self._config)
+
+        self.monitor = MonitorMaster(self._config.monitor_config)
+
+        see_memory_usage(
+            f"DeepSpeed Engine: Before configure distributed model",
+            force=self.memory_breakdown(),
+        )
+
+        self.pipeline_parallelism = isinstance(model, PipelineModule)
+
+        # Configure distributed model
+        self._configure_distributed_model(model)
+
+        # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict
+        self.param_names = {param: name for name, param in model.named_parameters()}
+
+        self._get_model_parameters()
+
+        see_memory_usage(f"DeepSpeed Engine: After configure distributed model")
+
+        # Configure wall clock timers
+        self.timers = SynchronizedWallClockTimer()
+        # Throughput timer
+        self.tput_timer = ThroughputTimer(
+            batch_size=self.train_batch_size(),
+            steps_per_output=self.steps_per_print(),
+            monitor_memory=False,
+        )
+
+        log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0])
+
+        if self.flops_profiler_enabled():
+            self.flops_profiler = FlopsProfiler(self.module, self, self.flops_profiler_recompute_fwd_factor())
+
+        if training_data:
+            self.training_dataloader = self.deepspeed_io(training_data)
+        else:
+            self.training_dataloader = None
+
+        # Configure optimizer and scheduler
+        self.optimizer = None
+        self.basic_optimizer = None
+        self.lr_scheduler = None
+        has_optimizer = False
+
+        if optimizer or self.optimizer_name():
+            has_optimizer = True
+        # If no parameters given by init default to module parameters
+        if model_parameters is None:
+            model_parameters = self.module.parameters()
+
+        # Convert model parameters from generator to list
+        if not isinstance(model_parameters, list):
+            model_parameters = list(model_parameters)
+
+        if has_optimizer:
+            self._configure_optimizer(optimizer, model_parameters)
+            self._configure_lr_scheduler(lr_scheduler)
+            self._report_progress(0)
+        elif self.zero_optimization():
+            # no optim selected but zero is enabled
+            self.optimizer = self._configure_zero_optimizer(optimizer=None)
+        elif self.bfloat16_enabled():
+            self.optimizer = self._configure_bf16_optimizer(optimizer=None)
+
+        # Hook optimizer for snip_momentum pruning
+        if hasattr(model, 'pruners'):
+            from deepspeed.compression.helper import rewrite_optimizer_step
+            self.optimizer.pruners = model.pruners
+            rewrite_optimizer_step(self.optimizer)
+
+        # Bookkeeping for sparse support
+        self.sparse_tensor_module_names = set()
+        # if self.sparse_gradients_enabled():
+        for name, module in self.module.named_modules():
+            if isinstance(module, (torch.nn.Embedding, torch.nn.EmbeddingBag)) and self.sparse_gradients_enabled():
+                self.sparse_tensor_module_names.add(name + ".weight")
+                logger.info("Will convert {} to sparse tensor during training".format(name))
+
+        self.save_non_zero_checkpoint = False
+        self.save_zero_checkpoint = False
+        if not isinstance(self.optimizer, DeepSpeedZeRoOffload):
+            self._configure_checkpointing(dist_init_required)
+
+        if self.eigenvalue_enabled():
+            self.eigenvalue = self._configure_eigenvalue()
+
+        if self.pld_enabled():
+            self.progressive_layer_drop = self._configure_progressive_layer_drop()
+
+        if self.curriculum_enabled_legacy():
+            self.curriculum_scheduler_legacy = self._configure_curriculum_scheduler_legacy()
+
+        if self.random_ltd_enabled():
+            random_ltd_config = self.random_ltd_config()
+            random_ltd_config[RANDOM_LTD_GLOBAL_BATCH_SIZE] = self.train_batch_size()
+            random_ltd_config[RANDOM_LTD_MICRO_BATCH_SIZE] = self.train_micro_batch_size_per_gpu()
+            self.random_ltd_scheduler = self._configure_random_ltd_scheduler(random_ltd_config)
+
+        # Engine timers
+
+        self.engine_timers = EngineTimers(enable_micro_timers=self.wall_clock_breakdown(),
+                                          enable_global_timers=self.wall_clock_breakdown()
+                                          or self.flops_profiler_enabled())
+
+        if self.global_rank == 0:
+            self._config.print("DeepSpeedEngine configuration")
+            if self.dump_state():
+                print_configuration(self, "DeepSpeedEngine")
+
+        # Use torch (un)flatten ops
+        self.flatten = _flatten_dense_tensors
+        self.unflatten = _unflatten_dense_tensors
+
+    def destroy(self):
+        if self.optimizer is not None and hasattr(self.optimizer, 'destroy'):
+            self.optimizer.destroy()
+
+    def _get_model_parameters(self):
+        if self.autotuning_profile_model_info():
+            self.autotuning_model_info = {}
+            num_params = 0
+            trainable_num_params = 0
+
+            for p in self.module.parameters():
+                # since user code might call deepspeed.zero.Init() before deepspeed.initialize(), need to check the attribute to check if the parameter is partitioned in zero 3 already or not
+                n = 0
+                if hasattr(p, "ds_tensor"):  # if the parameter is partitioned in zero 3
+                    n += p.ds_numel
+                else:  # if the parameter is not partitioned in zero 3 yet
+                    n += p.numel()
+                num_params += n
+                if p.requires_grad:
+                    trainable_num_params += n
+            if self.global_rank == 0:
+                self.autotuning_model_info["num_params"] = num_params * self.mp_world_size
+                self.autotuning_model_info["trainable_num_params"] = trainable_num_params * self.mp_world_size
+
+            logger.info(f"model parameter = {num_params}")
+
+    def get_batch_info(self):
+        """Get all training batch related settings.
+        Returns:
+            train_batch_size (int): The effective training batch size. This is the amount of data
+                samples that leads to one step of model update.
+            train_micro_batch_size_per_gpu (int): Batch size to be processed by one GPU in one
+                step (without gradient accumulation).
+            gradient_accumulation_steps (int): Number of training steps to accumulate gradients
+                before averaging and applying them.
+        """
+        return (
+            self.train_batch_size,
+            self.train_micro_batch_size_per_gpu,
+            self.gradient_accumulation_steps,
+        )
+
+    def set_train_batch_size(self, train_batch_size):
+        """Adjust the global batch size by increasing or decreasing the number of
+        micro-batches (i.e., gradient accumulation steps). The size of each micro-batch
+        (i.e., ``train_micro_batch_size_per_gpu``) is not changed.
+        Args:
+            train_batch_size (int): The new global batch size for training.
+        Raises:
+            ValueError: if ``train_batch_size`` is not divisible by the
+                configured micro-batch size and data parallelism.
+        """
+        if train_batch_size % (self.train_micro_batch_size_per_gpu() * self.dp_world_size) != 0:
+            #print(f'{train_batch_size=} {self.train_micro_batch_size_per_gpu()=} {self.dp_world_size=}')
+            raise ValueError(f'Train batch size must be divisible by micro-batch data parallelism')
+        new_gas = train_batch_size // (self.train_micro_batch_size_per_gpu() * self.dp_world_size)
+        # overwrite config
+        self._config.train_batch_size = train_batch_size
+        self._config.gradient_accumulation_steps = new_gas
+
+    def set_train_micro_batch_size(self, micro_batch_size):
+        """Adjust the micro batch size(i.e., the micro batch size in every data parallel group),
+        while keep the gradient accumulation steps the same.
+        Args:
+            micro_batch_size (int): The new micro batch size for training.
+        """
+        # overwrite config
+        new_global_batch_size = micro_batch_size * self._config.gradient_accumulation_steps * self.dp_world_size
+        self._config.train_batch_size = new_global_batch_size
+        self._config.train_micro_batch_size_per_gpu = micro_batch_size
+
+    def set_data_post_process_func(self, post_process_func):
+        if self.training_dataloader is not None:
+            self.training_dataloader.post_process_func = post_process_func
+
+    def set_custom_curriculum_learning_schedule(self, schedule_func_dict):
+        if self.training_dataloader is not None and self.curriculum_learning_enabled():
+            self.training_dataloader.data_sampler.set_custom_curriculum_learning_schedule(schedule_func_dict)
+
+    def get_global_grad_norm(self) -> float:
+        """Return the 2-norm of all gradients. If there is model parallelism,
+        the norm will be global.
+        The computed norm will be cached and reused until the next step() pass.
+        .. note::
+            In the presence of model parallelism, this is a collective call
+            and acts as a barrier among ``mpu.get_model_parallel_group()``.
+        Returns:
+            float: norm
+        """
+        return self._global_grad_norm
+
+    def __getattr__(self, name):
+        """
+        Pass through attributes defined in the model if they are not overridden by ds-engine.
+        """
+
+        _module = {}
+        if "module" in self.__dict__:
+            _module = self.__dict__['module']
+        if name in dir(self):
+            return getattr(self, name)
+        elif name in dir(_module):
+            return getattr(_module, name)
+        else:
+            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def checkpoint_tag_validation_enabled(self):
+        return self._config.checkpoint_tag_validation_enabled
+
+    def checkpoint_tag_validation_fail(self):
+        return self._config.checkpoint_tag_validation_fail
+
+    def elasticity_enabled(self):
+        return self._config.elasticity_enabled
+
+    def is_elastic_model_parallel_supported(self):
+        if self.elasticity_enabled():
+            # Add code for finding number of GPUs per node automatically
+            if self._config.num_gpus_per_node % self._config.elastic_model_parallel_size == 0:
+                return True
+            else:
+                return False
+
+    def pld_enabled(self):
+        return self._config.pld_enabled
+
+    def pld_params(self):
+        return self._config.pld_params
+
+    def pld_theta(self):
+        return self.pld_params()[PLD_THETA]
+
+    def pld_gamma(self):
+        return self.pld_params()[PLD_GAMMA]
+
+    def eigenvalue_enabled(self):
+        return self._config.eigenvalue_enabled
+
+    def eigenvalue_verbose(self):
+        return self._config.eigenvalue_verbose
+
+    def eigenvalue_max_iter(self):
+        return self._config.eigenvalue_max_iter
+
+    def eigenvalue_tol(self):
+        return self._config.eigenvalue_tol
+
+    def eigenvalue_stability(self):
+        return self._config.eigenvalue_stability
+
+    def eigenvalue_gas_boundary_resolution(self):
+        return self._config.eigenvalue_gas_boundary_resolution
+
+    def eigenvalue_layer_name(self):
+        return self._config.eigenvalue_layer_name
+
+    def eigenvalue_layer_num(self):
+        return self._config.eigenvalue_layer_num
+
+    def curriculum_enabled_legacy(self):
+        return self._config.curriculum_enabled_legacy
+
+    def curriculum_params_legacy(self):
+        return self._config.curriculum_params_legacy
+
+    def data_efficiency_enabled(self):
+        return self._config.data_efficiency_enabled
+
+    def data_efficiency_config(self):
+        return self._config.data_efficiency_config
+
+    def data_sampling_enabled(self):
+        return self._config.data_efficiency_config[DATA_SAMPLING][DATA_SAMPLING_ENABLED]
+
+    def data_sampling_config(self):
+        return self._config.data_efficiency_config[DATA_SAMPLING]
+
+    def curriculum_learning_enabled(self):
+        return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING][CURRICULUM_LEARNING_ENABLED]
+
+    def curriculum_learning_config(self):
+        return self._config.data_efficiency_config[DATA_SAMPLING][CURRICULUM_LEARNING]
+
+    def random_ltd_enabled(self):
+        return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD][RANDOM_LTD_ENABLED]
+
+    def random_ltd_config(self):
+        return self._config.data_efficiency_config[DATA_ROUTING][RANDOM_LTD]
+
+    def random_ltd_initialize(self):
+        assert self.random_ltd_enabled()
+        random_ltd_config = self.random_ltd_config()
+        random_ltd_queue = deque([x for x in sorted(random_ltd_config[RANDOM_LTD_LAYER_ID])])
+        count = 0
+        for name, layer in self.module.named_modules():
+            if isinstance(layer, RandomLayerTokenDrop):
+                if len(random_ltd_queue) != 0 and str(random_ltd_queue[0]) in name:  ###[1,2,3]
+                    layer.init_config(random_ltd_config, self.random_ltd_scheduler, count)
+                    random_ltd_queue.popleft()
+                    count += 1
+
+        if random_ltd_config[RANDOM_LTD_LAYER_NUM] != count:
+            raise ValueError(f'random_ltd_layer_num {random_ltd_config[RANDOM_LTD_LAYER_NUM]} must be \
+                equivalent to the len of random_ltd_layer_id {count}')
+
+        if random_ltd_config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]:
+            assert self.client_lr_scheduler is None
+            raise ValueError(f'not yet support')
+            #self.lr_scheduler = lr_schedules.WarmupLayerTokenDecayLR(self.optimizer, self.random_ltd_scheduler)
+
+    def wall_clock_breakdown(self):
+        return self._config.wall_clock_breakdown
+
+    def flops_profiler_enabled(self):
+        return self._config.flops_profiler_config.enabled or self.autotuning_enabled()
+
+    def flops_profiler_recompute_fwd_factor(self):
+        return self._config.flops_profiler_config.recompute_fwd_factor
+
+    def flops_profiler_profile_step(self):
+        step = self._config.flops_profiler_config.profile_step
+        if self._config.autotuning_config.enabled:
+            step = self.autotuning_start_profile_step()
+        return step
+
+    def flops_profiler_module_depth(self):
+        return self._config.flops_profiler_config.module_depth
+
+    def flops_profiler_top_modules(self):
+        return self._config.flops_profiler_config.top_modules
+
+    def flops_profiler_detailed(self):
+        if self._config.autotuning_config.enabled:
+            return False
+        return self._config.flops_profiler_config.detailed
+
+    def flops_profiler_output_file(self):
+        return self._config.flops_profiler_config.output_file
+
+    def memory_breakdown(self):
+        return self._config.memory_breakdown
+
+    def autotuning_enabled(self):
+        return self._config.autotuning_config.enabled
+
+    def autotuning_start_profile_step(self):
+        return self._config.autotuning_config.start_profile_step
+
+    def autotuning_end_profile_step(self):
+        return self._config.autotuning_config.end_profile_step
+
+    def autotuning_metric_path(self):
+        path = self._config.autotuning_config.metric_path
+        if not path:
+            path = os.path.join(os.getcwd(), "autotuning_metric.json")
+        return path
+
+    def autotuning_model_info_path(self):
+        path = self._config.autotuning_config.model_info_path
+        if not path:
+            path = os.path.join(os.getcwd(), "autotuning_model_info.json")
+        return path
+
+    def autotuning_metric(self):
+        return self._config.autotuning_config.metric
+
+    def autotuning_profile_model_info(self):
+        return self.autotuning_enabled(
+        ) and self._config.autotuning_config.model_info and self._config.autotuning_config.model_info.get(
+            "profile", False)
+
+    def sparse_gradients_enabled(self):
+        return self._config.sparse_gradients_enabled
+
+    def train_batch_size(self):
+        return self._config.train_batch_size
+
+    def train_micro_batch_size_per_gpu(self):
+        return self._config.train_micro_batch_size_per_gpu
+
+    def optimizer_name(self):
+        return (self.client_optimizer.__class__.__name__ if self.client_optimizer else self._config.optimizer_name)
+
+    def optimizer_params(self):
+        return self._config.optimizer_params
+
+    def optimizer_legacy_fusion(self):
+        return self._config.optimizer_legacy_fusion
+
+    def scheduler_name(self):
+        return self._config.scheduler_name
+
+    def scheduler_params(self):
+        return self._config.scheduler_params
+
+    def quantize_training(self):
+        return (
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS]
+            [WEIGHT_QUANTIZE_IN_FORWARD_ENABLED],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ENABLED],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_GROUPS],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS]
+            [WEIGHT_QUANTIZE_FP16_MIXED_QUANTIZE],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_CHANGE_RATIO],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_TYPE],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_ROUNDING],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_VERBOSE],
+            self._config.compression_config[WEIGHT_QUANTIZATION][SHARED_PARAMETERS][WEIGHT_QUANTIZE_KERNEL],
+        )
+
+    def zero_optimization(self):
+        return self._config.zero_enabled
+
+    def zero_allow_untested_optimizer(self):
+        return self._config.zero_allow_untested_optimizer
+
+    def zero_force_ds_cpu_optimizer(self):
+        return self._config.zero_force_ds_cpu_optimizer
+
+    def zero_reduce_scatter(self):
+        return self._config.zero_config.reduce_scatter
+
+    def zero_overlap_comm(self):
+        return self._config.zero_config.overlap_comm
+
+    def zero_offload_optimizer(self):
+        return self._config.zero_config.offload_optimizer
+
+    def zero_offload_param(self):
+        return self._config.zero_config.offload_param
+
+    def zero_use_cpu_optimizer(self):
+        if self._config.zero_config.offload_optimizer is not None:
+            return self._config.zero_config.offload_optimizer.device in [OffloadDeviceEnum.cpu, OffloadDeviceEnum.nvme]
+        return False
+
+    def zero_cpu_offload(self):
+        if self._config.zero_config.offload_optimizer is not None:
+            return self._config.zero_config.offload_optimizer.device == OffloadDeviceEnum.cpu
+        return False
+
+    def zero_partial_offload(self):
+        return getattr(self._config.zero_config.offload_optimizer, "ratio", 1.0)
+
+    def zero_sub_group_size(self):
+        return self._config.zero_config.sub_group_size
+
+    def zero_optimization_stage(self):
+        return self._config.zero_optimization_stage
+
+    def mics_shard_size(self):
+        return self._config.mics_shard_size
+
+    def zero_reduce_bucket_size(self):
+        return self._config.zero_config.reduce_bucket_size
+
+    def zero_multi_rank_bucket_allreduce(self):
+        return self._config.zero_config.use_multi_rank_bucket_allreduce
+
+    def zero_allgather_bucket_size(self):
+        return self._config.zero_config.allgather_bucket_size
+
+    def zero_optimization_partition_gradients(self):
+        return self.zero_optimization_stage() >= ZeroStageEnum.gradients
+
+    def zero_optimization_partition_weights(self):
+        return self.zero_optimization_stage() >= ZeroStageEnum.weights
+
+    def is_first_weights_partition_group(self):
+        ret = True if self.mics_shard_size() < 0 \
+            and self.zero_optimization_partition_weights() else False
+        if self.mics_shard_size() > 0 and self.global_rank < self.mics_shard_size():
+            ret = True
+        return ret
+
+    def zero_contiguous_gradients(self):
+        return self._config.zero_config.contiguous_gradients
+
+    def zero_load_from_fp32_weights(self):
+        return self._config.zero_config.load_from_fp32_weights
+
+    def zero_elastic_checkpoint(self):
+        return self._config.zero_config.elastic_checkpoint
+
+    def zero_max_live_parameters(self):
+        return self._config.zero_config.max_live_parameters
+
+    def zero_max_reuse_distance(self):
+        return self._config.zero_config.max_reuse_distance
+
+    def zero_prefetch_bucket_size(self):
+        return self._config.zero_config.prefetch_bucket_size
+
+    def zero_param_persistence_threshold(self):
+        return self._config.zero_config.param_persistence_threshold
+
+    def zero_model_persistence_threshold(self):
+        return self._config.zero_config.model_persistence_threshold
+
+    def zero_gather_16bit_weights_on_model_save(self):
+        return self._config.zero_config.gather_16bit_weights_on_model_save
+
+    def zero_grad_hooks(self):
+        return self._config.zero_config.grad_hooks
+
+    def zero_legacy_stage1(self):
+        return self._config.zero_config.legacy_stage1
+
+    def zero_ignore_unused_parameters(self):
+        return self._config.zero_config.ignore_unused_parameters
+
+    def graph_harvesting(self):
+        return self._config.graph_harvesting
+
+    def fp16_enabled(self):
+        return self._config.fp16_enabled
+
+    def bfloat16_enabled(self):
+        return self._config.bfloat16_enabled
+
+    def fp16_master_weights_and_gradients(self):
+        return self._config.fp16_master_weights_and_gradients
+
+    def amp_enabled(self):
+        return self._config.amp_enabled
+
+    def amp_params(self):
+        return self._config.amp_params
+
+    def fp16_auto_cast(self):
+        return self._config.fp16_auto_cast
+
+    def loss_scale(self):
+        return self._config.loss_scale
+
+    def gradient_accumulation_steps(self):
+        return self._config.gradient_accumulation_steps
+
+    def use_node_local_storage(self):
+        return self._config.use_node_local_storage
+
+    def load_universal_checkpoint(self):
+        return self._config.load_universal_checkpoint
+
+    @property
+    def communication_data_type(self):
+        res = self._config.communication_data_type
+        if res is not None:
+            return res
+
+        if self.fp16_enabled():
+            return torch.float16
+
+        if self.bfloat16_enabled():
+            return torch.bfloat16
+
+        return torch.float32
+
+    @communication_data_type.setter
+    def communication_data_type(self, value):
+        self._config.communication_data_type = value
+
+    def postscale_gradients(self):
+        return not self._config.prescale_gradients
+
+    def gradient_predivide_factor(self):
+        return self._config.gradient_predivide_factor
+
+    def steps_per_print(self):
+        return self._config.steps_per_print
+
+    def zero_allgather_partitions(self):
+        return self._config.zero_config.allgather_partitions
+
+    def zero_round_robin_gradients(self):
+        return self._config.zero_config.round_robin_gradients
+
+    def zero_hpz_partition_size(self):
+        return self._config.zero_config.zero_hpz_partition_size
+
+    def zero_quantized_weights(self):
+        return self._config.zero_config.zero_quantized_weights
+
+    def zero_quantized_nontrainable_weights(self):
+        return self._config.zero_config.zero_quantized_nontrainable_weights
+
+    def zero_quantized_gradients(self):
+        return self._config.zero_config.zero_quantized_gradients
+
+    def dump_state(self):
+        return self._config.dump_state
+
+    def gradient_clipping(self):
+        return self._config.gradient_clipping
+
+    def dynamic_loss_scale(self):
+        return self._config.loss_scale == 0
+
+    def initial_dynamic_scale(self):
+        return self._config.initial_dynamic_scale
+
+    def dynamic_loss_scale_args(self):
+        return self._config.dynamic_loss_scale_args
+
+    def swap_tensor_config(self):
+        return self._config.swap_tensor_config
+
+    def aio_config(self):
+        return self._config.aio_config
+
+    def get_data_types(self):
+        model_dtype = torch.float32
+        if self.fp16_enabled():
+            model_dtype = torch.float16
+        elif self.bfloat16_enabled():
+            model_dtype = torch.bfloat16
+
+        if self._config.grad_accum_dtype is None:
+            if model_dtype == torch.bfloat16 and not self.zero_optimization():
+                grad_accum_dtype = torch.float32
+            else:
+                grad_accum_dtype = model_dtype
+        else:
+            grad_accum_dtype = DtypeEnum(self._config.grad_accum_dtype).value
+
+        return (model_dtype, grad_accum_dtype)
+
+    def _optimizer_has_ckpt_event_prologue(self):
+        return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_prologue')
+
+    def _optimizer_has_ckpt_event_epilogue(self):
+        return self.optimizer is not None and hasattr(self.optimizer, 'checkpoint_event_epilogue')
+
+    def _configure_lr_scheduler(self, client_lr_scheduler):
+        # First check for scheduler in json configuration
+        lr_scheduler = self._scheduler_from_config(self.optimizer)
+        if lr_scheduler:
+            log_dist(f"DeepSpeed using configured LR scheduler = {self.scheduler_name()}", ranks=[0])
+            self.lr_scheduler = lr_scheduler
+        else:
+            if isinstance(client_lr_scheduler, Callable):
+                log_dist('DeepSpeed using client callable to create LR scheduler', ranks=[0])
+                self.lr_scheduler = client_lr_scheduler(self.basic_optimizer)
+            else:
+                log_dist('DeepSpeed using client LR scheduler', ranks=[0])
+                self.lr_scheduler = client_lr_scheduler
+
+        log_dist(f'DeepSpeed LR Scheduler = {self.lr_scheduler}', ranks=[0])
+
+    def _configure_checkpointing(self, dist_init_required):
+        self.checkpoint_engine = TorchCheckpointEngine()
+
+        if self._config is not None and self._config.nebula_config.enabled:
+            try:
+                from deepspeed.runtime.checkpoint_engine.nebula_checkpoint_engine import \
+                    NebulaCheckpointEngine
+                self.checkpoint_engine = NebulaCheckpointEngine(config_params=self._config.nebula_config)
+            except ImportError as err:
+                logger.error(f"No torch_nebula was found! Will fall back to torch.save. Details: {err}")
+                self.checkpoint_engine = TorchCheckpointEngine()
+
+        dp_rank = groups._get_sequence_data_parallel_rank()
+
+        rank = self.local_rank if self.use_node_local_storage() else dp_rank
+
+        # only the first data parallel process needs to store the model checkpoint
+        # if you want to use node local storage this must be done by rank 0 on each
+        # node
+        self.save_non_zero_checkpoint = (rank == 0) or (self.zero_optimization_partition_weights()
+                                                        and self.is_first_weights_partition_group())
+
+        if self.zero_optimization() or self.bfloat16_enabled():
+            param_rank = dist.get_rank(group=self.optimizer.zp_process_group)
+
+            # Only the first parameter parallel process needs to store the
+            # optimizer state checkpoints for zero
+            self.save_zero_checkpoint = param_rank == dp_rank
+
+    def _scheduler_from_config(self, optimizer):
+        scheduler_name = self.scheduler_name()
+        if scheduler_name is not None:
+            if hasattr(lr_schedules, scheduler_name):
+                scheduler = getattr(lr_schedules, scheduler_name)
+            else:
+                assert hasattr(torch.optim.lr_scheduler,
+                               scheduler_name), f"DeepSpeed does not recognize LR scheduler {scheduler_name}"
+
+                scheduler = getattr(torch.optim.lr_scheduler, scheduler_name)
+
+            scheduler_params = self.scheduler_params()
+            instantiated_scheduler = scheduler(optimizer, **scheduler_params)
+            return instantiated_scheduler
+        else:
+            return None
+
+    def _set_distributed_vars(self, args):
+        device_rank = args.device_rank if args is not None and hasattr(args, 'device_rank') else self.local_rank
+        if device_rank >= 0:
+            get_accelerator().set_device(device_rank)
+            self.device = torch.device(get_accelerator().device_name(), device_rank)
+            self.world_size = dist.get_world_size()
+            self.global_rank = dist.get_rank()
+        else:
+            self.world_size = 1
+            self.global_rank = 0
+            self.device = torch.device(get_accelerator().device_name())
+
+    # Configure based on command line arguments
+    def _configure_with_arguments(self, args, mpu):
+        # After the distributed backend is initialized we are guaranteed the LOCAL_RANK
+        # environment variable is set. We must align args.local_rank to this value for
+        # backwards compatibility with scripts relying on [args|self].local_rank containing
+        # the correct local rank info. _do_args_sanity_check will ensure this is the case.
+
+        if "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ:
+            ompi_local_rank = os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK")
+            local_rank = os.environ.get('LOCAL_RANK', ompi_local_rank)
+            assert ompi_local_rank == local_rank, f"LOCAL_RANK ({local_rank}) != OMPI_COMM_WORLD_LOCAL_RANK ({ompi_local_rank}), " \
+                "not sure how to proceed as we're seeing conflicting local rank info."
+            os.environ['LOCAL_RANK'] = local_rank
+
+        self.local_rank = int(os.environ['LOCAL_RANK'])
+        if hasattr(args, 'local_rank'):
+            args.local_rank = self.local_rank
+
+    # Validate command line arguments
+    def _do_args_sanity_check(self, args):
+        assert "LOCAL_RANK" in os.environ or "OMPI_COMM_WORLD_LOCAL_RANK" in os.environ, "DeepSpeed requires the LOCAL_RANK environment " \
+            "variable, it is set by the deepspeed launcher, deepspeed.init_distributed, or the torch's launcher. If using a " \
+            "different launcher please ensure LOCAL_RANK is set prior to initializing deepspeed."
+
+        if hasattr(args, 'local_rank') and args.local_rank is not None:
+            assert isinstance(args.local_rank,
+                              int), f"args.local_rank of {args.local_rank} is an unknown type {type(args.local_rank)}"
+            if args.local_rank >= 0:
+                env_local_rank = int(os.environ.get("LOCAL_RANK"))
+                assert (
+                    env_local_rank == args.local_rank
+                ), f"Mismatch in local rank setting, args.local_rank={args.local_rank} but env['LOCAL_RANK']={env_local_rank}."
+
+    def _is_supported_optimizer(self, optimizer_name):
+        return (optimizer_name in DEEPSPEED_OPTIMIZERS or getattr(torch.optim, optimizer_name, None) is not None)
+
+    def _supported_optims(self):
+        FairseqOptimizer = None
+        try:
+            from fairseq.optim.fairseq_optimizer import FairseqOptimizer
+        except ImportError:
+            pass
+
+        expected_optim_types = [Optimizer]
+        if FairseqOptimizer:
+            # fairseq optims are not torch.optim objects
+            expected_optim_types.append(FairseqOptimizer)
+        return expected_optim_types
+
+    # Validate configuration based on command line arguments
+    def _do_sanity_check(self):
+        expected_optim_types = self._supported_optims()
+        expected_optim_types += [type(None), Callable]
+        assert isinstance(self.client_optimizer, tuple(expected_optim_types)), \
+            f'Client Optimizer is of unexpected type {type(self.client_optimizer)}'
+
+        if not self.client_optimizer:
+            if self.optimizer_name() is not None:
+                assert self._is_supported_optimizer(
+                    self.optimizer_name()), "{} is not a supported DeepSpeed Optimizer".format(self.optimizer_name())
+
+        if (self.optimizer_name() == LAMB_OPTIMIZER or self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER):
+            assert (self.dynamic_loss_scale()), "DeepSpeed {} optimizer requires dynamic loss scaling".format(
+                self.optimizer_name())
+
+        # Detect invalid combinations of client optimizer and client scheduler
+        if isinstance(self.client_lr_scheduler, _LRScheduler):
+            assert isinstance(self.client_optimizer, Optimizer), \
+                f'Client Optimizer (type = {type(self.client_optimizer)} is not instantiated but Client LR Scheduler is instantiated'
+
+    def _broadcast_model(self):
+
+        def is_replicated(p):
+            if hasattr(p, "ds_status") and p.ds_status is not ZeroParamStatus.AVAILABLE:
+                return False
+            return True
+
+        for p in self.module.parameters():
+            # Broadcast the model for different parameters
+            if is_moe_param(p):
+                if torch.is_tensor(p) and is_replicated(p):
+                    dist.broadcast(p,
+                                   groups._get_expert_broadcast_src_rank(p.group_name),
+                                   group=self.expert_data_parallel_group[p.group_name])
+            else:
+                if torch.is_tensor(p) and is_replicated(p):
+                    dist.broadcast(p, groups._get_broadcast_src_rank(), group=self.seq_data_parallel_group)
+
+    @staticmethod
+    def __check_params(model: Module, dtype: torch.dtype) -> None:
+        return
+        if not all(param.dtype == dtype for param in model.parameters()) and dist.get_rank() == 0:
+            raise ValueError(f"{dtype} is enabled but the following parameters have dtype that is "
+                             f"not {dtype}: "
+                             f"{[(n, p.dtype) for n, p in model.named_parameters() if p.dtype != dtype]}")
+
+    def _set_client_model(self, model):
+        # register client model in _modules so that nn.module methods work correctly
+        modules = self.__dict__.get('_modules')
+        modules['module'] = model
+        # register module attribute in engine but avoid getattr
+        self.__dict__['module'] = model
+
+    def _configure_distributed_model(self, model):
+        self._set_client_model(model)
+        is_zero_init_model = self.zero_optimization_partition_weights() and any(
+            [hasattr(param, "ds_id") for param in self.module.parameters()])
+
+        if self.fp16_enabled():
+            if is_zero_init_model:
+                self.__check_params(self.module, torch.half)
+            self.module.half()
+        elif self.bfloat16_enabled():
+            if is_zero_init_model:
+                self.__check_params(self.module, torch.bfloat16)
+            self.module.bfloat16()
+        else:
+            self.__check_params(self.module, torch.float)
+
+        # zero.Init() handles device placement of model
+        if not (self.dont_change_device or is_zero_init_model):
+            self.module.to(self.device)
+
+        # MoE related initialization
+        for _, module in self.module.named_modules():
+            if isinstance(module, MoE):
+                self.has_moe_layers = True
+                self.num_experts.append(module.num_experts)
+
+        if self.has_moe_layers:
+            for _, module in self.module.named_modules():
+                if isinstance(module, TopKGate):
+                    self.gate_modules.append(module)
+                    if self.wall_clock_breakdown():
+                        module.wall_clock_breakdown = True
+                if isinstance(module, MOELayer):
+                    self.moe_layers.append(module)
+                    if self.wall_clock_breakdown():
+                        module.wall_clock_breakdown = True
+
+        # Pass the mpu from here to groups. For subsequent use, just query groups
+        if self.mpu is not None:
+            groups.mpu = self.mpu
+
+        # Set deepspeed parallelism spec. for the model including expert parallelism
+        for _, module in self.module.named_modules():
+            if hasattr(module, 'set_deepspeed_parallelism'):
+                module.set_deepspeed_parallelism(self._config.use_data_before_expert_parallel_)
+
+        # Query the groups module to get information about various parallel groups
+        self.local_all_to_all_group = None
+        if self.zero_quantized_gradients():
+            log_dist("Using quantized gradients", ranks=[0])
+            self.local_all_to_all_group = groups._get_local_all_to_all_group()
+        self.data_parallel_group = groups._get_data_parallel_group()
+        self.dp_world_size = groups._get_data_parallel_world_size()
+        self.zp_world_size = zp_manager.zp_size
+        self.seq_data_parallel_group = groups._get_sequence_data_parallel_group()
+        self.seq_dp_world_size = groups._get_sequence_data_parallel_world_size()
+        self.mp_world_size = groups._get_model_parallel_world_size()
+        self.expert_parallel_group = groups._get_expert_parallel_group_dict()
+        self.expert_data_parallel_group = groups._get_expert_data_parallel_group_dict()
+        self.sequence_parallel_size = groups._get_sequence_parallel_world_size()
+        if self.sequence_parallel_size > 1:
+            self.communication_data_type = self._config.seq_parallel_communication_data_type
+
+        if not (self.amp_enabled() or is_zero_init_model):
+            self._broadcast_model()
+
+    # check if parameters are duplicated in optimizer param_groups
+    def _check_for_duplicates(self, optimizer):
+        for name, param in self.module.named_parameters():
+            param_id = id(param)
+
+            def ids_list(group):
+                return [id(param) for param in group]
+
+            occurrence = sum([
+                ids_list(group['params']).count(param_id) if param_id in ids_list(group['params']) else 0
+                for group in optimizer.param_groups
+            ])
+            assert occurrence <= 1, f"Parameter with name: {name} occurs multiple times in optimizer.param_groups. Make sure it only appears once to prevent undefined behavior."
+
+    def _do_optimizer_sanity_check(self, basic_optimizer):
+        model_dtype, grad_accum_dtype = self.get_data_types()
+        zero_enabled = self.zero_optimization()
+        amp_enabled = self.amp_enabled()
+        # config based assertions
+        assert (
+            not (amp_enabled and zero_enabled)
+        ), "Amp and ZeRO are not currently compatible, please use (legacy) fp16 mode which performs similar to amp opt_mode=O2"
+        if zero_enabled:
+            if not is_zero_supported_optimizer(basic_optimizer):
+                assert (
+                    self.zero_allow_untested_optimizer()
+                ), 'You are using an untested ZeRO Optimizer. Please add <"zero_allow_untested_optimizer": true> in the configuration file to use it.'
+
+                if self.global_rank == 0:
+                    logger.warning("**** You are using ZeRO with an untested optimizer, proceed with caution *****")
+            if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage(
+            ) == 1 and not self.zero_cpu_offload():
+                return BFLOAT16
+            return ZERO_OPTIMIZATION
+        elif amp_enabled:
+            if model_dtype != grad_accum_dtype:
+                raise NotImplementedError(
+                    "Model data type and gradient accumulation data type must be equal to use Amp")
+            if model_dtype == torch.bfloat16 or model_dtype == torch.float16:
+                raise NotImplementedError("Cannot enable both amp with (legacy) fp16 or bfloat16 mode")
+            try:
+                logger.info("Initializing Apex amp from: {}".format(amp.__path__))
+            except NameError:
+                # If apex/amp is available it will be imported above
+                raise RuntimeError("Unable to import apex/amp, please make sure it is installed")
+            return AMP
+        # data type checks
+        elif model_dtype == grad_accum_dtype:
+            if model_dtype == torch.bfloat16:
+                if self.pipeline_parallelism:
+                    logger.warning(
+                        "**** BF16 gradient accumulation is not safe numerically with large number of accumulation steps, proceed with caution *****"
+                    )
+                    return BFLOAT16
+                else:
+                    raise NotImplementedError(
+                        "Bfloat16 wrapper must use a gradient accumulation type of fp32, enable ZeRO to use Bfloat16 gradient accumulation"
+                    )
+            if model_dtype == torch.float16:
+                return FP16
+            # else optimizer_wrapper = None
+        elif model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32:
+            return BFLOAT16
+        else:
+            raise NotImplementedError("unsupported mix of model dtype and gradient accumulation type")
+
+        return None
+
+    # Configure optimizer
+    def _configure_optimizer(self, client_optimizer, model_parameters):
+        if client_optimizer is None:
+            basic_optimizer = self._configure_basic_optimizer(model_parameters)
+            log_dist(f"Using DeepSpeed Optimizer param name {self.optimizer_name()} as basic optimizer", ranks=[0])
+        else:
+            if isinstance(client_optimizer, tuple(self._supported_optims())):
+                basic_optimizer = client_optimizer
+                log_dist('Using client Optimizer as basic optimizer', ranks=[0])
+            else:
+                basic_optimizer = client_optimizer(model_parameters)
+                log_dist('Using client callable to create basic optimizer', ranks=[0])
+
+            if self.zero_use_cpu_optimizer() and not isinstance(basic_optimizer, deepspeed.ops.adam.DeepSpeedCPUAdam):
+                if self.zero_force_ds_cpu_optimizer():
+                    msg = f'You are using ZeRO-Offload with a client provided optimizer ({type(basic_optimizer)}) which in most cases will yield poor performance. Please either use deepspeed.ops.adam.DeepSpeedCPUAdam or set an optimizer in your ds-config (https://www.deepspeed.ai/docs/config-json/#optimizer-parameters). If you really want to use a custom optimizer w. ZeRO-Offload and understand the performance impacts you can also set <"zero_force_ds_cpu_optimizer": false> in your configuration file.'
+                    raise ZeRORuntimeException(msg)
+
+        basic_optimizer.param_groups[:] = [pg for pg in basic_optimizer.param_groups if len(pg["params"]) != 0]
+        log_dist("Removing param_group that has no 'params' in the basic Optimizer", ranks=[0])
+
+        self._check_for_duplicates(basic_optimizer)
+
+        self.basic_optimizer = basic_optimizer
+        log_dist("DeepSpeed Basic Optimizer = {}".format(basic_optimizer.__class__.__name__), ranks=[0])
+
+        optimizer_wrapper = self._do_optimizer_sanity_check(basic_optimizer)
+
+        if optimizer_wrapper == ZERO_OPTIMIZATION:
+            self.optimizer = self._configure_zero_optimizer(basic_optimizer)
+        elif optimizer_wrapper == AMP:
+            amp_params = self.amp_params()
+            log_dist(f"Initializing AMP with these params: {amp_params}", ranks=[0])
+            model, self.optimizer = amp.initialize(self.module, basic_optimizer, **amp_params)
+            self._set_client_model(model)
+            self._broadcast_model()
+            # TODO: maybe need to broadcast experts differently?
+        elif optimizer_wrapper == FP16:
+            self.optimizer = self._configure_fp16_optimizer(basic_optimizer)
+        elif optimizer_wrapper == BFLOAT16:
+            self.optimizer = self._configure_bf16_optimizer(basic_optimizer)
+        else:
+            self.optimizer = basic_optimizer
+
+        log_dist("DeepSpeed Final Optimizer = {}".format(self.optimizer_name()), ranks=[0])
+
+        self.compression_scheduler = self._configure_compression_scheduler()
+        self.quantizer = self._configure_quantization()
+
+    def _configure_basic_optimizer(self, model_parameters):
+        optimizer_parameters = self.optimizer_params()
+        if optimizer_parameters is None:
+            optimizer_parameters = {}
+        # print(optimizer_parameters.keys())
+        if "max_grad_norm" in optimizer_parameters.keys():
+            raise ValueError(
+                "'max_grad_norm' is not supported as an optimizer parameter, please switch to using the deepspeed parameter 'gradient_clipping' see: https://www.deepspeed.ai/docs/config-json/#gradient-clipping for more details"
+            )
+
+        if self.optimizer_name() in [ADAM_OPTIMIZER, ADAMW_OPTIMIZER]:
+            torch_adam = optimizer_parameters.pop(TORCH_ADAM_PARAM, False)
+            adam_w_mode = optimizer_parameters.pop(ADAM_W_MODE, ADAM_W_MODE_DEFAULT)
+
+            # Optimizer name of Adam forces AdamW logic unless adam_w_mode is explicitly set
+            effective_adam_w_mode = self.optimizer_name() == ADAMW_OPTIMIZER or adam_w_mode
+
+            if torch_adam:
+                if not effective_adam_w_mode:
+                    optimizer = torch.optim.Adam(model_parameters, **optimizer_parameters)
+                else:
+                    optimizer = torch.optim.AdamW(model_parameters, **optimizer_parameters)
+            else:
+                if self.zero_use_cpu_optimizer():
+                    from deepspeed.ops.adam import DeepSpeedCPUAdam
+                    optimizer = DeepSpeedCPUAdam(model_parameters,
+                                                 **optimizer_parameters,
+                                                 adamw_mode=effective_adam_w_mode)
+                else:
+                    from deepspeed.ops.adam import FusedAdam
+
+                    optimizer = FusedAdam(
+                        model_parameters,
+                        **optimizer_parameters,
+                        adam_w_mode=effective_adam_w_mode,
+                    )
+
+        elif self.optimizer_name() == ADAGRAD_OPTIMIZER:
+            if self.zero_use_cpu_optimizer():
+                from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad
+                optimizer = DeepSpeedCPUAdagrad(model_parameters, **optimizer_parameters)
+            else:
+                optimizer = torch.optim.Adagrad(model_parameters, **optimizer_parameters)
+        elif self.optimizer_name() == LAMB_OPTIMIZER:
+            from deepspeed.ops.lamb import FusedLamb
+
+            optimizer = FusedLamb(model_parameters, **optimizer_parameters)
+        elif self.optimizer_name() == ONEBIT_ADAM_OPTIMIZER:
+            assert not self.zero_optimization(), "1bit-Adam is not compatible with ZeRO"
+            from deepspeed.runtime.fp16.onebit.adam import OnebitAdam
+
+            optimizer = OnebitAdam(model_parameters, self, **optimizer_parameters)
+            if not self.fp16_enabled():
+                logger.warning(f"Currently the convergence of 1-bit Adam is only verified under FP16")
+        elif self.optimizer_name() == ZERO_ONE_ADAM_OPTIMIZER:
+            assert not self.zero_optimization(), "0/1 Adam is not compatible with ZeRO"
+            from deepspeed.runtime.fp16.onebit.zoadam import ZeroOneAdam
+
+            optimizer = ZeroOneAdam(model_parameters, self, **optimizer_parameters)
+            if not self.fp16_enabled():
+                logger.warning(f'Currently the convergence of 0/1 Adam is only verified under FP16')
+        elif self.optimizer_name() == ONEBIT_LAMB_OPTIMIZER:
+            assert not self.zero_optimization(), "1bit-Lamb is not compatible with ZeRO"
+            from deepspeed.runtime.fp16.onebit.lamb import OnebitLamb
+
+            optimizer = OnebitLamb(model_parameters, self, **optimizer_parameters)
+            if not self.fp16_enabled():
+                logger.warning(f"Currently the convergence of 1-bit Lamb is only verified under FP16")
+        elif self.optimizer_name() == LION_OPTIMIZER:
+            if self.zero_use_cpu_optimizer():
+                from deepspeed.ops.lion import DeepSpeedCPULion
+                optimizer = DeepSpeedCPULion(model_parameters, **optimizer_parameters)
+            else:
+                from deepspeed.ops.lion import FusedLion
+                optimizer = FusedLion(model_parameters, **optimizer_parameters)
+        elif self.optimizer_name() == MUADAM_OPTIMIZER:
+            try:
+                from mup import MuAdam
+            except ImportError:
+                logger.error(f"Install mup to use MuAdam optimizer")
+            optimizer = MuAdam(model_parameters, **optimizer_parameters)
+        elif self.optimizer_name() == MUADAMW_OPTIMIZER:
+            try:
+                from mup import MuAdamW
+            except ImportError:
+                logger.error(f"Install mup to use MuAdamW optimizer")
+            optimizer = MuAdamW(model_parameters, **optimizer_parameters)
+        elif self.optimizer_name() == MUSGD_OPTIMIZER:
+            try:
+                from mup import MuSGD
+            except ImportError:
+                logger.error(f"Install mup to use MuSGD optimizer")
+            optimizer = MuSGD(model_parameters, **optimizer_parameters)
+        else:
+            torch_optimizer = getattr(torch.optim, self.optimizer_name())
+            optimizer = torch_optimizer(model_parameters, **optimizer_parameters)
+        return optimizer
+
+    def _configure_compression_scheduler(self):
+        return compression_scheduler(self.module, self._config.compression_config)
+
+    def _configure_random_ltd_scheduler(self, configs):
+        return RandomLTDScheduler(configs)
+
+    def _configure_quantization(self):
+        (
+            quantize_weight_in_forward,
+            quantize_enabled,
+            q_groups,
+            q_mixed_fp16,
+            q_change_ratio,
+            q_type,
+            q_rounding,
+            q_verbose,
+            use_quantizer_kernel,
+        ) = self.quantize_training()
+        if quantize_enabled and not quantize_weight_in_forward:
+            assert self.fp16_enabled(
+            ), "MoQ (quantize in optimization step) weight quantization is only supported for FP16"
+        quantizer = None
+        if quantize_enabled and not quantize_weight_in_forward:
+            from deepspeed.runtime.quantize import Quantizer
+
+            quantizer = Quantizer(
+                q_groups,
+                q_mixed_fp16,
+                q_change_ratio,
+                q_type,
+                q_rounding,
+                q_verbose,
+                self.eigenvalue_enabled(),
+                use_quantizer_kernel,
+                self.eigenvalue_layer_num() if self.eigenvalue_enabled() else 0,
+            )
+        return quantizer
+
+    def _configure_fp16_optimizer(self, optimizer):
+        initial_dynamic_scale = self.initial_dynamic_scale()
+        dynamic_loss_args = self.dynamic_loss_scale_args()
+        clip_grad = self.gradient_clipping()
+        if APEX_INSTALLED:
+            fused_opts = (apex.optimizers.FusedAdam, FusedAdam)
+        else:
+            fused_opts = FusedAdam
+        if isinstance(optimizer, fused_opts) \
+                or self.optimizer_name() in [ONEBIT_ADAM_OPTIMIZER, ZERO_ONE_ADAM_OPTIMIZER]:
+            if self.dynamic_loss_scale():
+                log_dist(f'Creating fp16 optimizer with dynamic loss scale', ranks=[0])
+                timers = self.timers if self.wall_clock_breakdown() else NoopTimer()
+                optimizer = FP16_Optimizer(
+                    optimizer,
+                    deepspeed=self,
+                    dynamic_loss_scale=True,
+                    initial_dynamic_scale=initial_dynamic_scale,
+                    dynamic_loss_args=dynamic_loss_args,
+                    mpu=self.mpu,
+                    clip_grad=clip_grad,
+                    fused_adam_legacy=self.optimizer_legacy_fusion(),
+                    timers=timers,
+                    has_moe_layers=self.has_moe_layers,
+                )
+            else:
+                log_dist(f'Creating fp16 optimizer with static loss scale: {self.loss_scale()}', ranks=[0])
+                optimizer = FP16_Optimizer(
+                    optimizer,
+                    deepspeed=self,
+                    static_loss_scale=self.loss_scale(),
+                    mpu=self.mpu,
+                    clip_grad=clip_grad,
+                    fused_adam_legacy=self.optimizer_legacy_fusion(),
+                    has_moe_layers=self.has_moe_layers,
+                )
+        else:
+            log_dist(f'Creating fp16 unfused optimizer with dynamic loss scale', ranks=[0])
+            optimizer = FP16_UnfusedOptimizer(
+                optimizer,
+                deepspeed=self,
+                static_loss_scale=self.loss_scale(),
+                dynamic_loss_scale=self.dynamic_loss_scale(),
+                dynamic_loss_args=dynamic_loss_args,
+                mpu=self.mpu,
+                clip_grad=clip_grad,
+                fused_lamb_legacy=self.optimizer_name() == LAMB_OPTIMIZER,
+            )
+
+        return optimizer
+
+    def _configure_bf16_optimizer(self, optimizer):
+        clip_grad = self.gradient_clipping()
+
+        if optimizer is None:
+            optimizer = DummyOptim(list(self.module.parameters()))
+
+        log_dist('Creating BF16 optimizer', ranks=[0])
+
+        timers = self.timers if self.wall_clock_breakdown() else NoopTimer()
+        optimizer = BF16_Optimizer(optimizer,
+                                   self.param_names,
+                                   mpu=self.mpu,
+                                   clip_grad=clip_grad,
+                                   allgather_bucket_size=self.zero_allgather_bucket_size(),
+                                   dp_process_group=self.seq_data_parallel_group,
+                                   timers=timers,
+                                   grad_acc_dtype=self.get_data_types()[1],
+                                   graph_harvesting=self.graph_harvesting())
+
+        return optimizer
+
+    def _configure_zero_optimizer(self, optimizer):
+        zero_stage = self.zero_optimization_stage()
+
+        mics_shard_size = self.mics_shard_size()
+        model_dtype, gradient_accumulation_dtype = self.get_data_types()
+
+        timers = self.timers if self.wall_clock_breakdown() else NoopTimer()
+
+        if optimizer is None:
+            optimizer = DummyOptim(list(self.module.parameters()))
+
+        if self.zero_legacy_stage1():
+            raise Exception(
+                "The deprecated version of ZeRO Stage 1 is not supported in deepspeed >= 0.5.9. Please downgrade to a version less than 0.5.9 if you need to use this deprecated version of ZeRO."
+            )
+
+        if zero_stage <= ZeroStageEnum.gradients:
+            overlap_comm = self.zero_overlap_comm()
+            contiguous_gradients = self.zero_contiguous_gradients()
+            round_robin_gradients = self.zero_round_robin_gradients()
+            assert not isinstance(optimizer, DummyOptim), "zero stage {} requires an optimizer".format(zero_stage)
+
+            log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0])
+            # Overlap and contiguous grads are meaningless in stage 1 and are ignored
+            if zero_stage == ZeroStageEnum.optimizer_states:
+                overlap_comm = False
+                round_robin_gradients = False
+                # Non-MoE requires contiguous grads to be disabled w. stage 1
+                if not self.has_moe_layers:
+                    contiguous_gradients = False
+
+            if isinstance(self.module, PipelineModule):
+                if overlap_comm:
+                    logger.warning("Pipeline parallelism does not support overlapped communication, will be disabled.")
+                    overlap_comm = False
+            optimizer = DeepSpeedZeroOptimizer(
+                optimizer,
+                self.param_names,
+                timers=timers,
+                static_loss_scale=self.loss_scale(),
+                dynamic_loss_scale=self.dynamic_loss_scale(),
+                dynamic_loss_args=self.dynamic_loss_scale_args(),
+                clip_grad=self.gradient_clipping(),
+                contiguous_gradients=contiguous_gradients,
+                reduce_bucket_size=self.zero_reduce_bucket_size(),
+                use_multi_rank_bucket_allreduce=self.zero_multi_rank_bucket_allreduce(),
+                allgather_bucket_size=self.zero_allgather_bucket_size(),
+                dp_process_group=self.seq_data_parallel_group,
+                expert_parallel_group=self.expert_parallel_group if self.has_moe_layers else None,
+                expert_data_parallel_group=self.expert_data_parallel_group if self.has_moe_layers else None,
+                reduce_scatter=self.zero_reduce_scatter(),
+                overlap_comm=overlap_comm,
+                offload_optimizer_config=self.zero_offload_optimizer(),
+                mpu=self.mpu,
+                postscale_gradients=self.postscale_gradients(),
+                gradient_predivide_factor=self.gradient_predivide_factor(),
+                gradient_accumulation_steps=self.gradient_accumulation_steps(),
+                ignore_unused_parameters=self.zero_ignore_unused_parameters(),
+                partition_grads=zero_stage == ZeroStageEnum.gradients,
+                round_robin_gradients=round_robin_gradients,
+                has_moe_layers=self.has_moe_layers,
+                fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(),
+                gradient_accumulation_dtype=gradient_accumulation_dtype,
+                communication_data_type=self.communication_data_type,
+                elastic_checkpoint=self.zero_elastic_checkpoint())
+
+        elif zero_stage == ZeroStageEnum.weights:
+            assert not self.has_moe_layers, "MoE not supported with Stage 3"
+            if isinstance(optimizer, DummyOptim):
+                log_dist("Creating ZeRO Offload", ranks=[0])
+                zero_param_parallel_group = groups._get_zero_param_intra_parallel_group()
+                if self.zero_hpz_partition_size() > 1 and zero_param_parallel_group is None:
+                    self._set_zero_group_parallelism()
+                    zero_param_parallel_group = groups._get_zero_param_intra_parallel_group()
+                optimizer = DeepSpeedZeRoOffload(
+                    self.module,
+                    timers=timers,
+                    ds_config=self.config,
+                    overlap_comm=self.zero_overlap_comm(),
+                    prefetch_bucket_size=self.zero_prefetch_bucket_size(),
+                    max_reuse_distance=self.zero_max_reuse_distance(),
+                    max_live_parameters=self.zero_max_live_parameters(),
+                    param_persistence_threshold=self.zero_param_persistence_threshold(),
+                    model_persistence_threshold=self.zero_model_persistence_threshold(),
+                    offload_param_config=self.zero_offload_param(),
+                    mpu=self.mpu,
+                    zero_param_parallel_group=zero_param_parallel_group,
+                    zero_quantized_weights=self.zero_quantized_weights(),
+                    zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(),
+                )
+            else:
+                log_dist(
+                    f'Creating fp16 ZeRO stage {zero_stage} optimizer,'
+                    f' MiCS is enabled {mics_shard_size>0},'
+                    f' Hierarchical params gather {self._config.mics_hierarchial_params_gather}',
+                    ranks=[0])
+                if mics_shard_size > 0:
+                    return self._return_mics_optimizer(optimizer, timers)
+
+                log_dist(f'Creating {model_dtype} ZeRO stage {zero_stage} optimizer', ranks=[0])
+                from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
+                optimizer = DeepSpeedZeroOptimizer_Stage3(
+                    self.module,
+                    optimizer,
+                    timers=timers,
+                    ds_config=self.config,
+                    static_loss_scale=self.loss_scale(),
+                    dynamic_loss_scale=self.dynamic_loss_scale(),
+                    dynamic_loss_args=self.dynamic_loss_scale_args(),
+                    clip_grad=self.gradient_clipping(),
+                    contiguous_gradients=self.zero_contiguous_gradients(),
+                    reduce_bucket_size=self.zero_reduce_bucket_size(),
+                    prefetch_bucket_size=self.zero_prefetch_bucket_size(),
+                    max_reuse_distance=self.zero_max_reuse_distance(),
+                    max_live_parameters=self.zero_max_live_parameters(),
+                    param_persistence_threshold=self.zero_param_persistence_threshold(),
+                    model_persistence_threshold=self.zero_model_persistence_threshold(),
+                    dp_process_group=self.seq_data_parallel_group,
+                    all2all_process_group=self.local_all_to_all_group,
+                    reduce_scatter=self.zero_reduce_scatter(),
+                    overlap_comm=self.zero_overlap_comm(),
+                    offload_optimizer_config=self.zero_offload_optimizer(),
+                    offload_param_config=self.zero_offload_param(),
+                    sub_group_size=self.zero_sub_group_size(),
+                    offload_ratio=self.zero_partial_offload(),
+                    mpu=self.mpu,
+                    postscale_gradients=self.postscale_gradients(),
+                    gradient_predivide_factor=self.gradient_predivide_factor(),
+                    gradient_accumulation_steps=self.gradient_accumulation_steps(),
+                    aio_config=self.aio_config(),
+                    gradient_accumulation_dtype=gradient_accumulation_dtype,
+                    communication_data_type=self.communication_data_type,
+                    zero_hpz_partition_size=self.zero_hpz_partition_size(),
+                    zero_quantized_weights=self.zero_quantized_weights(),
+                    zero_quantized_nontrainable_weights=self.zero_quantized_nontrainable_weights(),
+                )
+
+        else:
+            raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage))
+
+        return optimizer
+
+    def _return_mics_optimizer(self, basic_optimizer, timers):
+        from deepspeed.runtime.zero.mics import MiCS_Optimizer
+        model_dtype, gradient_accumulation_dtype = self.get_data_types()
+        optimizer = MiCS_Optimizer(self.module,
+                                   basic_optimizer,
+                                   timers=timers,
+                                   ds_config=self.config,
+                                   static_loss_scale=self.loss_scale(),
+                                   dynamic_loss_scale=self.dynamic_loss_scale(),
+                                   dynamic_loss_args=self.dynamic_loss_scale_args(),
+                                   clip_grad=self.gradient_clipping(),
+                                   contiguous_gradients=self.zero_contiguous_gradients(),
+                                   reduce_bucket_size=self.zero_reduce_bucket_size(),
+                                   prefetch_bucket_size=self.zero_prefetch_bucket_size(),
+                                   max_reuse_distance=self.zero_max_reuse_distance(),
+                                   max_live_parameters=self.zero_max_live_parameters(),
+                                   param_persistence_threshold=self.zero_param_persistence_threshold(),
+                                   model_persistence_threshold=self.zero_model_persistence_threshold(),
+                                   dp_process_group=self.seq_data_parallel_group,
+                                   reduce_scatter=self.zero_reduce_scatter(),
+                                   overlap_comm=self.zero_overlap_comm(),
+                                   offload_optimizer_config=self.zero_offload_optimizer(),
+                                   offload_param_config=self.zero_offload_param(),
+                                   sub_group_size=self.zero_sub_group_size(),
+                                   mpu=self.mpu,
+                                   postscale_gradients=self.postscale_gradients(),
+                                   gradient_predivide_factor=self.gradient_predivide_factor(),
+                                   gradient_accumulation_steps=self.gradient_accumulation_steps(),
+                                   aio_config=self.aio_config(),
+                                   gradient_accumulation_dtype=gradient_accumulation_dtype,
+                                   communication_data_type=self.communication_data_type)
+        return optimizer
+
+    def _configure_eigenvalue(self):
+        eigenvalue = Eigenvalue(
+            verbose=self.eigenvalue_verbose(),
+            max_iter=self.eigenvalue_max_iter(),
+            tol=self.eigenvalue_tol(),
+            stability=self.eigenvalue_stability(),
+            gas_boundary_resolution=self.eigenvalue_gas_boundary_resolution(),
+            layer_name=self.eigenvalue_layer_name(),
+            layer_num=self.eigenvalue_layer_num(),
+        )
+
+        return eigenvalue
+
+    def _configure_progressive_layer_drop(self):
+        pld = ProgressiveLayerDrop(theta=self.pld_theta(), gamma=self.pld_gamma())
+
+        return pld
+
+    def _configure_curriculum_scheduler_legacy(self):
+        scheduler = CurriculumScheduler(self.curriculum_params_legacy())
+        return scheduler
+
+    @staticmethod
+    def is_map_style_dataset(obj):
+        return hasattr(obj, "__getitem__") and hasattr(obj, "__len__")
+
+    @staticmethod
+    def is_iterable_style_dataset(obj):
+        return isinstance(obj, torch.utils.data.IterableDataset)  # hasattr(obj, "__iter__") should work as well
+
+    def dataloader_drop_last(self):
+        return self._config.dataloader_drop_last
+
+    def was_step_applied(self) -> bool:
+        """Returns True if the latest ``step()`` produced in parameter updates.
+        Note that a ``False`` return is not an error condition. Steps are frequently
+        no-ops, such as between gradient accumulation boundaries or when overflows
+        occur.
+        Returns:
+            bool: Whether the latest ``step()`` modified model parameters.
+        """
+        return self._step_applied
+
+    def deepspeed_io(self,
+                     dataset,
+                     batch_size=None,
+                     route=ROUTE_TRAIN,
+                     pin_memory=True,
+                     data_sampler=None,
+                     collate_fn=None,
+                     num_local_io_workers=None):
+        if not (self.is_map_style_dataset(dataset) or self.is_iterable_style_dataset(dataset)):
+            raise ValueError("Training data must be a torch Dataset")
+
+        if batch_size is None:
+            batch_size = self.train_micro_batch_size_per_gpu()
+
+        if collate_fn is None:
+            collate_fn = self.collate_fn
+
+        # Currently we only use timer in train route
+        deepspeed_io_timer = None
+        if route == ROUTE_TRAIN:
+            deepspeed_io_timer = self.tput_timer
+
+        # If mpu is provided, forward world size and parallel rank to sampler.
+        data_parallel_world_size = self.dp_world_size
+        data_parallel_rank = self.global_rank
+        if self.mpu is not None:
+            data_parallel_world_size = self.mpu.get_data_parallel_world_size()
+            data_parallel_rank = self.mpu.get_data_parallel_rank()
+
+        if data_sampler is None and (route == ROUTE_PREDICT or route == ROUTE_EVAL):
+            data_sampler = torch.utils.data.DistributedSampler(
+                dataset,
+                num_replicas=data_parallel_world_size,
+                rank=data_parallel_rank,
+                shuffle=False,
+            )
+
+        deepspeed_dataloader_config = {}
+        if self.curriculum_learning_enabled():
+            deepspeed_dataloader_config = {
+                CURRICULUM_LEARNING: self.curriculum_learning_enabled(),
+                DATA_EFFICIENCY: self.data_efficiency_config(),
+                DATA_PARALLEL_GROUP: self.data_parallel_group,
+                GRADIENT_ACCUMULATION_STEPS: self.gradient_accumulation_steps(),
+                GLOBAL_RANK: self.global_rank,
+                DATA_SAMPLING_NUM_WORKERS: self.data_sampling_config()[DATA_SAMPLING_NUM_WORKERS]
+            }
+
+        return DeepSpeedDataLoader(dataset=dataset,
+                                   batch_size=batch_size,
+                                   pin_memory=pin_memory,
+                                   collate_fn=collate_fn,
+                                   local_rank=self.local_rank,
+                                   tput_timer=deepspeed_io_timer,
+                                   num_local_io_workers=num_local_io_workers,
+                                   data_sampler=data_sampler,
+                                   data_parallel_world_size=data_parallel_world_size,
+                                   data_parallel_rank=data_parallel_rank,
+                                   dataloader_drop_last=self.dataloader_drop_last(),
+                                   deepspeed_dataloader_config=deepspeed_dataloader_config)
+
+    def train(self, mode=True):
+        r""""""
+
+        self.warn_unscaled_loss = True
+        self.module.train(mode)
+
+    def eval(self):
+        r""""""
+
+        self.warn_unscaled_loss = True
+        self.module.train(False)
+
+    def _scale_loss_by_gas(self, prescaled_loss):
+        if isinstance(prescaled_loss, torch.Tensor):
+            scaled_loss = prescaled_loss / self.gradient_accumulation_steps()
+        elif isinstance(prescaled_loss, tuple) or isinstance(prescaled_loss, list):
+            scaled_loss = []
+            for l in prescaled_loss:
+                if isinstance(l, torch.Tensor):
+                    scaled_loss.append(l / self.gradient_accumulation_steps())
+                else:
+                    scaled_loss.append(l)
+        else:
+            scaled_loss = prescaled_loss
+            if self.warn_unscaled_loss:
+                logger.warning(f"DeepSpeed unable to scale loss because of type: {type(prescaled_loss)}")
+                self.warn_unscaled_loss = False
+
+        return scaled_loss
+
+    @instrument_w_nvtx
+    def forward(self, *inputs, **kwargs):
+        r"""Execute forward propagation
+        Arguments:
+            *inputs: Variable length input list
+            **kwargs: variable length keyword arguments
+        """
+
+        if self.autotuning_profile_model_info():
+            ma = get_ma_status()
+        else:
+            see_memory_usage("Engine before forward", force=self.memory_breakdown())
+
+        flops_profiler_active = (self.flops_profiler_enabled()
+                                 and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0)
+
+        # used to check quantization happens at step 0!
+        if self.global_steps == 0 and hasattr(self, "compression_scheduler"):
+            self.compression_scheduler.step(step_zero_check=True)
+            if self.quantizer:
+                tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage(
+                ) == 2 else self.optimizer.fp16_groups
+                if self.compression_scheduler.weight_quantization_enabled:
+                    self.quantizer.quantize(
+                        tensor_to_quantize,
+                        (self.optimizer.overflow if self.fp16_enabled() else False),
+                        self.eigenvalue_enabled(),
+                        None,
+                    )
+
+        if flops_profiler_active:
+            self.flops_profiler.start_profile(ignore_list=None)
+
+        if self.module.training:
+            if self.progressive_layer_drop:
+                kwargs.update(self.progressive_layer_drop.get_state())
+
+        if self.__class__.__name__ != "PipelineEngine":
+            # TODO: The above if condition is a HACK since for PipelineEngine
+            # it's difficult to inject argument in forward pass.
+            if self.module.training and self.curriculum_enabled_legacy():
+                self.curriculum_scheduler_legacy.update_difficulty(self.global_steps + 1)
+                if self.curriculum_params_legacy()["curriculum_type"] == "seqlen":
+                    kwargs.update({"curriculum_seqlen": self.curriculum_scheduler_legacy.get_current_difficulty()})
+
+        if self.module.training and self.random_ltd_enabled():
+            self.random_ltd_scheduler.update_seq(self.global_steps)
+
+        if self.zero_optimization_partition_weights():
+            # Enable automated discovery of external parameters by indicating that
+            # we are in a forward pass.
+            for module in self.module.modules():
+                module._parameters._in_forward = True
+
+        self._start_timers(self.engine_timers.forward_timers)
+
+        if self.training_dataloader is None:
+            self.tput_timer.start()
+
+        if self.fp16_auto_cast():
+            inputs = self._cast_inputs_half(inputs)
+        # print(f"RANK[{self.global_rank}] self.fp16_auto_cast() is {self.fp16_auto_cast()}")
+
+        loss = self.module(*inputs, **kwargs)
+
+        # print(f"RANK[{self.global_rank}]'s loss is {loss}")
+
+        if self.zero_optimization_partition_weights():
+            # Disable automated discovery of external parameters
+            for module in self.module.modules():
+                module._parameters._in_forward = False
+
+        self._stop_timers(self.engine_timers.forward_timers)
+
+        if flops_profiler_active:
+            self.flops_profiler.stop_profile()
+
+        if self.autotuning_profile_model_info():
+            activation_mem = get_ma_status() - ma
+            self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem
+            print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path())
+            exit()
+        else:
+            see_memory_usage("Engine after forward", force=self.memory_breakdown())
+        return loss
+
+    def _cast_inputs_half(self, inputs):
+        if isinstance(inputs, (list, tuple)):
+            new_inputs = []
+            for v in inputs:
+                new_inputs.append(self._cast_inputs_half(v))
+            return inputs.__class__(new_inputs)
+        elif isinstance(inputs, dict):
+            new_inputs = {}
+            for k, v in inputs.items():
+                new_inputs[k] = self._cast_inputs_half(v)
+            return new_inputs
+        elif hasattr(inputs, 'half'):
+            return inputs.half()
+        else:
+            return inputs
+
+    def print_forward_breakdown(self, fwd_time):
+        gate_time = 0.0
+        moe_time = 0.0
+        falltoall = 0.0
+        salltoall = 0.0
+
+        for gate in self.gate_modules:
+            #logger.info(f"Individual TopK gate time: {gate.gate_time:.2f} ms")
+            gate_time += gate.gate_time
+
+        for l in self.moe_layers:
+            #logger.info(f"MoE layer; total: {l.time_moe:.2f} ms, first alltoall: {l.time_falltoall:.2f}, second alltoall: {l.time_salltoall:.2f}")
+            moe_time += l.time_moe
+            falltoall += l.time_falltoall
+            salltoall += l.time_salltoall
+
+        # TODO: Allreduce/average them across ranks for more accurate timing.
+
+        # if deepspeed.comm.get_rank() == 0:
+        log_dist(
+            f"time (ms) | fwd: {fwd_time:.2f} (fwd_moe: {moe_time:.2f}, 1st_a2a: {falltoall:.2f}, 2nd_a2a: {salltoall:.2f}, top_k: {gate_time:.2f})",
+            ranks=[0])
+
+    @instrument_w_nvtx
+    def allreduce_gradients(self, bucket_size=MEMORY_OPT_ALLREDUCE_SIZE):
+        assert not (self.bfloat16_enabled() and self.pipeline_parallelism), \
+            f'allreduce_gradients() is not valid when bfloat+pipeline_parallelism is enabled'
+
+        # Pass (PP) gas boundary flag to optimizer (required for zero)
+        self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
+        # ZeRO stage >= 2 communicates during non gradient accumulation boundaries as well
+        if self.zero_optimization_partition_gradients():
+            self.optimizer.overlapping_partition_gradients_reduce_epilogue()
+
+        # Communicate only at gradient accumulation boundaries
+        elif self.is_gradient_accumulation_boundary():
+            if self.zero_optimization_stage() == ZeroStageEnum.optimizer_states and hasattr(
+                    self.optimizer, 'reduce_gradients'):
+                self.optimizer.reduce_gradients(pipeline_parallel=self.pipeline_parallelism)
+            else:
+                self.buffered_allreduce_fallback(elements_per_buffer=bucket_size)
+
+    @instrument_w_nvtx
+    def backward(self, loss, allreduce_gradients=True, release_loss=False, retain_graph=False, scale_wrt_gas=True):
+        r"""Execute backward pass on the loss
+        Arguments:
+            loss: Torch tensor on which to execute backward propagation
+            allreduce_gradients: is deprecated, ignored, and will soon be removed'
+            retain_graph: bool, default: false
+                forward on user defined choice of retain_graph
+        """
+
+        see_memory_usage("Engine before backward", force=self.memory_breakdown())
+
+        if self.scale_wrt_gas is not None:
+            scale_wrt_gas = self.scale_wrt_gas
+
+        if not allreduce_gradients:
+            logger.warning(f"Argument `allreduce_gradients` is deprecated, ignored, and will soon be removed")
+
+        # scale loss w.r.t. gradient accumulation if needed
+        if self.gradient_accumulation_steps() > 1 and scale_wrt_gas:
+            loss = self._scale_loss_by_gas(loss.float())
+
+        # Log training loss
+        self.losses += loss.mean().item()
+        if self.monitor.enabled:
+            if self.is_gradient_accumulation_boundary():
+                if self.global_rank == 0:
+                    self.summary_events = [(
+                        f"Train/Samples/train_loss",
+                        self.losses,
+                        self.global_samples,
+                    )]
+                    self.monitor.write_events(self.summary_events)
+
+        self._start_timers(self.engine_timers.backward_timers)
+
+        assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \
+            "must provide optimizer during init in order to use backward"
+
+        self._start_timers(self.engine_timers.backward_inner_timers)
+
+        if self.zero_optimization():
+            self.optimizer.is_gradient_accumulation_boundary = self.is_gradient_accumulation_boundary()
+            self.optimizer.backward(loss, retain_graph=retain_graph)
+        elif self.amp_enabled():
+            # AMP requires delaying unscale when inside gradient accumulation boundaries
+            # https://nvidia.github.io/apex/advanced.html#gradient-accumulation-across-iterations
+            delay_unscale = not self.is_gradient_accumulation_boundary()
+            with amp.scale_loss(loss, self.optimizer, delay_unscale=delay_unscale) as scaled_loss:
+                scaled_loss.backward(retain_graph=retain_graph)
+        elif self.fp16_enabled():
+            if self.eigenvalue_enabled():
+                self.optimizer.backward(loss, create_graph=True, retain_graph=True)
+            else:
+                self.optimizer.backward(loss, retain_graph=retain_graph)
+        elif self.bfloat16_enabled():
+            self.optimizer.backward(loss)
+        else:
+            if self.eigenvalue_enabled():
+                loss.backward(create_graph=True, retain_graph=True)
+            else:
+                loss.backward(retain_graph=retain_graph)
+
+        self._stop_timers(self.engine_timers.backward_inner_timers)
+
+        self._start_timers(self.engine_timers.backward_reduce_timers)
+
+        if allreduce_gradients and self.enable_backward_allreduce:
+            # Traditional code path that allreduces the module parameter grads
+            self.allreduce_gradients()
+
+        self._stop_timers(self.engine_timers.backward_reduce_timers)
+
+        self._stop_timers(self.engine_timers.backward_timers)
+
+        if release_loss:
+            # loss.data = None
+            pass
+
+        see_memory_usage("Engine after backward", force=self.memory_breakdown())
+
+        return loss
+
+    def is_gradient_accumulation_boundary(self):
+        """
+        Query whether the current micro-batch is at the boundary of
+        gradient accumulation, and thus will trigger gradient reductions and
+        an optimizer step.
+
+        Returns:
+            bool: if the current step is a gradient accumulation boundary.
+
+        """
+        if self._is_gradient_accumulation_boundary is None:
+            return (self.micro_steps + 1) % \
+                self.gradient_accumulation_steps() == 0
+        else:
+            return self._is_gradient_accumulation_boundary
+
+    def set_gradient_accumulation_boundary(self, is_boundary):
+        """
+        Manually overrides the DeepSpeed engine's gradient accumulation boundary state, this is an optional
+        feature and should be used with care. The state should be set before to the intended
+        value before each forward/backward. The final forward/backward should have the
+        boundary state set to True. This style allows client code to only call engine.step() once after all
+        the gradient accumulation passes are complete. See example below:
+        .. code-block:: python
+        engine.set_gradient_accumulation_boundary(False)
+        for _ in range(gradient_accumulation_steps - 1):
+            micro_batch = next(data_loader)
+            loss = engine(micro_batch)
+            engine.backward(loss)
+        engine.set_gradient_accumulation_boundary(True)
+        micro_batch = next(data_loader)
+        loss = engine(micro_batch)
+        engine.backward(loss)
+        engine.step()
+        Arguments:
+            is_boundary (bool): are we at a gradient accumulation boundary or not?
+        """
+        self._is_gradient_accumulation_boundary = is_boundary
+        self.optimizer.is_gradient_accumulation_boundary = is_boundary
+
+    def zero_grad(self):
+        """
+        Zero parameter grads.
+        """
+        for param_name, param in self.module.named_parameters():
+            param.grad = None
+
+    def clip_fp32_gradients(self):
+        clip_grad_norm_(parameters=self.module.parameters(), max_norm=self.gradient_clipping(), mpu=self.mpu)
+
+    def _take_model_step(self, lr_kwargs, block_eigenvalue={}):
+        if self.gradient_clipping() > 0.0:
+            if not (self.fp16_enabled() or self.bfloat16_enabled() or self.amp_enabled() or self.zero_optimization()):
+                self.clip_fp32_gradients()
+            elif self.amp_enabled():
+                # AMP's recommended way of doing clipping
+                # https://nvidia.github.io/apex/advanced.html#gradient-clipping
+                master_params = amp.master_params(self.optimizer)
+                clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clipping(), mpu=self.mpu)
+        self.optimizer.step()
+
+        if hasattr(self.optimizer, '_global_grad_norm'):
+            self._global_grad_norm = self.optimizer._global_grad_norm
+
+        # Quantize the updated parameter if there is no overflow
+        if self.quantizer:
+            tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage(
+            ) == 2 else self.optimizer.fp16_groups
+            if self.compression_scheduler.weight_quantization_enabled:
+                self.quantizer.quantize(
+                    tensor_to_quantize,
+                    (self.optimizer.overflow if self.fp16_enabled() else False),
+                    self.eigenvalue_enabled(),
+                    block_eigenvalue,
+                )
+        # zero grad in basic optimizer could be unreliable and may not exhibit
+        # the behavior that we want
+        if self.bfloat16_enabled():
+            # TODO: Temporary until bf16_optimizer and zero_optimizer are integrated
+            if self.zero_optimization() and hasattr(self.optimizer, "zero_grad"):
+                self.optimizer.zero_grad()
+            else:
+                pass
+        elif self.zero_optimization() or self.fp16_enabled() or self.amp_enabled():
+            self.optimizer.zero_grad()
+        else:
+            self.zero_grad()
+
+        report_progress = self.global_rank == 0 if self.global_rank else True
+
+        # Check overflow here since in DS fp16 optimizer, the overflow is updated in above step() function.
+        overflow = False
+        if hasattr(self.optimizer, "overflow"):
+            overflow = self.optimizer.overflow
+        self._step_applied = not overflow
+
+        if overflow:
+            self.skipped_steps += 1
+        else:
+            self.compression_scheduler.step()
+            if self.lr_scheduler is not None:
+                try:
+                    self.lr_scheduler.step(**(lr_kwargs or {}))
+                except TypeError:
+                    # XXX Hack to work with Megatron 2.0 and DeepSpeed pipelines.
+                    # We don't currently have a way to specify lr_kwargs from
+                    # pipe_engine.train_batch()
+                    self.lr_scheduler.step(self.train_batch_size())
+
+        if report_progress and (self.global_steps + 1) % self.steps_per_print() == 0:
+            self._report_progress(self.global_steps + 1)
+
+        self.losses = 0.0
+        self.global_steps += 1
+        self.global_samples += self.train_batch_size()
+
+    def step(self, lr_kwargs=None):
+        r"""Execute the weight update step after forward and backward propagation
+        on effective_train_batch.
+        """
+        see_memory_usage("Engine before step", force=self.memory_breakdown())
+
+        # Check early because self.global_steps is incremented at some point here.
+        # TODO: Delay self.global_steps increment until very end of this function.
+        flops_profiler_active = self.flops_profiler_enabled(
+        ) and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0
+
+        self._start_timers(self.engine_timers.step_timers)
+
+        assert self.optimizer is not None and not isinstance(self.optimizer, DummyOptim), \
+            "must provide optimizer during init in order to use step"
+
+        report_progress = False
+
+        self._step_applied = False  # assume False, will flip to True
+
+        # Update the model when we reach gradient accumulation boundaries
+        if self.is_gradient_accumulation_boundary():
+            self.gas_boundary_ctr += 1
+
+            if (self.eigenvalue_enabled() and (self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution() == 0)
+                    and self.quantizer.any_precision_switch()):
+                log_dist(f"computing eigenvalue...", ranks=[0])
+                self.block_eigenvalue = self.eigenvalue.compute_eigenvalue(self.module, self.device,
+                                                                           self.optimizer.cur_scale)
+
+            if self.progressive_layer_drop:
+                self.progressive_layer_drop.update_state(self.global_steps)
+
+            if (self.eigenvalue_enabled() and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()
+                    and self.quantizer.any_precision_switch()):
+                self._take_model_step(lr_kwargs, self.block_eigenvalue)
+            else:
+                self._take_model_step(lr_kwargs)
+
+            report_progress = self.global_rank == 0 if self.global_rank else True
+
+        self.tput_timer.stop(global_step=self.is_gradient_accumulation_boundary(), report_speed=report_progress)
+
+        self._stop_timers(self.engine_timers.step_timers)
+
+        # Log learning rate
+        if self.monitor.enabled:
+            if self.is_gradient_accumulation_boundary():
+                if self.global_rank == 0:
+                    self.summary_events = [(f"Train/Samples/lr", self.get_lr()[0], self.global_samples)]
+
+                    if self.fp16_enabled() and hasattr(self.optimizer, "cur_scale"):
+                        self.summary_events.append((
+                            f"Train/Samples/loss_scale",
+                            self.optimizer.cur_scale,
+                            self.global_samples,
+                        ))
+
+                    if (self.eigenvalue_enabled()
+                            and not self.gas_boundary_ctr % self.eigenvalue_gas_boundary_resolution()):
+                        ev_values = self.block_eigenvalue.values()
+                        for i in range(len(ev_values)):
+                            self.summary_events.append((
+                                f"Train/Eigenvalues/ModelBlockParam_{i}",
+                                self.ev_values[i][0],
+                                self.global_samples,
+                            ))
+                    self.monitor.write_events(self.summary_events)
+
+        # Check flops profiling
+        if flops_profiler_active:
+            if self.autotuning_enabled():
+                self.flops = self.flops_profiler.get_total_flops() * 3
+                self.fwd_duration = self.flops_profiler.get_total_duration()
+            else:
+                self.flops_profiler.print_model_profile(
+                    profile_step=self.global_steps,
+                    module_depth=self.flops_profiler_module_depth(),
+                    top_modules=self.flops_profiler_top_modules(),
+                    detailed=self.flops_profiler_detailed(),
+                    output_file=self.flops_profiler_output_file(),
+                )
+            self.flops_profiler.end_profile()
+
+        if self.autotuning_enabled() and self.global_steps == (self.autotuning_end_profile_step() + 1):
+            self._autotuning_exit()
+
+        if self.wall_clock_breakdown():
+            # Log micro timing and reset
+            self.timers.log(names=self.engine_timers.micro_timers, memory_breakdown=self.memory_breakdown())
+
+        if self.wall_clock_breakdown() or self.flops_profiler_enabled():
+            # Log global timing and reset
+            if self.is_gradient_accumulation_boundary():
+                if self.monitor.enabled:
+                    self._write_monitor()
+
+                if self.has_moe_layers:
+                    fwd_time = self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False)
+                    self.print_forward_breakdown(fwd_time=fwd_time)
+
+                self.timers.log(self.engine_timers.global_timers)
+
+        self.micro_steps += 1
+        see_memory_usage("Engine after step", force=self.memory_breakdown())
+
+    def _start_timers(self, timer_names):
+        for name in timer_names:
+            self.timers(name).start()
+
+    def _stop_timers(self, timer_names):
+        record = self.is_gradient_accumulation_boundary() and \
+            self.flops_profiler_enabled() and \
+                (self.global_steps >= self.flops_profiler_profile_step())
+        for name in timer_names:
+            self.timers(name).stop(record=record)
+
+    def _autotuning_exit(self):
+        if self.global_rank == 0:
+            msg = self.timers.get_mean([
+                FORWARD_GLOBAL_TIMER,
+                BACKWARD_GLOBAL_TIMER,
+                STEP_GLOBAL_TIMER,
+            ], reset=False)
+            titer = 0.0
+            titer += msg[FORWARD_GLOBAL_TIMER] if FORWARD_GLOBAL_TIMER in msg else 0
+            titer += msg[BACKWARD_GLOBAL_TIMER] if BACKWARD_GLOBAL_TIMER in msg else 0
+            titer += msg[STEP_GLOBAL_TIMER] if STEP_GLOBAL_TIMER in msg else 0
+            titer *= self.gradient_accumulation_steps()
+            msg["latency"] = titer
+            msg["FLOPS_per_gpu"] = self.flops * 1_000_000 * self.gradient_accumulation_steps() / titer
+            msg["throughput"] = self.train_batch_size() * 1_000_000 / \
+                msg["latency"]
+            print_json_dist(msg, [0], path=self.autotuning_metric_path())
+            log_dist(
+                f"Wrote metrics to {self.autotuning_metric_path()}, {os.path.abspath(self.autotuning_metric_path())}",
+                ranks=[0])
+            import atexit
+            atexit.register(print, "Autotuning: done with running current ds config.")
+        exit()
+
+    def _write_monitor(self):
+        if self.global_rank == 0:
+            self.summary_events = [
+                (
+                    f"Train/Samples/elapsed_time_ms_forward",
+                    self.timers(FORWARD_GLOBAL_TIMER).elapsed(reset=False),
+                    self.global_samples,
+                ),
+                (
+                    f"Train/Samples/elapsed_time_ms_backward",
+                    self.timers(BACKWARD_GLOBAL_TIMER).elapsed(reset=False),
+                    self.global_samples,
+                ),
+                (
+                    f"Train/Samples/elapsed_time_ms_backward_inner",
+                    self.timers(BACKWARD_INNER_GLOBAL_TIMER).elapsed(reset=False),
+                    self.global_samples,
+                ),
+                (
+                    f"Train/Samples/elapsed_time_ms_backward_allreduce",
+                    self.timers(BACKWARD_REDUCE_GLOBAL_TIMER).elapsed(reset=False),
+                    self.global_samples,
+                ),
+                (
+                    f"Train/Samples/elapsed_time_ms_step",
+                    self.timers(STEP_GLOBAL_TIMER).elapsed(reset=False),
+                    self.global_samples,
+                ),
+            ]
+            self.monitor.write_events(self.summary_events)
+
+    def _get_optimizer_param(self, param_name):
+        result = []
+        if not self.optimizer:
+            return result
+        for group in self.optimizer.param_groups:
+            if param_name in group:
+                result.append(group[param_name])
+            else:
+                result.append(0.0)
+        return result
+
+    def get_lr(self):
+        return self._get_optimizer_param("lr")
+
+    def get_type(self):
+        return self._get_optimizer_param("type")
+
+    def get_mom(self):
+        if self.optimizer_name() in ["SGD", "RMSprop"]:
+            return self._get_optimizer_param("momentum")
+        else:
+            return self._get_optimizer_param("betas")
+
+    def get_pld_theta(self):
+        if self.progressive_layer_drop:
+            return self.progressive_layer_drop.get_theta()
+        else:
+            return None
+
+    def _report_progress(self, step):
+        lr = self.get_lr()
+        mom = self.get_mom()
+        log_dist(f"step={step}, skipped={self.skipped_steps}, lr={lr}, mom={mom}", ranks=[0])
+
+    def allreduce_bucket(self, bucket, dp_group):
+        tensor = self.flatten(bucket)
+
+        tensor_to_allreduce = tensor
+
+        if self.communication_data_type != tensor.dtype:
+            tensor_to_allreduce = tensor.to(self.communication_data_type)
+
+        if self.postscale_gradients():
+            if self.gradient_predivide_factor() != 1.0:
+                tensor_to_allreduce.mul_(1.0 / self.gradient_predivide_factor())
+
+            dist.all_reduce(tensor_to_allreduce, group=dp_group)
+            if self.gradient_average:
+                if self.gradient_predivide_factor() != dist.get_world_size(group=dp_group):
+                    tensor_to_allreduce.mul_(self.gradient_predivide_factor() / dist.get_world_size(group=dp_group))
+        else:
+            tensor_to_allreduce.mul_(1. / dist.get_world_size(group=dp_group))
+            dist.all_reduce(tensor_to_allreduce, group=dp_group)
+
+        if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce:
+            tensor.copy_(tensor_to_allreduce)
+
+        return tensor
+
+    def allreduce_and_copy(self, small_bucket, dp_group):
+        allreduced = self.allreduce_bucket(small_bucket, dp_group)
+        for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)):
+            buf.copy_(synced)
+
+    def allreduce_no_retain(self, bucket, dp_group, numel_per_bucket=500000000):
+        small_bucket = []
+        numel = 0
+        for tensor in bucket:
+            small_bucket.append(tensor)
+            numel = numel + tensor.numel()
+            if numel > numel_per_bucket:
+                self.allreduce_and_copy(small_bucket, dp_group)
+                small_bucket = []
+                numel = 0
+        if len(small_bucket) > 0:
+            self.allreduce_and_copy(small_bucket, dp_group)
+
+    def _get_gradients_for_reduction(self):
+        non_expert_grads = []
+        expert_grads = {}
+        if self.has_moe_layers:
+            for key in self.expert_data_parallel_group.keys():
+                expert_grads[key] = []
+
+        for param_name, param in self.module.named_parameters():
+            if not param.requires_grad:
+                continue
+
+            if param.grad is None:
+                # In cases where there is an imbalance of empty grads across
+                # ranks we must create empty grads, this will ensure that every
+                # rank is reducing the same size. In some cases it may make
+                # sense in the future to support the ability to average not
+                # w.r.t. world size but with a different value.
+                param.grad = torch.zeros(param.size(), dtype=param.dtype, device=param.device)
+
+            grad_data = param.grad.data
+            if param_name in self.sparse_tensor_module_names or grad_data.is_sparse:
+                # Call param.grad without data to avoid problem with setting of updated grads
+                grad_data = SparseTensor(param.grad)
+
+            if is_moe_param(param):
+                expert_grads[param.group_name].append(grad_data)
+            else:
+                non_expert_grads.append(grad_data)
+
+        return non_expert_grads, expert_grads
+
+    def _reduce_non_expert_gradients(self, grads, elements_per_buffer):
+        split_buckets = split_half_float_double_sparse(grads)
+        for _, bucket_tuple in enumerate(split_buckets):
+            bucket_type, bucket = bucket_tuple
+
+            if self.pipeline_parallelism:
+                dp_group = self.mpu.get_data_parallel_group()
+            else:
+                dp_group = groups._get_sequence_data_parallel_group()
+
+            if bucket_type == SparseTensor.type():
+                self.sparse_allreduce_no_retain(bucket, dp_group=dp_group)
+            else:
+                self.allreduce_no_retain(bucket, dp_group=dp_group, numel_per_bucket=elements_per_buffer)
+
+    def _reduce_expert_gradients(self, expert_grads, elements_per_buffer):
+        for ep_name, expert_grads_group in expert_grads.items():
+            expert_split_buckets = split_half_float_double_sparse(expert_grads_group)
+            for i, bucket_tuple in enumerate(expert_split_buckets):
+                bucket_type, bucket = bucket_tuple
+                if bucket_type == SparseTensor.type():
+                    self.sparse_allreduce_no_retain(bucket, groups._get_expert_data_parallel_group(ep_name))
+                else:
+                    # Separate between diff groups
+                    self.allreduce_no_retain(bucket,
+                                             dp_group=groups._get_expert_data_parallel_group(ep_name),
+                                             numel_per_bucket=elements_per_buffer)
+
+    def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000):
+        if grads is None:
+            non_expert_grads, expert_grads = self._get_gradients_for_reduction()
+        else:
+            assert not self.has_moe_layers, "attempting to reduce grads in unsupported way w.r.t. MoE"
+            non_expert_grads = grads
+
+        self._reduce_non_expert_gradients(non_expert_grads, elements_per_buffer)
+
+        if self.has_moe_layers:
+            self._reduce_expert_gradients(expert_grads, elements_per_buffer)
+
+    def sparse_allreduce_no_retain(self, bucket, dp_group):
+        allreduced_sparses = self.sparse_allreduce_bucket(bucket, dp_group)
+        # Densify sparse tensor and copy back to original location
+        for tensor in allreduced_sparses:
+            if tensor.is_sparse:
+                tensor.orig_dense_tensor.data = tensor.to_coo_tensor()
+            else:
+                tensor.orig_dense_tensor.copy_(tensor.to_dense())
+
+    def sparse_allreduce_bucket(self, bucket, dp_group):
+        sparse_list = []
+        for sparse in bucket:
+            sparse_list.append(self.sparse_allreduce(sparse, dp_group))
+        return sparse_list
+
+    def sparse_allreduce(self, sparse, dp_group):
+        original_data_type = sparse.values.dtype
+        if self.communication_data_type != sparse.values.dtype:
+            if self.communication_data_type in (torch.float16, torch.bfloat16):
+                indices = sparse.indices.to(torch.int32)
+            else:
+                indices = sparse.indices
+            values = sparse.values.to(self.communication_data_type)
+        else:
+            indices = sparse.indices
+            values = sparse.values
+
+        if self.postscale_gradients():
+            if self.gradient_average:
+                values.mul_(self.gradient_predivide_factor() /
+                            (dist.get_world_size(group=dp_group) / float(self.sequence_parallel_size)))
+        else:
+            values.mul_(1. / (dist.get_world_size(group=dp_group) / float(self.sequence_parallel_size)))
+
+        indices_device_list = self.sparse_all_gather(indices, dp_group)
+        values_device_list = self.sparse_all_gather(values, dp_group)
+
+        sparse.indices = torch.cat(indices_device_list).to(torch.long)
+        sparse.values = torch.cat(values_device_list).to(original_data_type)
+        return sparse
+
+    def sparse_all_gather(self, value, dp_group):
+        my_size = torch.LongTensor([value.size()[0]]).to(self.device)
+        all_sizes = self.all_gather_scalar(my_size, dp_group)
+        max_size = torch.cat(all_sizes).max()
+        fill_size = max_size - my_size
+
+        assert value.dim() in [1, 2]
+        if value.dim() == 1:
+            if fill_size > 0:
+                value = torch.cat([value, value.new_empty(fill_size)])
+            tensor_list = [value.new_empty(max_size) for _ in range(dist.get_world_size(group=dp_group))]
+        else:
+            if fill_size > 0:
+                value = torch.cat([value, value.new_empty(fill_size, value.size()[1])])
+            tensor_list = [
+                value.new_empty(max_size,
+                                value.size()[1]) for _ in range(dist.get_world_size(group=dp_group))
+            ]
+
+        dist.all_gather(tensor_list, value, group=dp_group)
+        tensors = []
+        for dev_idx, t in enumerate(tensor_list):
+            size = all_sizes[dev_idx][0]
+            tensors.append(t.index_select(0, torch.arange(size, dtype=torch.long, device=self.device)))
+
+        return tensors
+
+    def all_gather_scalar(self, value, dp_group):
+        tensor_list = [value.new_zeros(value.size()) for _ in range(dist.get_world_size(group=dp_group))]
+        dist.all_gather(tensor_list, value, group=dp_group)
+        return tensor_list
+
+    def module_state_dict(self, destination=None, prefix="", keep_vars=False, exclude_frozen_parameters=False):
+        sd = self.module.state_dict(destination, prefix, keep_vars)
+
+        # Remove frozen parameter weights from state_dict if specified
+        if exclude_frozen_parameters:
+            for n, p in self.module.named_parameters():
+                if not p.requires_grad and n in sd:
+                    del sd[n]
+
+        if self.random_ltd_enabled():
+            sd = remove_random_ltd_state_dict(sd)
+        return sd
+
+    @staticmethod
+    def load_moe_state_dict(checkpoint_path,
+                            tag,
+                            state_dict,
+                            old_moe_load,
+                            model=None,
+                            mpu=None,
+                            num_experts=1,
+                            checkpoint_engine=TorchCheckpointEngine()):
+        if old_moe_load:
+            expp_rank = groups._get_expert_data_parallel_rank(groups._get_max_expert_size_name())
+
+            num_local_experts = max(num_experts) // groups._get_expert_parallel_world_size(
+                groups._get_max_expert_size_name())
+            for local_expert_id in range(num_local_experts):
+                global_expert_id = expp_rank * num_local_experts + local_expert_id
+                expert_state_dict = checkpoint_engine.load(
+                    DeepSpeedEngine._get_expert_ckpt_name(
+                        checkpoint_path,
+                        -1,  # -1 means ignore layer_id
+                        global_expert_id,
+                        tag,
+                        mpu),
+                    map_location=torch.device('cpu'))
+
+                # Updating global -> local expert ids
+                moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.'
+                for key in list(expert_state_dict.keys()):
+                    local_key = key.replace(f'{moe_str_prefix}{global_expert_id}',
+                                            f'{moe_str_prefix}{local_expert_id}')
+                    expert_state_dict[local_key] = expert_state_dict.pop(key)
+                state_dict.update(expert_state_dict)
+
+        else:
+            moe_layer_id = 0
+            for n_module, module in model.named_modules():
+                if isinstance(module, MoE):  # and deepspeed.comm.get_rank() == 0:
+                    group_name = module.expert_group_name
+                    num_local_experts = module.num_local_experts
+                    expp_rank = groups._get_expert_parallel_rank(group_name)
+                    # loop all local_experts
+                    for local_expert_id in range(num_local_experts):
+                        global_expert_id = expp_rank * num_local_experts + local_expert_id
+                        expert_state_dict = checkpoint_engine.load(DeepSpeedEngine._get_expert_ckpt_name(
+                            checkpoint_path, moe_layer_id, global_expert_id, tag, mpu),
+                                                                   map_location=torch.device('cpu'))
+                        # print(expert_state_dict.keys())
+                        # Updating global -> local expert ids
+                        moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.'
+                        for key in list(expert_state_dict.keys()):
+                            local_key = key.replace(f'{moe_str_prefix}{global_expert_id}',
+                                                    f'{moe_str_prefix}{local_expert_id}')
+                            expert_state_dict[local_key] = expert_state_dict.pop(key)
+                        state_dict.update(expert_state_dict)
+                    moe_layer_id += 1
+
+    def load_module_state_dict(self, checkpoint, strict=True, custom_load_fn=None, fetch_z3_params=False):
+        if fetch_z3_params:
+            params_to_fetch = [
+                p for p in self.module.parameters()
+                if hasattr(p, 'ds_id') and p.ds_status == ZeroParamStatus.NOT_AVAILABLE
+            ]
+        else:
+            params_to_fetch = []
+
+        with deepspeed.zero.GatheredParameters(params_to_fetch, modifier_rank=0):
+            module_state_dict = checkpoint['module']
+            if custom_load_fn:
+                custom_load_fn(src=module_state_dict, dst=self.module)
+            else:
+                self.module.load_state_dict(
+                    module_state_dict,  # TODO
+                    strict=strict)
+
+        if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None:
+            saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS]
+            for param in self.module.parameters():
+                if param.requires_grad:
+                    continue
+                if param not in self.param_names:
+                    raise ValueError(f"failed to find frozen {param} in named params")
+                name = self.param_names[param]
+                if hasattr(param, 'ds_id'):
+                    param.ds_tensor.data.copy_(saved_frozen_params[name].data)
+                else:
+                    param.data.copy_(saved_frozen_params[name].data)
+
+    def _get_zero_ckpt_prefix(self, dp_rank, bf16_mode):
+        return f'{"bf16_" if bf16_mode else ""}zero_pp_rank_{dp_rank}'
+
+    def _get_rank_zero_ckpt_name(self, checkpoints_path, tag, mp_rank, dp_rank, bf16_mode):
+        file_prefix = self._get_zero_ckpt_prefix(dp_rank, bf16_mode=bf16_mode)
+        zero_ckpt_name = os.path.join(
+            checkpoints_path,
+            str(tag),
+            f"{file_prefix}_mp_rank_{mp_rank:02d}_optim_states.pt",
+        )
+        return zero_ckpt_name
+
+    def _get_zero_ckpt_name(self, checkpoints_path, tag):
+        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+        pp_rank = dist.get_rank(group=self.optimizer.zp_process_group)
+        bf16_mode = self.bfloat16_enabled()
+        return self._get_rank_zero_ckpt_name(checkpoints_path, tag, mp_rank, pp_rank, bf16_mode)
+
+    def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None):
+        if mp_placeholder is not None:
+            mp_rank_str = mp_placeholder
+        else:
+            mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+            mp_rank_str = f"{mp_rank:02d}"
+
+        if self.zero_optimization_partition_weights():
+            filename = "zero_pp_rank_{}".format(dist.get_rank(group=self.optimizer.zp_process_group))
+            ckpt_name = os.path.join(
+                checkpoints_path,
+                str(tag),
+                f"{filename}_mp_rank_{mp_rank_str}_model_states.pt",
+            )
+        else:
+            ckpt_name = os.path.join(
+                checkpoints_path,
+                str(tag),
+                "mp_rank_" + mp_rank_str + "_model_states.pt",
+            )
+        return ckpt_name
+
+    def _get_optimizer_ckpt_name(self, checkpoints_path, tag, expp_rank):
+        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+        ckpt_name = os.path.join(checkpoints_path, str(tag),
+                                 f'expp_rank_{expp_rank}_mp_rank_{mp_rank:02d}_optim_states.pt')
+        return ckpt_name
+
+    @staticmethod
+    def _get_expert_ckpt_name(checkpoints_path, layer_id, expert_id, tag, mpu=None):
+        mp_rank = 0 if mpu is None else mpu.get_model_parallel_rank()
+        if layer_id <= -1:
+            # Used to support old checkpoint loading
+            ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag),
+                                     f'expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt')
+        else:
+            # Used to support new checkpoint loading
+            ckpt_name = os.path.join(checkpoints_path, '' if tag is None else str(tag),
+                                     f'layer_{layer_id}_expert_{expert_id}_mp_rank_{mp_rank:02d}_model_states.pt')
+        return ckpt_name
+
+    def _get_all_ckpt_names(self, checkpoints_path, tag):
+        # It is required that (checkpoints_path, tag) are consistent among all ranks.
+        ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*")
+        import glob
+
+        ckpt_files = glob.glob(ckpt_file_pattern)
+        ckpt_files.sort()
+        return ckpt_files
+
+    def load_checkpoint(self,
+                        load_dir,
+                        tag=None,
+                        load_module_strict=True,
+                        load_optimizer_states=True,
+                        load_lr_scheduler_states=True,
+                        load_module_only=False,
+                        custom_load_fn=None):
+        """
+        Load training checkpoint
+
+        Arguments:
+            load_dir: Required. Directory to load the checkpoint from
+            tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file
+            load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match.
+            load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance
+            load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint.
+            load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting.
+            custom_load_fn: Optional. Custom model load function.
+
+        Returns:
+            A tuple of ``load_path`` and ``client_state``.
+            *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed.
+            *``client_state``: State dictionary used for loading required training states in the client code.
+
+        Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right
+        after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and
+        ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine
+        before ``load_checkpoint()``.
+
+        """
+
+        if tag is None:
+            latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest"
+            latest_path = os.path.join(load_dir, latest_tag)
+            if os.path.isfile(latest_path):
+                with open(latest_path, "r") as fd:
+                    tag = fd.read().strip()
+            else:
+                if self.load_universal_checkpoint():
+                    raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist')
+                else:
+                    logger.warning(
+                        f"Unable to find latest file at {latest_path}, if trying to load latest "
+                        "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint."
+                    )
+                    return None, None
+
+        if self._optimizer_has_ckpt_event_prologue():
+            # Prepare for checkpoint load by ensuring all parameters are partitioned
+            self.optimizer.checkpoint_event_prologue()
+
+        load_path, client_states = self._load_checkpoint(load_dir,
+                                                         tag,
+                                                         load_module_strict=load_module_strict,
+                                                         load_optimizer_states=load_optimizer_states,
+                                                         load_lr_scheduler_states=load_lr_scheduler_states,
+                                                         load_module_only=load_module_only,
+                                                         custom_load_fn=custom_load_fn)
+
+        load_zero_checkpoint = load_optimizer_states and load_path is not None and (self.zero_optimization()
+                                                                                    or self.bfloat16_enabled())
+        if load_zero_checkpoint:
+            success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states)
+            if not success:
+                self.optimizer._restore_from_bit16_weights()
+
+        if self._optimizer_has_ckpt_event_epilogue():
+            self.optimizer.checkpoint_event_epilogue()
+
+        if self.load_universal_checkpoint():
+            self.optimizer.update_lp_params()
+            if load_zero_checkpoint:
+                self.update_optimizer_step(step=client_states['iteration'] + 1)
+
+        return load_path, client_states
+
+    def _load_checkpoint(self,
+                         load_dir,
+                         tag,
+                         load_module_strict=True,
+                         load_optimizer_states=True,
+                         load_lr_scheduler_states=True,
+                         load_module_only=False,
+                         custom_load_fn=None):
+
+        from deepspeed.runtime.state_dict_factory import SDLoaderFactory
+
+        ckpt_list = self._get_all_ckpt_names(load_dir, tag)
+        sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine)
+
+        is_pipe_parallel = isinstance(self.module, PipelineModule)
+
+        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+        load_path, checkpoint, _ = sd_loader.load(self.mp_world_size, mp_rank, is_pipe_parallel=is_pipe_parallel)
+
+        if checkpoint is None:
+            return None, None
+
+        fetch_z3_params = False
+        if self.zero_optimization_partition_weights() and not load_optimizer_states:
+            checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir)
+            fetch_z3_params = True
+
+        if is_pipe_parallel:
+            # Pipeline parallelism uses this to load its own checkpoint files.
+            self._curr_ckpt_path = os.path.join(load_dir, tag)
+
+        if self.has_moe_layers:
+            # print(checkpoint.keys())
+            old_moe_load = False
+            if not isinstance(checkpoint['num_experts'], list):
+                old_moe_load = True
+            DeepSpeedEngine.load_moe_state_dict(load_dir,
+                                                tag,
+                                                state_dict=checkpoint['module'],
+                                                old_moe_load=old_moe_load,
+                                                model=self.module,
+                                                mpu=self.mpu,
+                                                num_experts=self.num_experts,
+                                                checkpoint_engine=self.checkpoint_engine)
+        if not self.load_universal_checkpoint():
+            self.load_module_state_dict(checkpoint=checkpoint,
+                                        strict=load_module_strict,
+                                        custom_load_fn=custom_load_fn,
+                                        fetch_z3_params=fetch_z3_params)
+
+        self.loaded_checkpoint_dp_world_size = checkpoint['dp_world_size']
+        if 'zp_world_size' not in checkpoint:
+            checkpoint['zp_world_size'] = self.zp_world_size
+        self.loaded_checkpoint_zp_world_size = checkpoint['zp_world_size']
+
+        optim_checkpoint = None
+        if load_module_only:
+            deepspeed_states = ['module']
+            if self.optimizer is not None and self.fp16_enabled():
+                self.optimizer.refresh_fp32_params()
+        else:
+            has_zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled()
+            if load_optimizer_states and self.optimizer is not None and not has_zero_optimizer_state:
+                if self.has_moe_layers:
+                    largest_group_name = groups._get_max_expert_size_name()
+                    expp_rank = groups._get_expert_parallel_rank(largest_group_name)
+                    optim_load_path = self._get_optimizer_ckpt_name(load_dir, tag, expp_rank)
+                    optim_checkpoint = self.checkpoint_engine.load(optim_load_path, map_location=torch.device('cpu'))
+                else:
+                    optim_checkpoint = checkpoint
+
+                if self.fp16_enabled() or self.bfloat16_enabled():
+                    self.optimizer.load_state_dict(optim_checkpoint['optimizer'],
+                                                   load_optimizer_states=load_optimizer_states)
+                else:
+                    optim_checkpoint = checkpoint
+
+                self.optimizer.load_state_dict(optim_checkpoint['optimizer'])
+
+            if load_lr_scheduler_states and self.lr_scheduler is not None:
+                self.lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+
+            if self.random_ltd_enabled() and self.random_ltd_scheduler is not None and 'random_ltd' in checkpoint:
+                self.random_ltd_scheduler.load_state_dict(checkpoint['random_ltd'])
+
+            if self.training_dataloader is not None and self.curriculum_learning_enabled(
+            ) and 'data_sampler' in checkpoint:
+                self.training_dataloader.data_sampler.load_state_dict(checkpoint['data_sampler'])
+
+            def get_sparse_tensor_module_names(original_set, loaded_set, original_parameters, loaded_parameters):
+                result = set()
+
+                for name in original_set:
+                    if name in loaded_parameters and name not in loaded_set:
+                        continue  # parameter existed in previous model and was not sparse
+                    result.add(name)
+
+                for name in loaded_set:
+                    if name in original_parameters:
+                        result.add(name)  # parameter exists in both configs and it was sparse
+
+                return result
+
+            if 'sparse_tensor_module_names' in checkpoint:
+                sparse_tensor_module_names = checkpoint['sparse_tensor_module_names']
+            elif 'csr_tensor_module_names' in checkpoint:
+                sparse_tensor_module_names = checkpoint['csr_tensor_module_names']
+            else:
+                sparse_tensor_module_names = None
+            if sparse_tensor_module_names is not None:
+                if load_module_strict:
+                    self.sparse_tensor_module_names = sparse_tensor_module_names
+                else:
+                    self.sparse_tensor_module_names = get_sparse_tensor_module_names(
+                        self.sparse_tensor_module_names, sparse_tensor_module_names,
+                        dict(self.module.named_parameters()), checkpoint["module"])
+
+            self.global_steps = checkpoint['global_steps']
+            self.global_samples = checkpoint.get('global_samples', self.global_steps * self.train_batch_size())
+            self.skipped_steps = checkpoint['skipped_steps']
+            self.loaded_checkpoint_mp_world_size = checkpoint['mp_world_size']
+            deepspeed_states = [
+                'module', 'sparse_tensor_module_names', 'skipped_steps', 'global_steps', 'zp_world_size',
+                'mp_world_size', 'data_sampler', 'random_ltd', 'dp_world_size',
+            ]
+        client_state = {}
+
+        if load_lr_scheduler_states:
+            deepspeed_states.append('lr_scheduler')
+        if load_optimizer_states:
+            deepspeed_states.append('optimizer')
+
+        client_state = {key: value for key, value in checkpoint.items() if not key in deepspeed_states}
+
+        if optim_checkpoint is not None:
+            client_state['optimizer'] = optim_checkpoint['optimizer']
+
+        return load_path, client_state
+
+    def _load_zero_checkpoint(self, load_dir, tag, load_optimizer_states=True):
+
+        load_serial = None
+        # When use loading checkpoint serial, checkpoint loading start from local rank 0,
+        # all other local rank would be paused, waiting for its rank-1 peer ready and its notification.
+        if self._config.zero_config.pipeline_loading_checkpoint:
+            assert self.zero_optimization_stage(
+            ) == ZeroStageEnum.weights, "Only stage3 support for pipeline checkpoint loading"
+            load_serial = torch.zeros(1).to(self.device)
+            if dist.get_local_rank() != 0:
+                dist.recv(tensor=load_serial, src=dist.get_rank() - 1)
+        if self.load_universal_checkpoint():
+            zero_sd_list = None
+            checkpoint_folder = f'{os.path.join(load_dir, tag)}'
+        else:
+            if load_optimizer_states and self.zp_world_size != self.loaded_checkpoint_zp_world_size:
+                raise ZeRORuntimeException("The checkpoint being loaded used a DP " \
+                    f"world size of {self.loaded_checkpoint_zp_world_size} but the " \
+                    f"current world size is {self.zp_world_size}. Automatic adjustment " \
+                    "of ZeRO's optimizer state partitioning with a new world size is not " \
+                    "currently supported.")
+            checkpoint_folder = None
+            zero_sd_list = self._get_all_zero_checkpoints(load_dir, tag)
+            if zero_sd_list is None:
+                return False
+
+        self.optimizer.load_state_dict(state_dict_list=zero_sd_list,
+                                       load_optimizer_states=load_optimizer_states,
+                                       load_from_fp32_weights=self.zero_load_from_fp32_weights(),
+                                       checkpoint_folder=checkpoint_folder,
+                                       load_serial=load_serial)
+
+        if self.load_universal_checkpoint():
+            logger.info(f'loaded universal zero checkpoints from {checkpoint_folder} for rank {self.global_rank}')
+        else:
+            logger.info(f"loading {len(zero_sd_list)} zero partition checkpoints for rank {self.global_rank}")
+        return True
+
+    def update_optimizer_step(self, step):
+
+        def set_step(d):
+            if isinstance(d['step'], torch.Tensor):
+                d['step'] = torch.tensor(step, dtype=d['step'].dtype, device=d['step'].device)
+            else:
+                d['step'] = step
+
+        optimizer = self.optimizer
+        base_optimizer = optimizer.optimizer
+        state = base_optimizer.state
+        for group in optimizer.param_groups:
+            if 'step' in group:
+                set_step(group)
+            for p in group['params']:
+                if p in state and len(state[p]) > 0 and 'step' in state[p]:
+                    set_step(state[p])
+
+    def _get_mp_rank_zero_checkpoint_names(self, load_dir, tag, mp_rank, dp_world_size, bf16_mode):
+        zero_ckpt_names = []
+        for dp_rank in range(dp_world_size):
+            ckpt_name = self._get_rank_zero_ckpt_name(checkpoints_path=load_dir,
+                                                      tag=tag,
+                                                      mp_rank=mp_rank,
+                                                      dp_rank=dp_rank,
+                                                      bf16_mode=bf16_mode)
+            zero_ckpt_names.append(ckpt_name)
+
+        return zero_ckpt_names
+
+    def _get_all_zero_checkpoint_names(self, load_dir, tag, bf16_mode):
+        mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+        zero_ckpt_names = self._get_mp_rank_zero_checkpoint_names(load_dir=load_dir,
+                                                                  tag=tag,
+                                                                  mp_rank=mp_rank,
+                                                                  dp_world_size=self.loaded_checkpoint_dp_world_size,
+                                                                  bf16_mode=bf16_mode)
+        for i, ckpt_name in enumerate(zero_ckpt_names):
+            if not os.path.exists(ckpt_name):
+                # transparently handle the old file pattern for optim_states
+                if "optim_states.pt" in ckpt_name:
+                    ckpt_name_try = ckpt_name.replace("_optim_states.pt", "optim_states.pt")
+                    if os.path.exists(ckpt_name_try):
+                        zero_ckpt_names[i] = ckpt_name_try
+                        continue
+
+        return zero_ckpt_names
+
+    def _get_all_zero_checkpoint_state_dicts(self, zero_ckpt_names):
+        zero_sd_list = []
+        for i, ckpt_name in enumerate(zero_ckpt_names):
+            _state = None
+            if ckpt_name is None:
+                _state = {OPTIMIZER_STATE_DICT: None}
+            # Fully load state for current rank
+            elif self.zero_elastic_checkpoint() or dist.get_rank(group=self.optimizer.zp_process_group) == i:
+                _state = self.checkpoint_engine.load(
+                    ckpt_name,
+                    map_location='cpu',
+                )
+            else:
+                _state = {OPTIMIZER_STATE_DICT: None}
+            zero_sd_list.append(_state)
+
+        zero_optimizer_sd = [sd[OPTIMIZER_STATE_DICT] for sd in zero_sd_list]
+        logger.info(f"successfully read {len(zero_optimizer_sd)} ZeRO state_dicts for rank {self.global_rank}")
+        return zero_optimizer_sd
+
+    def _get_all_zero_checkpoints(self, load_dir, tag):
+        for bf16_mode in [self.bfloat16_enabled(), not self.bfloat16_enabled()]:
+            zero_ckpt_names = self._get_all_zero_checkpoint_names(load_dir, tag, bf16_mode)
+            if zero_ckpt_names is not None:
+                # Warn if loading checkpoint of different bit16 type
+                if bf16_mode is not self.bfloat16_enabled():
+                    checkpoint_bit16 = BFLOAT16 if bf16_mode else FP16
+                    engine_bit16 = BFLOAT16 if self.bfloat16_enabled() else FP16
+                    logger.warn(f'Loading {checkpoint_bit16} zero checkpoints into {engine_bit16} training engine')
+                return self._get_all_zero_checkpoint_state_dicts(zero_ckpt_names)
+
+        return None
+
+    def _checkpoint_tag_validation(self, tag):
+        if self.checkpoint_tag_validation_enabled():
+            s_hash = hashlib.sha1(tag.encode())
+            bhash = torch.ByteTensor([s_hash.digest()]).flatten().to(self.device)
+            max_bhash = bhash.clone()
+            min_bhash = bhash.clone()
+            dist.all_reduce(max_bhash, op=dist.ReduceOp.MAX)
+            dist.all_reduce(min_bhash, op=dist.ReduceOp.MIN)
+            valid = all(min_bhash == bhash) and all(max_bhash == bhash)
+            msg = (f"[rank={dist.get_rank()}] The checkpoint tag name '{tag}' is not consistent across "
+                   "all ranks. Including rank unique information in checkpoint tag could cause issues when "
+                   "restoring with different world sizes.")
+            if self.checkpoint_tag_validation_fail():
+                assert valid, msg
+            elif not valid:
+                logger.warning(msg)
+
+    def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False):
+        """Save training checkpoint
+
+        Arguments:
+            save_dir: Required. Directory for saving the checkpoint
+            tag: Optional. Checkpoint tag used as a unique identifier for the checkpoint, global step is
+                used if not provided. Tag name must be the same across all ranks.
+            client_state: Optional. State dictionary used for saving required training states in the client code.
+            save_latest: Optional. Save a file 'latest' pointing to the latest saved checkpoint.
+            exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state.
+        Important: all processes must call this method and not just the process with rank 0. It is
+        because each process needs to save its master weights and scheduler+optimizer states. This
+        method will hang waiting to synchronize with other processes if it's called just for the
+        process with rank 0.
+
+        """
+        if self._optimizer_has_ckpt_event_prologue():
+            # Custom preparation for checkpoint save, if applicable
+            self.optimizer.checkpoint_event_prologue()
+
+        rank = self.local_rank if self.use_node_local_storage() else self.global_rank
+
+        # This is to make sure the checkpoint names are created without collision
+        # There seems to be issue creating them in parallel
+
+        # Ensure save_dir directory exists
+        if rank == 0:
+            self.checkpoint_engine.makedirs(save_dir, exist_ok=True)
+        dist.barrier()
+
+        if tag is None:
+            tag = f"global_step{self.global_steps}"
+
+        # Ensure tag is a string
+        tag = str(tag)
+        self.checkpoint_engine.create(tag)
+
+        # Ensure checkpoint tag is consistent across ranks
+        self._checkpoint_tag_validation(tag)
+
+        if self.has_moe_layers:
+            self.save_non_zero_checkpoint = False
+            self._create_checkpoint_file(save_dir, tag, False)
+            self._save_moe_checkpoint(save_dir,
+                                      tag,
+                                      client_state=client_state,
+                                      exclude_frozen_parameters=exclude_frozen_parameters)
+
+        # We distribute the task of saving layer checkpoint files among
+        # data parallel instances, so all procs should call _save_checkpoint.
+        # All procs then call module_state_dict(), but only procs of data
+        # parallel rank 0 save the general model params.
+        if not self.has_moe_layers:
+            self._create_checkpoint_file(save_dir, tag, False)
+            self._save_checkpoint(save_dir,
+                                  tag,
+                                  client_state=client_state,
+                                  exclude_frozen_parameters=exclude_frozen_parameters)
+
+        if self.save_zero_checkpoint:
+            self._create_zero_checkpoint_files(save_dir, tag)
+            self._save_zero_checkpoint(save_dir, tag)
+
+        if self._optimizer_has_ckpt_event_epilogue():
+            self.optimizer.checkpoint_event_epilogue()
+
+        # Save latest checkpoint tag
+        self.checkpoint_engine.commit(tag)
+        if save_latest and rank == 0:
+            with open(os.path.join(save_dir, 'latest'), 'w') as fd:
+                fd.write(tag)
+
+        dist.barrier()
+
+        return True
+
+    def _get_non_moe_state_dict(self, full_state_dict):
+        """
+            Get the state dict of the non-moe layers
+        """
+        for key in list(full_state_dict.keys()):
+            if 'expert' in key and 'moe.gate.wg.weight' not in key:
+                full_state_dict.pop(key)
+
+        return full_state_dict
+
+    def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False):
+        save_path = self._get_ckpt_name(save_dir, tag)
+        # A hack to save the checkpointing directory. Pipeline parallelism overrides
+        # module_state_dict() and uses this path to save the model. module_state_dict()
+        # then instead just returns None.
+
+        # Using layer_#_export_# to save the model's expert state_dict
+        moe_layer_id = 0
+        for n_module, module in self.module.named_modules():
+            if isinstance(module, MoE):  # and deepspeed.comm.get_rank() == 0:
+                group_name = module.expert_group_name
+                num_local_experts = module.num_local_experts
+                expp_rank = groups._get_expert_parallel_rank(group_name)
+                exp_dp_rank = groups._get_expert_data_parallel_rank(group_name)
+                # print(expp_rank, exp_dp_rank)
+                if exp_dp_rank != 0:
+                    moe_layer_id += 1
+                    continue
+
+                # get all moe parameters
+                moe_state_dict = {}
+                for n, p in module.state_dict().items():
+                    if 'expert' in n and 'moe.gate.wg.weight' not in n:
+                        moe_state_dict[n_module + '.' + n] = p
+                moe_str_prefix = '.deepspeed_moe.experts.deepspeed_experts.'
+                # print(moe_state_dict.keys()) # until now, everything is fine. So the bug happens at next few lines
+                # Reorder the moe name rank, so that each checkpoint only has one expert
+                experts_state_dict = defaultdict(dict)
+                for key in list(moe_state_dict.keys()):
+                    m = re.match(f".*{moe_str_prefix}([0-9]+).*", key)
+
+                    local_expert_id = None
+                    if not m:
+                        logger.warn(f'No expert found in key {key}.')
+                    else:
+                        local_expert_id = m.group(1)
+
+                    global_expert_id = expp_rank * \
+                        num_local_experts + int(local_expert_id)
+                    expert_key = key.replace(f'{moe_str_prefix}{local_expert_id}',
+                                             f'{moe_str_prefix}{global_expert_id}')
+                    # truncating extra tensor (shared) storage
+                    truncated = moe_state_dict.pop(key).clone().detach()
+                    experts_state_dict[str(global_expert_id)][expert_key] = truncated
+
+                # let save the moe parameters
+                for global_expert_id, expert_state_dict in experts_state_dict.items():
+                    # save the moe parameters
+                    moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu)
+                    if self.random_ltd_enabled():
+                        expert_state_dict = remove_random_ltd_state_dict(expert_state_dict)
+                    self.checkpoint_engine.save(expert_state_dict, moe_save_path)
+                moe_layer_id += 1
+
+        self._curr_ckpt_path = os.path.join(save_dir, tag)
+
+        largest_group_name = groups._get_max_expert_size_name()
+        expp_rank = groups._get_expert_parallel_rank(largest_group_name)
+        exp_dp_rank = groups._get_expert_data_parallel_rank(largest_group_name)
+
+        # In the case of E + D parallelism, only the
+        # first expert parallel group should save the expert weights
+        # since each expert parallel group is a copy of the model's experts
+        if exp_dp_rank != 0:
+            return
+
+        # Save optimizer states. They are different across each exp parallel rank.
+        optimizer_state = {
+            'optimizer': self.optimizer.state_dict() if self.optimizer and not self.zero_optimization() else None
+        }
+        # TODO: why use BufferedWriter not the path
+        file_path = self._get_optimizer_ckpt_name(save_dir, tag, expp_rank)
+        self.checkpoint_engine.save(optimizer_state, file_path)
+
+        # get non-moe parameters
+        model_state_dict = self._get_non_moe_state_dict(
+            self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters))
+
+        if expp_rank == 0:
+            # TODO: update num experts info,.. in checkpoint
+            state = {
+                'module':
+                model_state_dict,
+                'lr_scheduler':
+                self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None,
+                'data_sampler':
+                self.training_dataloader.data_sampler.state_dict() if
+                (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None,
+                'random_ltd':
+                self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None,
+                'sparse_tensor_module_names':
+                self.sparse_tensor_module_names,
+                'skipped_steps':
+                self.skipped_steps,
+                'global_steps':
+                self.global_steps,
+                'global_samples':
+                self.global_samples,
+                'zp_world_size':
+                self.zp_world_size,
+                'dp_world_size':
+                self.dp_world_size,
+                'mp_world_size':
+                self.mp_world_size,
+                'num_experts':
+                self.num_experts
+            }
+            state.update(client_state)
+            logger.info(f'Saving model checkpoint: {save_path}')
+            self.checkpoint_engine.save(state, save_path)
+
+    def _create_checkpoint_file(self, save_dir, tag, zero_checkpoint):
+        name_function = (self._get_zero_ckpt_name if zero_checkpoint else self._get_ckpt_name)
+        try:
+            checkpoint_name = name_function(save_dir, tag)
+            path = os.path.dirname(checkpoint_name)
+            self.checkpoint_engine.makedirs(path, exist_ok=True)
+        except:
+            logger.error(f"Failed saving model checkpoint to {save_dir} with tag {tag}")
+            return False
+
+        return True
+
+    def _create_zero_checkpoint_files(self, save_dir, tag):
+        success = True
+        # zero checkpoint files are created sequentially
+        for rank in range(dist.get_world_size(self.optimizer.zp_process_group)):
+            if rank == self.global_rank:
+                success = self._create_checkpoint_file(save_dir, tag, True)
+
+            dist.barrier(group=self.optimizer.zp_process_group)
+
+        return success
+
+    def _save_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_parameters=False):
+
+        save_path = self._get_ckpt_name(save_dir, tag)
+
+        zero_optimizer_state = self.zero_optimization() or self.bfloat16_enabled()
+
+        save_frozen_param = self.zero_optimization_partition_gradients() and not exclude_frozen_parameters
+
+        # A hack to save the checkpointing directory. Pipeline parallelism overrides
+        # module_state_dict() and uses this path to save the model. module_state_dict()
+        # then instead just returns None.  The module_state_dict() implementation in
+        # PipelineEngine expects the save path to be set in self._curr_ckpt_path.
+        self._curr_ckpt_path = os.path.join(save_dir, tag)
+        module = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters)
+        self._curr_ckpt_path = None
+
+        state = dict(module=module,
+                     buffer_names=self._get_buffer_names(),
+                     optimizer=self.optimizer.state_dict() if self.optimizer and not zero_optimizer_state else None,
+                     param_shapes=self._get_zero_param_shapes() if self.optimizer and zero_optimizer_state else None,
+                     frozen_param_shapes=self._get_zero_frozen_param_attributes(self._get_param_shape_func)
+                     if save_frozen_param else None,
+                     shared_params=self._get_shared_params() if self.optimizer and zero_optimizer_state else None,
+                     frozen_param_fragments=self._get_zero_frozen_param_attributes(self._get_param_fragment_func)
+                     if save_frozen_param else None,
+                     lr_scheduler=self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None,
+                     data_sampler=self.training_dataloader.data_sampler.state_dict() if
+                     (self.training_dataloader is not None and self.curriculum_learning_enabled()) else None,
+                     random_ltd=self.random_ltd_scheduler.state_dict() if self.random_ltd_enabled() else None,
+                     sparse_tensor_module_names=self.sparse_tensor_module_names,
+                     skipped_steps=self.skipped_steps,
+                     global_steps=self.global_steps,
+                     global_samples=self.global_samples,
+                     dp_world_size=self.seq_dp_world_size,
+                     mp_world_size=self.mp_world_size,
+                     ds_config=self.config,
+                     ds_version=version)
+        state.update(client_state)
+
+        if self.save_non_zero_checkpoint:
+            log_dist(message=f'Saving model checkpoint: {save_path}', ranks=[0, 1])
+            self.checkpoint_engine.save(state, save_path)
+
+    def _get_buffer_names(self):
+        buffer_names = []
+
+        # we save buffer names so that we could extract later the real buffers from the saved
+        # state_dict["module"] in the non-zero checkpoint - the buffers are already there but they
+        # are intermixed with param placeholders
+
+        # have to traverse the tree to be able to skip non-persistent buffers
+        def get_layer_named_buffers(module, prefix=""):
+            for name, buf in module.named_buffers(recurse=False):
+                if buf is not None and name not in module._non_persistent_buffers_set:
+                    buffer_names.append(prefix + name)
+
+            for name, child in module.named_children():
+                if child is not None:
+                    get_layer_named_buffers(child, prefix + name + ".")
+
+        get_layer_named_buffers(self.module, prefix="")
+
+        return buffer_names
+
+    def _get_param_shape_func(self, param):
+        return param.ds_shape if hasattr(param, 'ds_id') else param.shape
+
+    def _get_param_fragment_func(self, param):
+        return param.ds_tensor.detach().cpu() if hasattr(param, 'ds_id') else param.detach().cpu()
+
+    def _get_zero_frozen_param_attributes(self, attr_func):
+        frozen_param_fragments = OrderedDict()
+
+        for param in self.module.parameters():
+            if param.requires_grad:
+                continue
+            if param not in self.param_names:
+                raise ValueError(f"failed to find frozen {param} in named params")
+            name = self.param_names[param]
+            frozen_param_fragments[name] = attr_func(param)
+
+        return frozen_param_fragments
+
+    def _get_zero_param_shapes(self):
+        """Returns a dict of name to shape mapping, only for the flattened fp32 weights saved by the
+        optimizer. the names are exactly as in state_dict. The order is absolutely important, since
+        the saved data is just flattened data with no identifiers and requires reconstruction in the
+        same order it was saved.
+        We can't rely on self.module.named_parameters() to get the saved tensors, as some params
+        will be missing and others unsaved and then it'd be impossible to reconstruct state_dict
+        from the flattened weights.
+        optimizer.bit16_groups seems to be the easiest to use as it's in all zeroX versions.
+        """
+        param_group_shapes = []
+        cnt = 0
+        numel = 0
+
+        # zero2 started using a round_robin_bit16_groups which is a shuffled version of bit16_groups -
+        # if we don't use it, we get parameters ordered incorrectly
+        if hasattr(self.optimizer, "round_robin_bit16_groups"):
+            bit16_groups = self.optimizer.round_robin_bit16_groups
+        elif self.bfloat16_enabled() and hasattr(self.optimizer, "bf16_groups"):
+            bit16_groups = self.optimizer.bf16_groups
+        else:
+            bit16_groups = self.optimizer.bit16_groups if self.zero_optimization_stage(
+            ) == 2 else self.optimizer.fp16_groups
+
+        for bit16_group in bit16_groups:
+            param_shapes = OrderedDict()
+            for param in bit16_group:
+                cnt += 1
+                numel += param.ds_numel if hasattr(param, "ds_numel") else param.numel()
+                shape = param.ds_shape if hasattr(param, "ds_shape") else param.shape
+                if param not in self.param_names:
+                    raise ValueError(f"failed to find optimizer param in named params")
+                name = self.param_names[param]
+                param_shapes[name] = shape
+
+                # uncomment to debug zero_to_fp32.py problems
+                # if self.global_rank == 0: print(f"saving param {name} {shape} (numel={shape.numel()})")
+            param_group_shapes.append(param_shapes)
+        # if self.global_rank == 0: print(f"Total saved {numel} numels in {cnt} params")
+
+        return param_group_shapes
+
+    def _get_shared_params(self):
+        """
+        Returns a dict of shared params, which can later be used to reconstruct the original state dict,
+        e.g. in `zero_to_fp32`. Each dict entry is a pair of param names, where the key is the name
+        of the variable that isn't stored and the value is the actual param holding data.
+        """
+        shared_index = {}
+        shared_params_by_full_name = {}
+
+        is_zero3_model = (self.zero_optimization_partition_weights()
+                          and any(hasattr(param, "ds_id") for param in self.module.parameters()))
+
+        def get_layer_state_dict(module, prefix=""):
+            # handle params
+            for name, param in module.named_parameters(recurse=False):
+                if param is None or (is_zero3_model and not hasattr(param, "ds_id")):
+                    continue
+                key = prefix + name
+
+                # When weights are manged by stage 3, we can't rely on param.data_ptr() as it will be reused
+                # as weights get gathered and reduced, but param.ds_id is unique across all zero weights
+                # (and shared params will have the same param.ds_id)
+                param_id = param.ds_id if is_zero3_model else param.data_ptr()
+
+                if param_id in shared_index:
+                    # shared weights
+                    #print(f"`{key}` is shared with `{shared_index[param_id]}`")
+                    shared_params_by_full_name[key] = shared_index[param_id]
+                else:
+                    shared_index[param_id] = key
+
+            for name, child in module.named_children():
+                if child is not None:
+                    get_layer_state_dict(child, prefix + name + ".")
+
+        if dist.get_rank() == 0:
+            get_layer_state_dict(self.module, prefix="")
+
+        return shared_params_by_full_name
+
+    def _copy_recovery_script(self, save_path):
+        base_dir = os.path.dirname(os.path.dirname(__file__))
+        script = "zero_to_fp32.py"
+        src = os.path.join(base_dir, "utils", script)
+        dst = os.path.join(save_path, script)
+        #logger.info(f"creating recovery script {dst}")
+        copyfile(src, dst)
+        self._change_recovery_script_permissions(dst)
+
+    def _change_recovery_script_permissions(self, dst):
+        # make executable (safeguard for file shares - Azure as example)
+        try:
+            os.chmod(dst, os.stat(dst).st_mode | stat.S_IEXEC)
+        except (FileNotFoundError, PermissionError) as e:
+            #this message is used in unit test TestZeRONonDistributed
+            logger.info(
+                f'Warning: Could not change permissions for {dst} due to error: {e}. Continuing without changing permissions.'
+            )
+
+    def _save_zero_checkpoint(self, save_path, tag):
+        zero_checkpoint_name = self._get_zero_ckpt_name(save_path, tag)
+        zero_sd = dict(optimizer_state_dict=self.optimizer.state_dict(), ds_config=self.config, ds_version=version)
+        self.checkpoint_engine.save(zero_sd, zero_checkpoint_name)
+
+        if self.global_rank == 0:
+            self._copy_recovery_script(save_path)
+        ckpt_type = 'zero' if self.zero_optimization() else 'bf16_zero'
+        logger.info(f'{ckpt_type} checkpoint saved {zero_checkpoint_name}')
+
+    def _zero3_consolidated_16bit_state_dict(self):
+        """
+        Get a full non-partitioned state_dict with fp16 weights on cpu.
+        Important: this function must be called on all ranks and not just rank 0.
+        This is similar to nn.Module.state_dict (modelled after _save_to_state_dict), but:
+        1. consolidates the weights from different partitions on gpu0
+        2. works on one layer at a time to require as little gpu0 memory as possible, by
+        moving the already consolidated weights to cpu
+        3. takes care to keep the shared params shared when gradually copying the params to cpu
+        Returns:
+            a consolidated fp16 ``state_dict`` on cpu on rank 0, ``None`` on other ranks
+        """
+        if not self.zero_optimization_partition_weights():
+            raise ValueError("this function requires ZeRO-3 mode")
+
+        state_dict = OrderedDict() if dist.get_rank() == 0 else None
+        shared_params = {}
+
+        def get_layer_state_dict(module, prefix=""):
+            # gather one layer at a time to be memory-efficient
+            # must use modifier_rank=0 to release GPU memory after each layer gathered
+            #see_memory_usage("before GatheredParameters", force=True)
+            with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0):
+                if dist.get_rank() == 0:
+                    # handle params
+                    for name, param in module.named_parameters(recurse=False):
+                        if param is None:
+                            continue
+                        key = prefix + name
+                        # can't rely on param.data_ptr() as it will be reused as weights gets
+                        # gathered and reduced, but param.ds_id is unique across all zero weights
+                        # (and shared params will have the same param.ds_id)
+                        if param.ds_id in shared_params:
+                            # shared weights
+                            #print(f"`{key}` is shared with `{shared_params[param.ds_id]}`")
+                            state_dict[key] = state_dict[shared_params[param.ds_id]]
+                        else:
+                            state_dict[key] = param.detach().cpu()
+                            shared_params[param.ds_id] = key
+                        #print(f"param {param.ds_id} {param.shape} {key} ")
+
+                    # now buffers - not sure if need to take care of potentially shared weights here
+                    for name, buf in module.named_buffers(recurse=False):
+                        if (buf is not None and name not in module._non_persistent_buffers_set):
+                            state_dict[prefix + name] = buf.detach().cpu()
+            #see_memory_usage("after GatheredParameters", force=True)
+
+            for name, child in module.named_children():
+                if child is not None:
+                    get_layer_state_dict(child, prefix + name + ".")
+
+        # Prepare for checkpoint save by ensuring all parameters are partitioned
+        if self._optimizer_has_ckpt_event_prologue():
+            self.optimizer.checkpoint_event_prologue()
+
+        see_memory_usage("before get_layer_state_dict", force=False)
+        get_layer_state_dict(self.module, prefix="")
+        see_memory_usage("after get_layer_state_dict", force=False)
+
+        if self._optimizer_has_ckpt_event_epilogue():
+            self.optimizer.checkpoint_event_epilogue()
+
+        return state_dict
+
+    def save_fp16_model(self, save_dir, save_filename="pytorch_model.bin"):
+        """has been renamed to save_16bit_model, keeping this around for backwards
+        compatibility"""
+        return self.save_16bit_model(save_dir, save_filename)
+
+    def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin"):
+        """
+        Save 16bit model weights
+
+        This method saves the 16bit model weights at the desired destination.
+
+        Arguments:
+            save_dir: Required. Directory for saving the model
+            save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin``
+
+        Returns:
+            ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if
+            stage3_gather_16bit_weights_on_model_save is ``False``.
+
+        Important: all processes must call this method and not just the process with rank 0. It is
+        because the processes need to work in sync to gather the weights. This method will hang
+        waiting to synchronize with other processes if it's called just for the process with rank 0.
+
+        """
+
+        path = os.path.join(save_dir, save_filename)
+
+        if self.zero_optimization_partition_weights():
+            if self.zero_gather_16bit_weights_on_model_save():
+                # consolidation is expensive in time and memory and therefore isn't a default
+                state_dict = self._zero3_consolidated_16bit_state_dict()
+            else:
+                # the model will be bogus if not consolidated so don't confuse the user by saving it
+                logger.info(
+                    f"Did not save the model {path} because `stage3_gather_16bit_weights_on_model_save` is False")
+                return False
+        else:
+            state_dict = self.module.state_dict()
+
+        tag = f"global_step{self.global_steps}"
+        tag = str(tag)
+        self.checkpoint_engine.create(tag)
+
+        if dist.get_rank() == 0:
+            self.checkpoint_engine.makedirs(save_dir, exist_ok=True)
+            logger.info(f"Saving model weights to {path}, tag: {tag}")
+            self.checkpoint_engine.save(state_dict, path)
+
+        self.checkpoint_engine.commit(tag)
+
+        return True
+
+    def empty_partition_cache(self):
+        """
+        Release GPU memory consumed by offloaded model parameters.
+        """
+        if hasattr(self.optimizer, 'empty_partition_cache'):
+            self.optimizer.empty_partition_cache()
+            gc.collect()
+            get_accelerator().empty_cache()
diff --git a/opensora/adaptor/modules.py b/opensora/adaptor/modules.py
new file mode 100644
index 000000000..e53599d55
--- /dev/null
+++ b/opensora/adaptor/modules.py
@@ -0,0 +1,23 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def fp32_layer_norm_forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    origin_dtype = inputs.dtype
+    return F.layer_norm(inputs.float(), self.normalized_shape, self.weight.float() if self.weight is not None else None,
+                        self.bias.float() if self.bias is not None else None, self.eps).to(origin_dtype)
+
+
+def fp32_silu_forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.silu(inputs.float(), inplace=self.inplace).to(inputs.dtype)
+
+
+def fp32_gelu_forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.gelu(inputs.float(), approximate=self.approximate).to(inputs.dtype)
+
+
+def replace_with_fp32_forwards():
+    nn.GELU.forward = fp32_gelu_forward
+    nn.SiLU.forward = fp32_silu_forward
+    nn.LayerNorm.forward = fp32_layer_norm_forward
diff --git a/opensora/adaptor/stage_1_and_2.py b/opensora/adaptor/stage_1_and_2.py
new file mode 100644
index 000000000..35f9a6d55
--- /dev/null
+++ b/opensora/adaptor/stage_1_and_2.py
@@ -0,0 +1,2522 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import os
+import pdb
+from deepspeed import comm as dist
+from packaging import version as pkg_version
+from collections import OrderedDict
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from deepspeed.runtime import ZeROOptimizer
+from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler
+from deepspeed.runtime.utils import (bwc_tensor_model_parallel_rank, get_global_norm, empty_cache, see_memory_usage,
+                                     inf, is_model_parallel_parameter, align_dense_tensors, all_gather_dp_groups)
+
+from deepspeed.runtime.zero.config import ZeroStageEnum
+from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
+from deepspeed.ops.adam import DeepSpeedCPUAdam
+from deepspeed.utils import logger
+from deepspeed.moe.utils import is_moe_param
+from deepspeed.git_version_info import version
+
+from deepspeed.runtime.constants import PIPE_REPLICATED
+from deepspeed.accelerator import get_accelerator
+
+from deepspeed.checkpoint.constants import (DS_VERSION, GROUP_PADDINGS, PARTITION_COUNT, LOSS_SCALER,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS, BASE_OPTIMIZER_STATE,
+                                            BASE_OPTIMIZER_STATE_STEP, CLIP_GRAD, ZERO_STAGE, PARAM_SLICE_MAPPINGS)
+from deepspeed.utils import link_hp_params
+from deepspeed.checkpoint import enable_universal_checkpoint
+
+from deepspeed.utils import groups
+
+from opensora.adaptor.zp_manager import zp_manager
+
+# Toggle this to true to enable correctness test
+# with gradient partitioning and without
+pg_correctness_test = False
+
+OPTIMIZER_ALLGATHER_TIMER = 'optimizer_allgather'
+OPTIMIZER_GRADIENTS_TIMER = 'optimizer_gradients'
+OPTIMIZER_STEP_TIMER = 'optimizer_step'
+OPTIMIZER_TIMERS = [OPTIMIZER_ALLGATHER_TIMER, OPTIMIZER_GRADIENTS_TIMER, OPTIMIZER_STEP_TIMER]
+
+
+def input(msg):
+    return
+
+
+def split_half_float_double(tensors):
+    device_type = get_accelerator().device_name()
+    dtypes = [
+        "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type),
+        "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type)
+    ]
+    buckets = []
+    for i, dtype in enumerate(dtypes):
+        bucket = [t for t in tensors if t.type() == dtype]
+        if bucket:
+            buckets.append(bucket)
+    return buckets
+
+
+def isclose(a, b, rtol=1e-09, atol=0.0):
+    return abs(a - b) <= max(rtol * max(abs(a), abs(b)), atol)
+
+
+def lcm(x, y):
+    from fractions import gcd  # or can import gcd from `math` in Python 3
+    return x * y // gcd(x, y)
+
+
+def get_alignment_padding(tensor_list, alignment):
+    num_elements = sum([tensor.numel() for tensor in tensor_list])
+    remainder = num_elements % alignment
+    return (alignment - remainder) if remainder else remainder
+
+
+def move_to_cpu(tensor_list):
+    for tensor in tensor_list:
+        tensor.data = tensor.data.cpu()
+
+
+def print_rank_msg(msg):
+    print(f"rank {dist.get_rank()} - {msg}")
+
+
+def _get_padded_tensor(src_tensor, size):
+    if src_tensor.numel() >= size:
+        return src_tensor
+    padded_tensor = torch.zeros(size, dtype=src_tensor.dtype, device=src_tensor.device)
+    slice_tensor = torch.narrow(padded_tensor, 0, 0, src_tensor.numel())
+    slice_tensor.data.copy_(src_tensor.data)
+    return padded_tensor
+
+
+def contigous_flatten(tensors):
+    return _flatten_dense_tensors([tensor.contiguous() for tensor in tensors])
+
+
+def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, zp_process_group):
+    for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)):
+        partition_id = dist.get_rank(group=zp_process_group[group_id])
+        dp_world_size = dist.get_world_size(group=zp_process_group[group_id])
+        if dp_world_size == 1:
+            # no groups share optimizer states
+            # pipeline parallel with bf16 will default call this even if dp size = 1.
+            continue
+        input_tensor = partitioned_params[partition_id].contiguous()
+        # print(f"call all_gather_into_tensor_dp_groups, input size is {input_tensor.size()}, "
+        #       f"output size is {group_flat.size()}")
+        #
+        # print(f"groups_flat.size = {groups_flat.numel()}")
+        # print(f"partitioned_param_groups = {sum([v.numel() for v in partitioned_param_groups])}")
+        dist.all_gather_into_tensor(group_flat, input_tensor, zp_process_group[group_id])
+
+
+class DeepSpeedZeroOptimizer(ZeROOptimizer):
+    """
+    DeepSpeedZeroOptimizer designed to reduce the memory footprint
+    required for training large deep learning models.
+
+    For more details please see ZeRO: Memory Optimization Towards Training A Trillion Parameter Models
+    https://arxiv.org/abs/1910.02054
+
+    For usage examples, refer to TODO: DeepSpeed Tutorial
+
+    """
+
+    def __init__(self,
+                 init_optimizer,
+                 param_names,
+                 timers,
+                 static_loss_scale=1.0,
+                 dynamic_loss_scale=False,
+                 dynamic_loss_args=None,
+                 verbose=True,
+                 contiguous_gradients=True,
+                 reduce_bucket_size=500000000,
+                 use_multi_rank_bucket_allreduce=True,
+                 allgather_bucket_size=5000000000,
+                 dp_process_group=None,
+                 expert_parallel_group=None,
+                 expert_data_parallel_group=None,
+                 reduce_scatter=True,
+                 overlap_comm=False,
+                 offload_optimizer_config=None,
+                 mpu=None,
+                 clip_grad=0.0,
+                 gradient_accumulation_dtype=torch.float32,
+                 communication_data_type=torch.float16,
+                 postscale_gradients=True,
+                 gradient_predivide_factor=1.0,
+                 gradient_accumulation_steps=1,
+                 ignore_unused_parameters=True,
+                 partition_grads=True,
+                 round_robin_gradients=False,
+                 has_moe_layers=False,
+                 fp16_master_weights_and_gradients=False,
+                 elastic_checkpoint=False):
+
+        if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none:
+            self.cpu_offload = True
+            self.cpu_offload_pin_memory = offload_optimizer_config.pin_memory
+        else:
+            self.cpu_offload = False
+            self.cpu_offload_pin_memory = False
+
+        if dist.get_rank() == 0:
+            logger.info(f"Reduce bucket size {reduce_bucket_size}")
+            logger.info(f"Allgather bucket size {allgather_bucket_size}")
+            logger.info(f"CPU Offload: {self.cpu_offload}")
+            logger.info(f'Round robin gradient partitioning: {round_robin_gradients}')
+        # The fused optimizer does all the work. We need this layer for two reason:
+        # 1. maintain same user API from apex.fp16_utils
+        # 2. keep common stuff here in case we need to add ne552w fused optimizer later
+
+        self.elastic_checkpoint = elastic_checkpoint
+        self.param_names = param_names
+        self.mpu = mpu
+        # differences from apex.fp16_utils:
+        # - assume all model params in fp16
+        # - assume all params requires grad
+        # - flat by groups, not keeping state. TODO: remove state explicitly?
+        # - master grad and unflat master weight never exist. TODO: a way to save out unflat master?
+        if not get_accelerator().is_available():
+            raise SystemError("Accelerator is not detected, cannot perform low precision training (e.g., fp16, bf16).")
+        self.optimizer = init_optimizer
+
+        # Use torch (un)flatten ops
+        self.flatten = contigous_flatten
+        self.unflatten = _unflatten_dense_tensors
+
+        # ZeRO stage 1 (False) or 2 (True)
+        self.partition_gradients = partition_grads
+        self.zero_stage_string = "ZeRO-2" if partition_grads else "ZeRO-1"
+
+        self.timers = timers
+
+        self.reduce_scatter = reduce_scatter
+
+        self.overlap_comm = overlap_comm
+
+        self.deepspeed_adam_offload = self.cpu_offload
+
+        self.device = get_accelerator().current_device_name() if not self.cpu_offload else 'cpu'
+
+        zp_manager.init_group()
+        self.zp_process_group = zp_manager.zp_group
+        zp_rank = dist.get_rank(group=self.zp_process_group)
+        zp_size = dist.get_world_size(group=self.zp_process_group)
+        print(f"zp rank is {zp_rank}, zp_size={zp_size}")
+
+        self.dp_process_group = dp_process_group
+
+        self.sequence_parallel_size = groups._get_sequence_parallel_world_size()
+        # expert parallel group
+        self.ep_process_group = expert_parallel_group
+
+        # data parallel group for experts
+        self.expert_dp_process_group = expert_data_parallel_group
+
+        # data parallel size for non-experts
+        dp_size = dist.get_world_size(group=self.dp_process_group)
+
+        # For MoE models this maybe different for different param group
+        # It will be modified during MoE setup later in the init
+        self.real_zp_process_group = [self.zp_process_group for i in range(len(self.optimizer.param_groups))]
+        self.real_dp_process_group = [self.dp_process_group for i in range(len(self.optimizer.param_groups))]
+        self.partition_count = [zp_manager.zp_size for i in range(len(self.optimizer.param_groups))]
+
+        self.is_gradient_accumulation_boundary = True
+
+        # CPU-Offload requires contiguous gradients
+        self.contiguous_gradients = contiguous_gradients or self.cpu_offload
+
+        self.has_moe_layers = has_moe_layers
+        if self.has_moe_layers:
+            self._configure_moe_settings()
+        self._global_grad_norm = 0.
+
+        if mpu is None:
+            self.model_parallel_group = None
+            self.model_parallel_world_size = 1
+            self.model_parallel_rank = 0
+        else:
+            self.model_parallel_group = mpu.get_model_parallel_group()
+            self.model_parallel_world_size = mpu.get_model_parallel_world_size()
+            self.model_parallel_rank = bwc_tensor_model_parallel_rank(mpu)
+
+        self.overflow = False
+        self.clip_grad = clip_grad
+        self.communication_data_type = communication_data_type
+        self.gradient_predivide_factor = gradient_predivide_factor
+        self.postscale_gradients = postscale_gradients
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.micro_step_id = 0
+        self.ignore_unused_parameters = ignore_unused_parameters
+        self.round_robin_gradients = round_robin_gradients
+
+        self.extra_large_param_to_reduce = None
+        self.fp16_master_weights_and_gradients = fp16_master_weights_and_gradients
+
+        if self.fp16_master_weights_and_gradients:
+            assert self.cpu_offload and type(self.optimizer) in [DeepSpeedCPUAdam], \
+                f"fp16_master_and_gradients requires optimizer to support keeping fp16 master and gradients while keeping the optimizer states in fp32." \
+                f"Currently only supported using ZeRO-Offload with DeepSpeedCPUAdam. But current setting is ZeRO-Offload:{self.cpu_offload} and optimizer type {type(self.optimizer)}." \
+                f"Either disable fp16_master_weights_and_gradients or enable {self.zero_stage_string} Offload with DeepSpeedCPUAdam."
+
+        if self.reduce_scatter:
+            valid_reduce_scatter_dtypes = (torch.float16, torch.bfloat16, torch.float32)
+            assert self.communication_data_type in valid_reduce_scatter_dtypes, f"{self.zero_stage_string} supports {valid_reduce_scatter_dtypes} communication_data_type with reduce scatter enabled. Got: '{self.communication_data_type}'"
+            assert self.gradient_predivide_factor == 1.0, "gradient_predivide_factor != 1.0 is not yet supported with {self.zero_stage_string} with reduce scatter enabled"
+            assert self.postscale_gradients, "pre-scale gradients is not yet supported with {self.zero_stage_string} with reduce scatter enabled"
+
+        # param flattened by groups
+        self.bit16_groups = []
+        self.bit16_groups_flat = []
+
+        # param partitioned by data parallel degree
+        # this will contain a list of equal sized tensors
+        # each of which will be updated by a different process
+        self.parallel_partitioned_bit16_groups = []
+
+        # a single 32-bit partition of the parallel partitioned parameters
+        # that this process will update
+        self.single_partition_of_fp32_groups = []
+
+        # param partition info
+
+        # These are the parameters in each group that will not be updated by this process directly
+        self.params_not_in_partition = []
+
+        # These are the parameters that will be updated by this process directly
+        self.params_in_partition = []
+
+        # Offset from the first parameter in the self.params_in_partition
+        # the parameter boundaries may not align with partition boundaries
+        # so we need to keep track of the offset
+        self.first_offset = []
+
+        # number of elements per partition in each group
+        self.partition_size = []
+
+        # align nccl all-gather send buffers to 4-byte boundary
+        self.nccl_start_alignment_factor = 16  # 4-byte alignment/sizeof(fp16) = 2
+
+        assert (
+                allgather_bucket_size % self.nccl_start_alignment_factor == 0
+        ), f"allgather_bucket_size must be a multiple of nccl_start_alignment_factor, {self.nccl_start_alignment_factor} "
+
+        self.all_reduce_print = False
+        self.dtype = self.optimizer.param_groups[0]['params'][0].dtype
+        self.gradient_accumulation_dtype = gradient_accumulation_dtype
+
+        if self.dtype != self.gradient_accumulation_dtype:
+            self.use_separate_grad_accum = True
+        else:
+            self.use_separate_grad_accum = False
+        if self.use_separate_grad_accum and not self.partition_gradients:
+            self.use_grad_accum_attribute = True
+        else:
+            self.use_grad_accum_attribute = False
+
+        self.round_robin_bit16_groups = []
+        self.round_robin_bit16_indices = []
+
+        # Use different parallel to do all_to_all_reduce related things
+        # padding on each partition for alignment purposes
+        self.groups_padding = []
+        # loop to deal with groups
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            partition_id = dist.get_rank(group=self.real_zp_process_group[i])
+
+            # push this group to list before modify
+            # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group
+            trainable_parameters = []
+            for param in param_group['params']:
+                if param.requires_grad:
+                    param.grad_accum = None
+                    trainable_parameters.append(param)
+            self.bit16_groups.append(trainable_parameters)
+
+            # not sure why apex was cloning the weights before flattening
+            # removing cloning here
+
+            see_memory_usage(f"Before moving param group {i} to CPU")
+            # move all the parameters to cpu to free up GPU space for creating flat buffer
+            move_to_cpu(self.bit16_groups[i])
+            empty_cache()
+            see_memory_usage(f"After moving param group {i} to CPU", force=False)
+
+            # Reorder group parameters for load balancing of gradient partitioning during backward among ranks.
+            # This ensures that gradients are reduced in a fashion such that ownership round robins among the ranks.
+            # For example, rather than 3 gradients (g_n+2, g_n+1, g_n) that are reduced consecutively belonging
+            # to the same rank, instead they will belong to 3 ranks (r_m+2, r_m+1, r_m).
+            if self.round_robin_gradients:
+                round_robin_tensors, round_robin_indices = self._round_robin_reorder(
+                    self.bit16_groups[i], dist.get_world_size(group=self.real_zp_process_group[i]))
+            else:
+                round_robin_tensors = self.bit16_groups[i]
+                round_robin_indices = list(range(len(self.bit16_groups[i])))
+
+            self.round_robin_bit16_groups.append(round_robin_tensors)
+            self.round_robin_bit16_indices.append(round_robin_indices)
+
+            # create flat buffer in CPU and move to GPU
+            self.bit16_groups_flat.append(
+                self.flatten_dense_tensors_aligned(
+                    self.round_robin_bit16_groups[i],
+                    self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_zp_process_group[i])).to(
+                    get_accelerator().current_device_name()))
+            see_memory_usage(f"After flattening and moving param group {i} to GPU", force=False)
+
+            # Record padding required for alignment
+            if partition_id == dist.get_world_size(group=self.real_zp_process_group[i]) - 1:
+                padding = self.bit16_groups_flat[i].numel() - sum(
+                    [t.numel() for t in self.round_robin_bit16_groups[i]])
+            else:
+                padding = 0
+            self.groups_padding.append(padding)
+
+            if dist.get_rank(group=self.real_zp_process_group[i]) == 0:
+                see_memory_usage(f"After Flattening and after emptying param group {i} cache", force=False)
+
+            # set model bit16 weight to slices of flattened buffer
+            self._update_model_bit16_weights(i)
+
+            # divide the flat weights into near equal partition equal to the data parallel degree
+            # each process will compute on a different part of the partition
+            data_parallel_partitions = self.get_data_parallel_partitions(self.bit16_groups_flat[i], i)
+            self.parallel_partitioned_bit16_groups.append(data_parallel_partitions)
+
+            # print(f"self.bit16_groups_flat[i].size = {self.bit16_groups_flat[i].numel()}")
+            # print(f"data_parallel_partitions = {sum([v.numel() for v in data_parallel_partitions])}")
+
+            # verify that data partition start locations are 4-byte aligned
+            for partitioned_data in data_parallel_partitions:
+                assert (partitioned_data.data_ptr() % (2 * self.nccl_start_alignment_factor) == 0)
+
+            # A partition of the fp32 master weights that will be updated by this process.
+            # Note that the params in single_partition_of_fp32_groups is cloned and detached
+            # from the origin params of the model.
+            if not fp16_master_weights_and_gradients:
+                self.single_partition_of_fp32_groups.append(self.parallel_partitioned_bit16_groups[i][partition_id].to(
+                    self.device).clone().float().detach())
+            else:
+                self.single_partition_of_fp32_groups.append(self.parallel_partitioned_bit16_groups[i][partition_id].to(
+                    self.device).clone().half().detach())
+
+            # Set local optimizer to have flat params of its own partition.
+            # After this, the local optimizer will only contain its own partition of params.
+            # In that case, the local optimizer only saves the states(momentum, variance, etc.) related to its partition's params(zero stage1).
+            self.single_partition_of_fp32_groups[
+                i].requires_grad = True  # keep this in case internal optimizer uses it
+            param_group['params'] = [self.single_partition_of_fp32_groups[i]]
+
+            partition_size = len(self.bit16_groups_flat[i]) / dist.get_world_size(group=self.real_zp_process_group[i])
+            params_in_partition, params_not_in_partition, first_offset = self.get_partition_info(
+                self.round_robin_bit16_groups[i], partition_size, partition_id)
+
+            self.partition_size.append(partition_size)
+            self.params_in_partition.append(params_in_partition)
+            self.params_not_in_partition.append(params_not_in_partition)
+            self.first_offset.append(first_offset)
+
+        self.reduce_bucket_size = int(reduce_bucket_size)
+        self.use_multi_rank_bucket_allreduce = use_multi_rank_bucket_allreduce
+        self.allgather_bucket_size = int(allgather_bucket_size)
+
+        self.reduction_stream = None if get_accelerator().is_synchronized_device() else get_accelerator().Stream()
+        # self.copy_grad_stream = get_accelerator().Stream()
+        self.callback_queued = False
+
+        self.param_dict = {}
+
+        # map between param_id and bool to specify if a param is in this partition
+        self.is_param_in_current_partition = {}
+
+        self.grads_in_ipg_bucket = []
+        self.params_in_ipg_bucket = []
+        self.elements_in_ipg_bucket = 0
+        self.params_already_reduced = []
+        self._release_ipg_buffers()
+        self.previous_reduced_grads = None
+        self.ipg_bucket_has_moe_params = False
+
+        # simplified param id
+        self.param_id = {}
+
+        # interesting code: unique ids being assigned to individual parameters
+        largest_param_numel = 0
+        count = 0
+        for i, params_group in enumerate(self.bit16_groups):
+            for param in params_group:
+                unique_id = id(param)
+                self.param_id[unique_id] = count
+                self.param_dict[count] = param
+                self.params_already_reduced.append(False)
+                if param.numel() > largest_param_numel:
+                    largest_param_numel = param.numel()
+                count = count + 1
+
+        for param_group in self.params_in_partition:
+            for param in param_group:
+                self.is_param_in_current_partition[self.get_param_id(param)] = True
+
+        for param_group in self.params_not_in_partition:
+            for param in param_group:
+                self.is_param_in_current_partition[self.get_param_id(param)] = False
+
+        if self.cpu_offload:
+            self.accumulated_grads_in_cpu = {}
+            self.norm_for_param_grads = {}
+            self.local_overflow = False
+            self.grad_position = {}
+            self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel,
+                                                                device=self.device,
+                                                                dtype=self.dtype)
+            if self.cpu_offload_pin_memory:
+                self.temp_grad_buffer_for_cpu_offload = get_accelerator().pin_memory(
+                    self.temp_grad_buffer_for_cpu_offload)
+            self.temp_grad_buffer_for_gpu_offload = torch.zeros(largest_param_numel,
+                                                                device=get_accelerator().current_device_name(),
+                                                                dtype=self.dtype)
+            for i, params_group in enumerate(self.bit16_groups):
+                self.get_grad_position(i, self.params_in_partition[i], self.first_offset[i], self.partition_size[i])
+
+        # mapping from parameter to partition that it belongs to
+        self.param_to_partition_ids = {}
+
+        # stores if a partition has been reduced in this step
+        self.is_partition_reduced = {}
+
+        # number of grads in partition that still need to be computed
+        self.remaining_grads_in_partition = {}
+
+        # total number of grads in partition
+        self.total_grads_in_partition = {}
+
+        # stores if a grad in a partition has been computed or not
+        self.is_grad_computed = {}
+
+        # stores the offset at which a parameter gradient needs to be inserted in a partition
+        self.grad_partition_insertion_offset = {}
+
+        # the offset in the gradient at which it must be inserted at the beginning of the partition
+        self.grad_start_offset = {}
+
+        # will store the averaged gradients required by this partition
+        self.averaged_gradients = {}
+
+        # For cpu_offload, will store the averaged gradients required by this partition
+        self.offload_gradient_dict = {}
+
+        # store index of first parameter in each partition
+        self.first_param_index_in_partition = {}
+
+        # initializes all data structures for implementing gradient partitioning
+        self.initialize_gradient_partitioning_data_structures()
+
+        # resets the data structure value for the next backward propagation
+        self.reset_partition_gradient_structures()
+
+        # creates backward hooks for gradient partitioning
+        if self.partition_gradients or self.overlap_comm:
+            self.create_reduce_and_remove_grad_hooks()
+
+        self.custom_loss_scaler = False
+        self.external_loss_scale = None
+
+        # we may have a way of fusing dynamic scale. Do not support for now
+        self.loss_scaler = CreateLossScaler(dtype=self.dtype,
+                                            static_loss_scale=static_loss_scale,
+                                            dynamic_scaling=dynamic_loss_scale,
+                                            dynamic_loss_args=dynamic_loss_args)
+        self.dynamic_loss_scale = self.loss_scaler.dynamic
+
+        if self.dtype != torch.float16:
+            # Only fp16 should use dynamic loss scaling
+            assert self.loss_scaler.cur_scale == 1.0
+            assert not self.dynamic_loss_scale
+
+        see_memory_usage("Before initializing optimizer states", force=True)
+        self.initialize_optimizer_states()
+        see_memory_usage("After initializing optimizer states", force=True)
+
+        if dist.get_rank() == 0:
+            logger.info(f"optimizer state initialized")
+
+        if dist.get_rank(group=self.zp_process_group) == 0:
+            see_memory_usage(f"After initializing ZeRO optimizer", force=True)
+
+        self._link_all_hp_params()
+        self._enable_universal_checkpoint()
+        self._param_slice_mappings = self._create_param_mapping()
+
+    def _enable_universal_checkpoint(self):
+        for lp_param_group in self.bit16_groups:
+            enable_universal_checkpoint(param_list=lp_param_group)
+
+    def _create_param_mapping(self):
+        param_mapping = []
+        for i, _ in enumerate(self.optimizer.param_groups):
+            param_mapping_per_group = OrderedDict()
+            for lp in self.bit16_groups[i]:
+                if lp._hp_mapping is not None:
+                    lp_name = self.param_names[lp]
+                    param_mapping_per_group[lp_name] = lp._hp_mapping.get_hp_fragment_address()
+            param_mapping.append(param_mapping_per_group)
+
+        return param_mapping
+
+    def _link_all_hp_params(self):
+        dp_world_size = dist.get_world_size(group=self.zp_process_group)
+        if self.cpu_offload:
+            self._get_offload_gradient_dict()
+
+        for i, _ in enumerate(self.optimizer.param_groups):
+            # Link bit16 and fp32 params in partition
+            partition_id = dist.get_rank(group=self.real_zp_process_group[i])
+            partition_size = self.bit16_groups_flat[i].numel() // dp_world_size
+            flat_hp_partition = self.single_partition_of_fp32_groups[i]
+            link_hp_params(lp_param_list=self.bit16_groups[i],
+                           flat_hp_partition=flat_hp_partition,
+                           gradient_dict=self.averaged_gradients,
+                           offload_gradient_dict=self.offload_gradient_dict,
+                           use_offload=self.cpu_offload,
+                           param_group_index=i,
+                           partition_start=partition_id * partition_size,
+                           partition_size=partition_size,
+                           partition_optimizer_state=self.optimizer.state[flat_hp_partition],
+                           dp_group=self.real_zp_process_group[i])
+
+    def is_moe_group(self, group):
+        return 'moe' in group and group['moe']
+
+    def _configure_moe_settings(self):
+        # if we're using ZeRO stage 2, ensure contiguous gradients are used
+        if self.partition_gradients:
+            assert self.contiguous_gradients, "Contiguous Gradients in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE"
+        # NOTE: To run ZeRO stage 1 with MoE, we need to set self.contiguous_gradients to True or ignore the assertion
+        if not self.partition_gradients and not self.contiguous_gradients:
+            logger.warn(
+                "ZeRO Stage 1 has not been thoroughly tested with MoE. This configuration is still experimental.")
+        assert self.reduce_scatter, "Reduce Scatter in ZeRO Stage 2 must be set to True for MoE. Other code paths are not tested with MoE"
+
+        assert any(
+            [self.is_moe_group(group) for group in self.optimizer.param_groups]
+        ), "The model has moe layers, but None of the param groups are marked as MoE. Create a param group with 'moe' key set to True before creating optimizer"
+        self.is_moe_param_group = []
+        for i, group in enumerate(self.optimizer.param_groups):
+            if self.is_moe_group(group):
+                assert all([is_moe_param(param)
+                            for param in group['params']]), "All params in MoE group must be MoE params"
+                self.real_zp_process_group[i] = self.expert_dp_process_group[group['name']]
+                self.partition_count[i] = dist.get_world_size(group=self.expert_dp_process_group[group['name']])
+                self.is_moe_param_group.append(True)
+            else:
+                self.is_moe_param_group.append(False)
+
+        assert self.expert_dp_process_group is not None, "Expert data parallel group should be configured with MoE"
+        assert self.ep_process_group is not None, "Expert parallel group should be configured with MoE"
+
+    def _update_model_bit16_weights(self, group_index):
+        updated_params = self.unflatten(self.bit16_groups_flat[group_index],
+                                        self.round_robin_bit16_groups[group_index])
+        for p, q in zip(self.round_robin_bit16_groups[group_index], updated_params):
+            p.data = q.data
+
+        # set model fp16 weight to slices of reordered flattened buffer
+        for param_index, param in enumerate(self.bit16_groups[group_index]):
+            new_index = self.round_robin_bit16_indices[group_index][param_index]
+            param.data = self.round_robin_bit16_groups[group_index][new_index].data
+
+    def _round_robin_reorder(self, tensor_list, num_partitions):
+
+        # disable round robin if need to debug something
+        # return tensor_list, list(range(len(tensor_list)))
+
+        partition_tensors = {}
+
+        for i, tensor in enumerate(tensor_list):
+            j = i % num_partitions
+            if not j in partition_tensors:
+                partition_tensors[j] = []
+            partition_tensors[j].append((i, tensor))
+
+        reordered_tensors = []
+        reordered_indices = {}
+
+        for partition_index in partition_tensors.keys():
+            for i, (original_index, tensor) in enumerate(partition_tensors[partition_index]):
+                reordered_indices[original_index] = len(reordered_tensors)
+                reordered_tensors.append(tensor)
+
+        return reordered_tensors, reordered_indices
+
+    def _release_ipg_buffers(self):
+        if self.contiguous_gradients:
+            self.ipg_buffer = None
+            self.grads_in_partition = None
+            self.grads_in_partition_offset = 0
+
+    def initialize_optimizer_states(self):
+
+        for i, group in enumerate(self.bit16_groups):
+            single_grad_partition = torch.zeros(int(self.partition_size[i]),
+                                                dtype=self.single_partition_of_fp32_groups[i].dtype,
+                                                device=self.device)
+            self.single_partition_of_fp32_groups[i].grad = get_accelerator().pin_memory(
+                single_grad_partition) if self.cpu_offload_pin_memory else single_grad_partition
+
+        # Initialize the optimizer states with the flattened fp32 partition.
+        # State initialization for the Adagrad optimizer occurs at construction as opposed to other optimizers
+        # which do lazy initialization of the state at the first call to step.
+        if isinstance(self.optimizer, torch.optim.Adagrad):
+            self.optimizer = torch.optim.Adagrad(self.single_partition_of_fp32_groups, **self.optimizer.defaults)
+        else:
+            self.optimizer.step()
+
+        if not self.cpu_offload:
+            for group in self.single_partition_of_fp32_groups:
+                group.grad = None  # class init
+
+        return
+
+    #########################################################################
+    #################### ZeRO Stage 1 - reduce gradients ####################
+    #########################################################################
+    def reduce_gradients(self, pipeline_parallel=False):
+        world_size = dist.get_world_size(self.zp_process_group)
+        my_rank = dist.get_rank(self.zp_process_group)
+
+        # with PP we must create ipg buffer, since backward is handled outside zero
+        if pipeline_parallel and self.contiguous_gradients:
+            self.ipg_buffer = []
+            buf_0 = torch.empty(int(self.reduce_bucket_size),
+                                dtype=self.dtype,
+                                device=get_accelerator().current_device_name())
+            self.ipg_buffer.append(buf_0)
+            self.ipg_index = 0
+
+        if not self.overlap_comm:
+            for i, group in enumerate(self.bit16_groups):
+                for param in group:
+                    grad_reduc = self.get_gradient_for_reduction(param)
+                    if grad_reduc is not None:
+                        self.reduce_ready_partitions_and_remove_grads(param, i)
+        # reduce any pending grads in either hook/non-hook case
+        self.overlapping_partition_gradients_reduce_epilogue()
+
+    #########################################################################
+    #########################ZeRO Partition Gradients########################
+    #########################################################################
+
+    def get_first_param_index(self, group_id, param_group, partition_id):
+        for index, param in enumerate(param_group):
+            param_id = self.get_param_id(param)
+            if partition_id in self.param_to_partition_ids[group_id][param_id]:
+                return index
+        return None
+
+    def initialize_gradient_partitioning_data_structures(self):
+
+        for i, param_group in enumerate(self.round_robin_bit16_groups):
+            total_partitions = dist.get_world_size(group=self.real_zp_process_group[i])
+
+            self.param_to_partition_ids[i] = {}
+            self.is_partition_reduced[i] = {}
+            self.total_grads_in_partition[i] = {}
+            self.remaining_grads_in_partition[i] = {}
+            self.is_grad_computed[i] = {}
+            self.grad_partition_insertion_offset[i] = {}
+            self.grad_start_offset[i] = {}
+            self.first_param_index_in_partition[i] = {}
+
+            for partition_id in range(total_partitions):
+                self.is_grad_computed[i][partition_id] = {}
+                self.grad_partition_insertion_offset[i][partition_id] = {}
+                self.grad_start_offset[i][partition_id] = {}
+                self.total_grads_in_partition[i][partition_id] = 0
+                self.initialize_gradient_partition(i, param_group, partition_id)
+                self.is_partition_reduced[i][partition_id] = False
+                self.first_param_index_in_partition[i][partition_id] = self.get_first_param_index(
+                    i, param_group, partition_id)
+
+    def independent_gradient_partition_epilogue(self):
+        self.report_ipg_memory_usage(f"In ipg_epilogue before reduce_ipg_grads", 0)
+        self.reduce_ipg_grads()
+        self.report_ipg_memory_usage(f"In ipg_epilogue after reduce_ipg_grads", 0)
+
+        # if dist.get_rank() == 0:
+        #    logger.info("Params already reduced %s", self.params_already_reduced)
+        for i in range(len(self.params_already_reduced)):
+            self.params_already_reduced[i] = False
+
+        if self.overlap_comm:
+            get_accelerator().synchronize()
+            # It is safe to clear previously reduced grads of other partitions
+            self._clear_previous_reduced_grads()
+
+        if self.cpu_offload is False:
+            for i, _ in enumerate(self.bit16_groups):
+
+                if not i in self.averaged_gradients or self.averaged_gradients[i] is None:
+                    self.averaged_gradients[i] = self.get_flat_partition(
+                        self.params_in_partition[i],
+                        self.first_offset[i],
+                        self.partition_size[i],
+                        dtype=self.gradient_accumulation_dtype,
+                        device=get_accelerator().current_device_name(),
+                        return_tensor_list=True)
+                else:
+                    avg_new = self.get_flat_partition(self.params_in_partition[i],
+                                                      self.first_offset[i],
+                                                      self.partition_size[i],
+                                                      dtype=self.gradient_accumulation_dtype,
+                                                      device=get_accelerator().current_device_name(),
+                                                      return_tensor_list=True)
+
+                    for accumulated_grad, new_avg_grad in zip(self.averaged_gradients[i], avg_new):
+                        accumulated_grad.add_(new_avg_grad)
+
+        self._release_ipg_buffers()
+
+        # No need to keep the gradients anymore.
+        # All gradients required by the step
+        # are in self.averaged_gradients
+        self.zero_grad(set_to_none=True)
+        see_memory_usage(f"End ipg_epilogue")
+
+    # resets all partition to no reduced
+    # sets remaining grads to the total number of grads in each partition
+    # set is grad computed to false for all grads in partition
+    def reset_partition_gradient_structures(self):
+        for i, _ in enumerate(self.bit16_groups):
+            total_partitions = dist.get_world_size(group=self.real_zp_process_group[i])
+            for partition_id in range(total_partitions):
+                self.is_partition_reduced[i][partition_id] = False
+                self.remaining_grads_in_partition[i][partition_id] = self.total_grads_in_partition[i][partition_id]
+
+                for param_id in self.is_grad_computed[i][partition_id]:
+                    self.is_grad_computed[i][partition_id][param_id] = False
+
+    def initialize_gradient_partition(self, i, param_group, partition_id):
+
+        def set_key_value_list(dictionary, key, value):
+            if key in dictionary:
+                dictionary[key].append(value)
+            else:
+                dictionary[key] = [value]
+
+        def increment_value(dictionary, key):
+            if key in dictionary:
+                dictionary[key] += 1
+            else:
+                dictionary[key] = 1
+
+        partition_size = self.partition_size[i]
+
+        start_index = partition_size * partition_id
+        end_index = partition_size * (partition_id + 1)
+
+        current_index = 0
+        first_offset = 0
+
+        for param in param_group:
+
+            param_size = param.numel()
+            param_id = self.get_param_id(param)
+
+            if start_index <= current_index < end_index:
+                set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id)
+                increment_value(self.total_grads_in_partition[i], partition_id)
+
+                self.is_grad_computed[i][partition_id][param_id] = False
+
+                self.grad_partition_insertion_offset[i][partition_id][param_id] = current_index - start_index
+                self.grad_start_offset[i][partition_id][param_id] = 0
+
+            elif current_index < start_index < (current_index + param_size):
+                assert (first_offset == 0
+                        ), "This can happen either zero or only once as this must be the first tensor in the partition"
+                first_offset = start_index - current_index
+
+                set_key_value_list(self.param_to_partition_ids[i], param_id, partition_id)
+                increment_value(self.total_grads_in_partition[i], partition_id)
+
+                self.is_grad_computed[i][partition_id][param_id] = False
+
+                self.grad_partition_insertion_offset[i][partition_id][param_id] = 0
+                self.grad_start_offset[i][partition_id][param_id] = first_offset
+
+            current_index = current_index + param_size
+
+    def overlapping_partition_gradients_reduce_epilogue(self):
+        self.independent_gradient_partition_epilogue()
+
+    def fill_grad_accum_attribute(self):
+        for group in self.bit16_groups:
+            for param in group:
+                if param.grad is not None:
+                    if param.grad_accum is None:
+                        param.grad_accum = param.grad.to(self.gradient_accumulation_dtype)
+                    else:
+                        param.grad_accum.add_(
+                            param.grad.to(self.gradient_accumulation_dtype).view(param.grad_accum.shape))
+                    param.grad = None
+
+    def get_gradient_for_reduction(self, param):
+        if self.use_grad_accum_attribute:
+            return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None
+        else:
+            return param.grad
+
+    def get_param_gradient_attribute(self, param):
+        return param.grad_accum if self.use_grad_accum_attribute else param.grad
+
+    # Clear the tensor the reduction gradient attribute is pointing to
+    def clear_grad_attribute(self, param):
+        if self.use_grad_accum_attribute:
+            param.grad_accum = None
+        else:
+            param.grad = None
+
+    def create_reduce_and_remove_grad_hooks(self):
+        self.grad_accs = []
+        for i, param_group in enumerate(self.bit16_groups):
+            for param in param_group:
+                if param.requires_grad:
+                    def wrapper(param, i):
+                        param_tmp = param.expand_as(param)
+                        grad_acc = param_tmp.grad_fn.next_functions[0][0]
+
+                        def reduce_partition_and_remove_grads(*notneeded):
+                            self.reduce_ready_partitions_and_remove_grads(param, i)
+
+                        grad_acc.register_hook(reduce_partition_and_remove_grads)
+                        self.grad_accs.append(grad_acc)
+
+                    wrapper(param, i)
+
+    def get_param_id(self, param):
+        unique_id = id(param)
+        return self.param_id[unique_id]
+
+    def report_ipg_memory_usage(self, tag, param_elems):
+        elem_count = self.elements_in_ipg_bucket + param_elems
+        percent_of_bucket_size = (100.0 * elem_count) // self.reduce_bucket_size
+        see_memory_usage(
+            f"{tag}: elems in_bucket {self.elements_in_ipg_bucket} param {param_elems} max_percent {percent_of_bucket_size}"
+        )
+
+    # create a flat tensor aligned at the alignment boundary
+    def flatten_dense_tensors_aligned(self, tensor_list, alignment):
+        return self.flatten(align_dense_tensors(tensor_list, alignment))
+
+    ############### Independent Partition Gradient ########################
+    def reduce_independent_p_g_buckets_and_remove_grads(self, param, i):
+
+        grad_reduc = self.get_gradient_for_reduction(param)
+        if self.elements_in_ipg_bucket + param.numel() > self.reduce_bucket_size:
+            self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads", param.numel())
+            self.reduce_ipg_grads()
+            if self.contiguous_gradients and self.overlap_comm:
+                # Swap ipg_index between 0 and 1
+                self.ipg_index = 1 - self.ipg_index
+            self.report_ipg_memory_usage("In ipg_remove_grads after reduce_ipg_grads", param.numel())
+
+        param_id = self.get_param_id(param)
+        assert self.params_already_reduced[param_id] == False, \
+            f"The parameter {param_id} has already been reduced. \
+            Gradient computed twice for this partition. \
+            Multiple gradient reduction is currently not supported"
+
+        if self.contiguous_gradients:
+            if param.numel() > self.reduce_bucket_size:
+                self.extra_large_param_to_reduce = param
+            else:
+                # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening
+                new_grad_tensor = self.ipg_buffer[self.ipg_index].narrow(0, self.elements_in_ipg_bucket, param.numel())
+                new_grad_tensor.copy_(grad_reduc.view(-1))
+                grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc)
+
+        self.elements_in_ipg_bucket += param.numel()
+
+        assert grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient"
+
+        self.grads_in_ipg_bucket.append(grad_reduc)
+        self.params_in_ipg_bucket.append((i, param, param_id))
+
+        # make sure the average tensor function knows how to average the gradients
+        if is_moe_param(param):
+            self.ipg_bucket_has_moe_params = True
+
+        self.report_ipg_memory_usage("End ipg_remove_grads", 0)
+
+    def print_rank_0(self, message):
+        if dist.get_rank() == 0:
+            logger.info(message)
+
+    def gradient_reduction_w_predivide(self, tensor):
+
+        dp_world_size = dist.get_world_size(group=self.dp_process_group)
+
+        tensor_to_allreduce = tensor
+
+        if self.communication_data_type != tensor.dtype:
+            tensor_to_allreduce = tensor.to(self.communication_data_type)
+
+        if self.postscale_gradients:
+            if self.gradient_predivide_factor != 1.0:
+                tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor)
+
+            dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group)
+
+            if self.gradient_predivide_factor != dp_world_size:
+                tensor_to_allreduce.mul_(self.gradient_predivide_factor /
+                                         (dp_world_size / float(self.sequence_parallel_size)))
+        else:
+            tensor_to_allreduce.div_(dp_world_size / float(self.sequence_parallel_size))
+            dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group)
+
+        if self.communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce:
+            tensor.copy_(tensor_to_allreduce)
+
+        return tensor
+
+    def allreduce_and_copy_with_multiple_ranks(self,
+                                               small_bucket,
+                                               log=None,
+                                               divide=True,
+                                               process_group=None,
+                                               bucket_ranks=None):
+        process_group = self.zp_process_group if process_group is None else process_group
+        allreduced = self.allreduce_bucket(small_bucket, log=log, divide=divide, process_group=process_group)
+        for buf, synced, bucket_rank in zip(small_bucket, self.unflatten(allreduced, small_bucket), bucket_ranks):
+            if dist.get_rank(group=process_group) == bucket_rank:
+                buf.copy_(synced)
+
+    def allreduce_and_scatter(self, bucket, numel_per_bucket=500000000, log=None, divide=True, process_group=None):
+        small_bucket = []
+        small_bucket_ranks = []
+        numel = 0
+        allreduce_sizes = []
+
+        for i, bucket_elem in enumerate(bucket):
+            rank, tensor = bucket_elem
+            small_bucket.append(tensor)
+            small_bucket_ranks.append(rank)
+            numel = numel + tensor.numel()
+            if numel > numel_per_bucket:
+                self.allreduce_and_copy_with_multiple_ranks(small_bucket,
+                                                            log=None,
+                                                            divide=divide,
+                                                            process_group=process_group,
+                                                            bucket_ranks=small_bucket_ranks)
+                small_bucket = []
+                small_bucket_ranks = []
+                numel = 0
+
+        if len(small_bucket) > 0:
+            self.allreduce_and_copy_with_multiple_ranks(small_bucket,
+                                                        log=None,
+                                                        divide=divide,
+                                                        process_group=process_group,
+                                                        bucket_ranks=small_bucket_ranks)
+
+    def average_tensor(self, tensor):
+        if self.overlap_comm:
+            stream = self.reduction_stream
+            if not get_accelerator().is_synchronized_device():
+                stream.wait_stream(get_accelerator().current_stream())
+        else:
+            stream = get_accelerator().current_stream()
+
+        with get_accelerator().stream(stream):
+            if not self.reduce_scatter:
+                self.gradient_reduction_w_predivide(tensor)
+                return
+
+            # Accumulate destination ranks and bucket offsets for each gradient slice.
+            # Note: potential future optimization, record access pattern of parameters
+            # in backward pass and partition gradients w.r.t. access pattern so that our
+            # bucket is guaranteed to be contiguous w.r.t. ranks
+            rank_and_offsets = []
+            real_dp_process_group = []
+            curr_size = 0
+            prev_id, prev_process_group = -1, None
+
+            process_group = self.zp_process_group
+            # count = 0
+            for i, param, param_id in self.params_in_ipg_bucket:
+
+                process_group = self.zp_process_group
+                grad_reduc = self.get_gradient_for_reduction(param)
+                # Averages gradients at parameter level if ipg has a moe param
+                # Otherwise averaging is done at the entire buffer level at the end of the loop
+                # MoE param have different groups
+                if self.ipg_bucket_has_moe_params:
+                    process_group = self.expert_dp_process_group[param.group_name] if is_moe_param(
+                        param) else self.zp_process_group
+                    grad_reduc.data.div_(dist.get_world_size(group=process_group) / float(self.sequence_parallel_size))
+
+                partition_ids = self.param_to_partition_ids[i][param_id]
+                assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids
+                            ]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}"
+                partition_size = self.partition_size[i]
+                # Get all partition ids + their offsets
+                partition_ids_w_offsets = []
+                for partition_id in partition_ids:
+                    offset = self.grad_start_offset[i][partition_id][param_id]
+                    partition_ids_w_offsets.append((partition_id, offset))
+                partition_ids_w_offsets.sort(key=lambda t: t[1])
+
+                # Calculate rank and offsets for grad slices
+                for idx in range(len(partition_ids_w_offsets)):
+                    partition_id, offset = partition_ids_w_offsets[idx]
+
+                    # if dist.get_rank() == 0 and count < 100:
+                    #     print(f"Rank {dist.get_rank()} rank offset id {idx} calculated dp size {dist.get_world_size(group=process_group)} real dp size {dist.get_world_size(self.real_dp_process_group[i])} and dst: {partition_id}")
+                    # count += 1
+
+                    # Calculate numel for grad slice depending on partition location
+                    if idx == len(partition_ids_w_offsets) - 1:
+                        # Last partition_id uses its own offset
+                        numel = param.numel() - offset
+                    else:
+                        # Set numel to next partition's offset
+                        numel = partition_ids_w_offsets[idx + 1][1] - offset
+
+                    # Merge bucket ranges if they belong to the same rank
+                    if partition_id == prev_id and process_group == prev_process_group:
+                        prev_pid, prev_size, prev_numel = rank_and_offsets[-1]
+                        rank_and_offsets[-1] = (prev_pid, prev_size, prev_numel + numel)
+                    else:
+                        rank_and_offsets.append((partition_id, curr_size, numel))
+                        real_dp_process_group.append(process_group)
+                    curr_size += numel
+                    prev_id, prev_process_group = partition_id, process_group
+
+            if not self.ipg_bucket_has_moe_params:
+                tensor.div_(dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size))
+
+            buckets = {}
+            for i, (dst, bucket_offset, numel) in enumerate(rank_and_offsets):
+                grad_slice = tensor.narrow(0, int(bucket_offset), int(numel))
+                bucket_key = real_dp_process_group[i] if self.use_multi_rank_bucket_allreduce else (
+                    dst, real_dp_process_group[i])
+                if bucket_key not in buckets:
+                    buckets[bucket_key] = []
+                if self.use_multi_rank_bucket_allreduce:
+                    buckets[bucket_key].append((dst, grad_slice))
+                else:
+                    buckets[bucket_key].append(grad_slice)
+
+            for bucket_key in buckets:
+                if self.use_multi_rank_bucket_allreduce:
+                    self.allreduce_and_scatter(buckets[bucket_key],
+                                               numel_per_bucket=self.reduce_bucket_size,
+                                               divide=self.ipg_bucket_has_moe_params,
+                                               process_group=bucket_key)
+                else:
+                    dst, process_group = bucket_key
+                    self.allreduce_no_retain(buckets[bucket_key],
+                                             numel_per_bucket=self.reduce_bucket_size,
+                                             rank=dst,
+                                             divide=self.ipg_bucket_has_moe_params,
+                                             process_group=process_group)
+
+    ##############################################################################
+    ############################# CPU Offload Methods#############################
+    ##############################################################################
+    def get_grad_position(self, group_id, tensor_list, first_offset, partition_size):
+        current_offset = 0
+
+        for i, tensor in enumerate(tensor_list):
+            param_id = self.get_param_id(tensor)
+            param_start_offset = 0
+
+            num_elements = tensor.numel()
+
+            # we need to offset to get to the right element
+            if i == 0 and first_offset > 0:
+                tensor_offset = first_offset
+                num_elements = num_elements - tensor_offset
+                param_start_offset = first_offset
+
+            # we dont need all elements of the tensor
+            if num_elements > (partition_size - current_offset):
+                num_elements = partition_size - current_offset
+
+            self.grad_position[param_id] = [
+                int(group_id), int(param_start_offset),
+                int(current_offset), int(num_elements)
+            ]
+            current_offset += num_elements
+
+    def update_overflow_tracker_for_param_grad(self, param):
+        grad_accum = self.get_param_gradient_attribute(param)
+        if grad_accum is not None and self._has_inf_or_nan(grad_accum.data):
+            self.local_overflow = True
+
+    def _get_offload_gradient_dict(self):
+        for param_group_index, _ in enumerate(self.optimizer.param_groups):
+            self.offload_gradient_dict[param_group_index] = []
+            for lp_param in self.params_in_partition[param_group_index]:
+                param_id = self.get_param_id(lp_param)
+                [_, _, dest_offset, num_elements] = self.grad_position[param_id]
+                dest_tensor = self.single_partition_of_fp32_groups[param_group_index].grad.view(-1).narrow(
+                    0, dest_offset, num_elements)
+                self.offload_gradient_dict[param_group_index].append(dest_tensor)
+
+    def async_accumulate_grad_in_cpu_via_gpu(self, param):
+        param_id = self.get_param_id(param)
+
+        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
+
+        # copy to a preexisiting buffer to avoid memory allocation penalty
+        dest_buffer = self.temp_grad_buffer_for_gpu_offload.view(-1).narrow(0, 0, param.numel())
+
+        # buffer for storing gradients for this parameter in CPU
+        def buffer_to_accumulate_to_in_cpu():
+            if not self.fp16_master_weights_and_gradients:
+                buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device)
+                return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer
+            else:
+                return self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements)
+
+        # accumulate gradients into param.grad_accum or parts of it that belongs to this partition
+        def accumulate_gradients():
+            grad_accum = self.get_param_gradient_attribute(param)
+            if not self.fp16_master_weights_and_gradients:
+                dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True)
+                grad_accum.data.view(-1).add_(dest_buffer)
+            else:
+                dest_buffer.narrow(0, source_offset,
+                                   num_elements).copy_(self.accumulated_grads_in_cpu[param_id].view(-1),
+                                                       non_blocking=True)
+                grad_accum.data.view(-1).narrow(0, source_offset,
+                                                num_elements).add_(dest_buffer.narrow(0, source_offset, num_elements))
+
+        # move accumulated gradients back to CPU
+        def copy_gradients_to_cpu():
+            grad_accum = self.get_param_gradient_attribute(param)
+            if not self.fp16_master_weights_and_gradients:
+                self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True)
+            else:
+                self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1).narrow(
+                    0, source_offset, num_elements),
+                    non_blocking=True)
+
+        if param_id not in self.accumulated_grads_in_cpu:
+            self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu()
+
+        if self.micro_step_id > 0:
+            accumulate_gradients()
+
+        # at the boundary we will send 32bit directly
+        if not self.is_gradient_accumulation_boundary:
+            copy_gradients_to_cpu()
+
+    def set_norm_for_param_grad(self, param):
+        param_id = self.get_param_id(param)
+        grad_accum = self.get_param_gradient_attribute(param)
+        accumulated_grad = self.accumulated_grads_in_cpu[
+            param_id] if self.gradient_accumulation_steps > 1 else grad_accum
+
+        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
+
+        start = source_offset
+        accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements)
+
+        self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2)
+
+    def set_norm_for_param_grad_in_gpu(self, param):
+        param_id = self.get_param_id(param)
+        grad_accum = self.get_param_gradient_attribute(param)
+        if grad_accum is None:
+            accumulated_grad = param.grad
+        else:
+            accumulated_grad = grad_accum
+
+        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
+
+        start = source_offset
+        accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements)
+
+        self.norm_for_param_grads[param_id] = accumulated_grad.data.double().norm(2)
+
+    def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param):
+        param_id = self.get_param_id(param)
+
+        [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
+
+        dest_tensor = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements)
+
+        grad_accum = self.get_param_gradient_attribute(param)
+        if grad_accum is None:
+            src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
+        else:
+            src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
+        if not self.fp16_master_weights_and_gradients:
+            src_tensor = src_tensor.float()
+
+        dest_tensor.copy_(src_tensor, non_blocking=True)
+        param.grad = None  # offload only
+
+    def complete_grad_norm_calculation_for_cpu_offload(self, params):
+        total_norm = 0.0
+        norm_type = 2.0
+        for p in params:
+            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+            if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
+                continue
+
+            if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
+                param_id = self.get_param_id(p)
+                # as some model have trainable parameters but skipped in training,
+                # their backward hooks in self.create_reduce_and_remove_grad_hooks() will not run,
+                # so they have no norm_for_param_grads
+                if param_id in self.norm_for_param_grads:
+                    param_norm = self.norm_for_param_grads[param_id]
+                    total_norm += param_norm.item() ** 2
+                else:
+                    # As unused parameters in modules may not be expected sometimes,
+                    # add an explicit error msg when it occurred and an option to
+                    # avoid the error
+                    assert self.ignore_unused_parameters, """
+                        This assert indicates that your module has parameters that
+                        were not used in producing loss.
+                        You can avoid this assert by
+                        (1) enable ignore_unused_parameters option in zero_optimization config;
+                        (2) making sure all trainable parameters and `forward` function
+                            outputs participate in calculating loss.
+                    """
+
+        # Sum across all model parallel GPUs.
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group)
+
+        self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM)
+
+        total_norm = total_norm_cuda[0].item() ** (1. / norm_type)
+
+        if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
+            total_norm = -1
+
+        return total_norm
+
+    ############################################################################################
+    def copy_grads_in_partition(self, param):
+        if self.cpu_offload:
+
+            if self.gradient_accumulation_steps > 1:
+                self.async_accumulate_grad_in_cpu_via_gpu(param)
+
+            if self.is_gradient_accumulation_boundary:
+                self.set_norm_for_param_grad_in_gpu(param)
+
+                self.update_overflow_tracker_for_param_grad(param)
+
+                self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param)
+
+            return
+        # print(f"ID {self.get_param_id(param)} grad norm {param.grad.norm()}")
+        if self.grads_in_partition is None:
+            self.grads_in_partition_offset = 0
+            total_size = 0
+            for group in self.params_in_partition:
+                for param_in_partition in group:
+                    total_size += param_in_partition.numel()
+
+            see_memory_usage(f"before copying {total_size} gradients into partition")
+            self.grads_in_partition = torch.empty(int(total_size),
+                                                  dtype=self.dtype,
+                                                  device=get_accelerator().current_device_name())
+            see_memory_usage(f"after copying {total_size} gradients into partition")
+
+        grad_reduc = self.get_gradient_for_reduction(param)
+        # The allreduce buffer will be rewritten. Copy the gradients in partition to a new buffer
+        new_grad_tensor = self.grads_in_partition.view(-1).narrow(0, self.grads_in_partition_offset, param.numel())
+        new_grad_tensor.copy_(grad_reduc.view(-1))
+        grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc)
+        # print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}")
+        self.grads_in_partition_offset += param.numel()
+
+    def reduce_ipg_grads(self):
+        if self.contiguous_gradients:
+            if self.extra_large_param_to_reduce is not None:
+                assert len(self.params_in_ipg_bucket) == 1, "more than 1 param in ipg bucket, this shouldn't happen"
+                _, _, param_id = self.params_in_ipg_bucket[0]
+                assert self.get_param_id(self.extra_large_param_to_reduce
+                                         ) == param_id, "param in ipg bucket does not match extra-large param"
+                extra_large_grad_reduc = self.get_gradient_for_reduction(self.extra_large_param_to_reduce)
+                self.average_tensor(extra_large_grad_reduc.view(-1))
+                self.extra_large_param_to_reduce = None
+            else:
+                self.average_tensor(self.ipg_buffer[self.ipg_index])
+        else:
+            self.buffered_reduce_fallback(None,
+                                          self.grads_in_ipg_bucket,
+                                          elements_per_buffer=self.elements_in_ipg_bucket)
+
+        if self.overlap_comm:
+            stream = self.reduction_stream
+        elif self.cpu_offload:
+            # TODO: copy_grad_stream is disabled because of race with reduce. This hurts perf and should be fixed.
+            #            get_accelerator().synchronize()
+            #            stream = self.copy_grad_stream
+            stream = get_accelerator().current_stream()
+        else:
+            stream = get_accelerator().current_stream()
+
+        with get_accelerator().stream(stream):
+            for _, param, param_id in self.params_in_ipg_bucket:
+
+                assert self.params_already_reduced[param_id] == False, \
+                    f"The parameter {param_id} has already been reduced. \
+                    Gradient computed twice for this partition. \
+                    Multiple gradient reduction is currently not supported"
+
+                self.params_already_reduced[param_id] = True
+                if self.partition_gradients:
+                    if not self.is_param_in_current_partition[param_id]:
+                        if self.overlap_comm and self.contiguous_gradients is False:
+                            # Clear grads of other partitions during the next reduction
+                            # to avoid clearing them before the reduction is complete.
+                            if self.previous_reduced_grads is None:
+                                self.previous_reduced_grads = []
+                            self.previous_reduced_grads.append(param)
+                        else:
+                            self.clear_grad_attribute(param)
+                    elif self.contiguous_gradients:
+                        self.copy_grads_in_partition(param)
+                else:  # zero stage 1 - partition only optimizer state
+                    if self.contiguous_gradients and self.is_param_in_current_partition[param_id]:
+                        self.copy_grads_in_partition(param)
+
+        self.grads_in_ipg_bucket = []
+        self.params_in_ipg_bucket = []
+        self.ipg_bucket_has_moe_params = False
+        self.elements_in_ipg_bucket = 0
+        #####################################################################
+
+    def reduce_ready_partitions_and_remove_grads(self, param, i):
+        if self.partition_gradients or self.is_gradient_accumulation_boundary:
+            self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
+
+    def zero_reduced_gradients(self, partition_id, i):
+
+        def are_all_related_partitions_reduced(params_id):
+            for partition_id in self.param_to_partition_ids[i][params_id]:
+                if not self.is_partition_reduced[i][partition_id]:
+                    return False
+            return True
+
+        for params_id in self.is_grad_computed[i][partition_id]:
+            if are_all_related_partitions_reduced(params_id):
+                self.param_dict[params_id].grad = None  # dead code
+
+    def flatten_and_print(self, message, tensors, start=0, n=5):
+        flatten_tensor = self.flatten(tensors)
+
+        def print_func():
+            logger.info(flatten_tensor.contiguous().view(-1).narrow(0, start, n))
+
+        self.sequential_execution(print_func, message)
+
+    def get_grads_to_reduce(self, i, partition_id):
+
+        def get_reducible_portion(key):
+            grad = self.param_dict[key].grad
+            total_elements = grad.numel()
+            start = self.grad_start_offset[i][partition_id][key]
+            num_elements = min(total_elements - start,
+                               self.partition_size[i] - self.grad_partition_insertion_offset[i][partition_id][key])
+            if not pg_correctness_test:
+                if num_elements == total_elements:
+                    return grad
+                else:
+                    return grad.contiguous().view(-1).narrow(0, int(start), int(num_elements))
+            else:
+                if num_elements == total_elements:
+                    return grad.clone()
+                else:
+                    return grad.clone().contiguous().view(-1).narrow(0, int(start), int(num_elements))
+
+        grads_to_reduce = []
+        for key in self.is_grad_computed[i][partition_id]:
+            grad = get_reducible_portion(key)
+            grads_to_reduce.append(grad)
+        return grads_to_reduce
+
+    def sequential_execution(self, function, message, group=None):
+        if group is None:
+            group = self.zp_process_group
+        if dist.get_rank(group=group) == 0:
+            logger.info(message)
+        for id in range(dist.get_world_size(group=group)):
+            if id == dist.get_rank(group=group):
+                function()
+            dist.barrier(group=group)
+
+    def set_none_gradients_to_zero(self, i, partition_id):
+        for param_id in self.is_grad_computed[i][partition_id]:
+            param = self.param_dict[param_id]
+            if param.grad is None:
+                param.grad = torch.zero_like(param)
+
+    ######################Reduction Related Methods##############################
+    def allreduce_bucket(self, bucket, rank=None, log=None, divide=True, process_group=None):
+        rank = None
+        tensor = self.flatten(bucket)
+
+        process_group = self.zp_process_group if process_group is None else process_group
+
+        tensor_to_allreduce = tensor
+
+        if pg_correctness_test or self.sequence_parallel_size > 1:
+            communication_data_type = torch.float32
+        else:
+            communication_data_type = self.communication_data_type
+
+        if communication_data_type != tensor.dtype:
+            tensor_to_allreduce = tensor.to(communication_data_type)
+
+        if divide:
+            tensor_to_allreduce.div_(
+                dist.get_world_size(group=self.dp_process_group) / float(self.sequence_parallel_size))
+
+        tensor_to_allreduce = tensor_to_allreduce.contiguous()
+        if rank is None:
+            #    "All Reducing"
+            dist.all_reduce(tensor_to_allreduce, group=self.dp_process_group)
+        else:
+            global_rank = dist.get_global_rank(self.dp_process_group, rank)
+            dist.reduce(tensor_to_allreduce, global_rank, group=self.dp_process_group)
+
+        if communication_data_type != tensor.dtype and tensor is not tensor_to_allreduce:
+            if rank is None or rank == dist.get_rank(group=process_group):
+                tensor.copy_(tensor_to_allreduce)
+
+        return tensor
+
+    def _clear_previous_reduced_grads(self):
+        if self.previous_reduced_grads is not None:
+            for param in self.previous_reduced_grads:
+                self.clear_grad_attribute(param)
+            self.previous_reduced_grads = None
+
+    # if rank is specified do a reduction instead of an allreduce
+    def allreduce_and_copy(self, small_bucket, rank=None, log=None, divide=True, process_group=None):
+        process_group = self.zp_process_group if process_group is None else process_group
+        if self.overlap_comm:
+            get_accelerator().synchronize()
+            # It is safe to clear the previously reduced grads of other partitions
+            self._clear_previous_reduced_grads()
+            stream = self.reduction_stream
+        else:
+            stream = get_accelerator().current_stream()
+
+        with get_accelerator().stream(stream):
+            allreduced = self.allreduce_bucket(
+                small_bucket,
+                rank=rank,
+                log=log,
+                divide=divide,
+                process_group=process_group,
+            )
+            if rank is None or rank == dist.get_rank(group=self.zp_process_group):
+                for buf, synced in zip(small_bucket, self.unflatten(allreduced, small_bucket)):
+                    buf.copy_(synced)
+
+    def allreduce_no_retain(
+            self,
+            bucket,
+            numel_per_bucket=500000000,
+            rank=None,
+            log=None,
+            divide=True,
+            process_group=None,
+    ):
+        small_bucket = []
+        numel = 0
+        for tensor in bucket:
+            small_bucket.append(tensor)
+            numel = numel + tensor.numel()
+            if numel > numel_per_bucket:
+                self.allreduce_and_copy(small_bucket, rank=rank, log=None, divide=divide, process_group=process_group)
+                small_bucket = []
+                numel = 0
+
+        if len(small_bucket) > 0:
+            self.allreduce_and_copy(small_bucket, rank=rank, log=log, divide=divide, process_group=process_group)
+
+    # allows using reduction of gradients instead of using all_reduce
+
+    def buffered_reduce_fallback(self, rank, grads, elements_per_buffer=500000000, log=None):
+        split_buckets = split_half_float_double(grads)
+
+        for i, bucket in enumerate(split_buckets):
+            self.allreduce_no_retain(bucket, numel_per_bucket=elements_per_buffer, rank=rank, log=log)
+
+    #############################################################################
+    #############################################################################
+    #############################################################################
+
+    # views the tensor as multiple partitions and returns
+    # those partitions
+    def get_data_parallel_partitions(self, tensor, group_id):
+        partitions = []
+
+        dp = dist.get_world_size(group=self.real_zp_process_group[group_id])
+        # dp_id = dist.get_rank(group=self.real_dp_process_group[group_id])
+
+        total_num_elements = tensor.numel()
+
+        base_size = total_num_elements // dp
+        remaining = total_num_elements % dp
+
+        start = 0
+        for id in range(dp):
+            partition_size = base_size
+            if id < remaining:
+                partition_size = partition_size + 1
+            partitions.append(tensor.narrow(0, start, partition_size))
+            start = start + partition_size
+        return partitions
+
+    def get_partition_info(self, tensor_list, partition_size, partition_id):
+        params_in_partition = []
+        params_not_in_partition = []
+
+        start_index = partition_size * partition_id
+        end_index = partition_size * (partition_id + 1)
+
+        current_index = 0
+        first_offset = 0
+
+        for tensor in tensor_list:
+
+            tensor_size = tensor.numel()
+
+            if start_index <= current_index < end_index:
+                params_in_partition.append(tensor)
+
+            elif current_index < start_index < (current_index + tensor_size):
+                params_in_partition.append(tensor)
+
+                assert (first_offset == 0
+                        ), "This can happen either zero or only once as this must be the first tensor in the partition"
+                first_offset = start_index - current_index
+
+            else:
+                params_not_in_partition.append(tensor)
+
+            current_index = current_index + tensor_size
+
+        return params_in_partition, params_not_in_partition, first_offset
+
+    def zero_grad(self, set_to_none=True):
+        """
+        Zero FP16 parameter grads.
+        """
+        # FP32 grad should never exist.
+        # For speed, set model fp16 grad to None by default
+        # zero all pointers to grad tensors
+        for group in self.bit16_groups:
+            for p in group:
+                if set_to_none:
+                    p.grad = None  # epilogue and in step
+                    p.grad_accum = None
+                else:
+                    if p.grad is not None:
+                        p.grad.detach_()
+                        p.grad.zero_()
+
+    def _model_parallel_all_reduce(self, tensor, op):
+        """ Perform all reduce within model parallel group, if any.
+        """
+        if self.model_parallel_group is None or self.model_parallel_world_size == 1:
+            pass
+        else:
+            dist.all_reduce(tensor=tensor, op=op, group=self.model_parallel_group)
+
+    def get_grad_norm_direct(self, gradients, params, norm_type=2):
+        """Clips gradient norm of an iterable of parameters.
+
+        This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
+        added functionality to handle model parallel parameters. Note that
+        the gradients are modified in place.
+
+        Arguments:
+            parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+                single Tensor that will have gradients normalized
+            max_norm (float or int): max norm of the gradients
+            norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+                infinity norm.
+
+        Returns:
+            Total norm of the parameters (viewed as a single vector).
+        """
+        norm_type = float(norm_type)
+        if norm_type == inf:
+            total_norm = max(g.data.abs().max() for g in gradients)
+            total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=self.dp_process_group)
+
+            # Take max across all GPUs.
+            self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.MAX)
+            total_norm = total_norm_cuda[0].item()
+        else:
+            total_norm = 0.0
+            # if dist.get_rank() == 0:
+            #    logger.info(f"Total Norm beginning {total_norm}")
+            for g, p in zip(gradients, params):
+                # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+                if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
+                    continue
+                if is_model_parallel_parameter(p) or (self.model_parallel_rank == 0):
+                    param_norm = g.data.double().norm(2)
+                    total_norm += param_norm.item() ** 2
+            # Sum across all model parallel GPUs.
+            total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=self.dp_process_group)
+
+            self._model_parallel_all_reduce(tensor=total_norm_cuda, op=dist.ReduceOp.SUM)
+
+            total_norm = total_norm_cuda[0].item() ** (1. / norm_type)
+
+        if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
+            total_norm = -1
+
+        return total_norm
+
+    # creates a flat fused tensor from the tensor list starting at the first_offset
+    # in the first tensor of the list. If there are not enough elements in the tensor
+    # list then the flat tensor will be padded with zeros
+    def get_flat_partition(self, tensor_list, first_offset, partition_size, dtype, device, return_tensor_list=False):
+        flat_tensor_list = []
+        current_size = 0
+
+        for i, tensor in enumerate(tensor_list):
+            grad_accum = self.get_param_gradient_attribute(tensor)
+            if grad_accum is None:
+                grad_accum = torch.zeros_like(tensor, dtype=dtype)
+
+            tensor = grad_accum
+            num_elements = tensor.numel()
+            tensor_offset = 0
+
+            # we need to offset to get to the right element
+            if i == 0 and first_offset > 0:
+                tensor_offset = first_offset
+                num_elements = num_elements - tensor_offset
+
+            # we dont need all elements of the tensor
+            if num_elements > (partition_size - current_size):
+                num_elements = partition_size - current_size
+
+            # we need a narrow view of the tensor based on the tensor offset and number of elements that
+            # we need from this tensor
+            if tensor_offset > 0 or num_elements < tensor.numel():
+                flat_tensor_list.append(tensor.contiguous().view(-1).narrow(0, int(tensor_offset), int(num_elements)))
+            else:
+                flat_tensor_list.append(tensor)
+
+            current_size = current_size + num_elements
+
+        # this means its the last partition and does not align with the dp boundary. We need to pad before flattening
+        if current_size < partition_size:
+            flat_tensor_list.append(torch.zeros(int(partition_size - current_size), dtype=dtype, device=device))
+
+        if return_tensor_list:
+            return flat_tensor_list
+
+        return self.flatten(flat_tensor_list)
+
+    def free_grad_in_param_list(self, param_list):
+        for p in param_list:
+            p.grad = None  # in step
+            p.grad_accum = None
+
+    def reset_cpu_buffers(self):
+        self.norm_for_param_grads = {}
+        self.local_overflow = False
+
+    def set_lr(self, lr):
+        """Set the learning rate."""
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = lr
+
+    def get_lr(self):
+        """Return the current learning rate."""
+        return self.optimizer.param_groups[0]["lr"]
+
+    def override_loss_scale(self, loss_scale):
+        if loss_scale != self.external_loss_scale:
+            logger.info(f'[deepspeed] setting loss scale from {self.external_loss_scale} -> {loss_scale}')
+        self.custom_loss_scaler = True
+        self.external_loss_scale = loss_scale
+
+    def scaled_global_norm(self, norm_type=2):
+        assert norm_type == 2, "only L2 norm supported"
+        norm_groups = []
+        for i, group in enumerate(self.bit16_groups):
+            partition_id = dist.get_rank(group=self.real_zp_process_group[i])
+            if self.cpu_offload:
+                norm_groups.append(self.complete_grad_norm_calculation_for_cpu_offload(self.params_in_partition[i]))
+                single_grad_partition = self.single_partition_of_fp32_groups[i].grad
+            else:
+                norm_groups.append(self.get_grad_norm_direct(self.averaged_gradients[i], self.params_in_partition[i]))
+
+        if self.has_moe_layers:
+            self._average_expert_grad_norms(norm_groups)
+
+        # note that the get_global_norm function only supports l2 norm
+        return get_global_norm(norm_list=norm_groups)
+
+    def get_bit16_param_group(self, group_no):
+        bit16_partitions = self.parallel_partitioned_bit16_groups[group_no]
+        partition_id = dist.get_rank(group=self.real_zp_process_group[group_no])
+        return [bit16_partitions[dist.get_rank(group=self.real_zp_process_group[group_no])]]
+
+    def _optimizer_step(self, group_no):
+        original_param_groups = self.optimizer.param_groups
+        self.optimizer.param_groups = [original_param_groups[group_no]]
+        # Disabling this as the C++ side copy & synchronize is not working correctly
+        # from deepspeed.ops.adam import DeepSpeedCPUAdam
+        # if type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half:
+        #    self.optimizer.step(fp16_param_groups=[self.get_bit16_param_group(group_no)])
+        # else:
+        #    self.optimizer.step()
+        self.optimizer.step()
+        self.optimizer.param_groups = original_param_groups
+
+    def step(self, closure=None):
+        """
+        Not supporting closure.
+        """
+        self.micro_step_id = -1
+
+        see_memory_usage(f"In step before checking overflow")
+
+        # First compute norm for all group so we know if there is overflow
+        if self.dtype == torch.float16:
+            self.check_overflow()
+
+        prev_scale = self.loss_scale
+        self._update_scale(self.overflow)
+        if self.overflow:
+            see_memory_usage('After overflow before clearing gradients')
+            self.zero_grad(set_to_none=True)
+            if self.cpu_offload:
+                self.reset_cpu_buffers()
+            else:
+                self.averaged_gradients = {}
+
+            see_memory_usage('After overflow after clearing gradients')
+
+            for timer in OPTIMIZER_TIMERS:
+                self.timers(timer).start()
+                self.timers(timer).stop()
+            return
+
+        # Step 1:- Calculate gradient norm using bit-16 grads
+        see_memory_usage('Before norm calculation')
+        scaled_global_grad_norm = self.scaled_global_norm()
+        self._global_grad_norm = scaled_global_grad_norm / prev_scale
+        see_memory_usage('After norm before optimizer')
+
+        # Step 2:- run optimizer and upscaling simultaneously
+        for i, group in enumerate(self.bit16_groups):
+            self.timers(OPTIMIZER_GRADIENTS_TIMER).start()
+            partition_id = dist.get_rank(group=self.real_zp_process_group[i])
+            if self.cpu_offload:
+                single_grad_partition = self.single_partition_of_fp32_groups[i].grad
+                self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
+
+                self.timers(OPTIMIZER_GRADIENTS_TIMER).stop()
+                self.timers(OPTIMIZER_STEP_TIMER).start()
+                self._optimizer_step(i)
+
+                # Disabled, this is not currently working
+                # from deepspeed.ops.adam import DeepSpeedCPUAdam
+                # if not (type(self.optimizer) == DeepSpeedCPUAdam and self.dtype == torch.half):
+                #    bit16_partitions = self.parallel_partitioned_bit16_groups[i]
+                #    fp32_partition = self.single_partition_of_fp32_groups[i]
+                #    bit16_partitions[partition_id].data.copy_(fp32_partition.data)
+                bit16_partitions = self.parallel_partitioned_bit16_groups[i]
+                fp32_partition = self.single_partition_of_fp32_groups[i]
+                bit16_partitions[partition_id].data.copy_(fp32_partition.data)
+
+                self.timers(OPTIMIZER_STEP_TIMER).stop()
+            else:
+                # free gradients for all the parameters that are not updated by this process(ZeRO stage2)
+                self.free_grad_in_param_list(self.params_not_in_partition[i])
+
+                # create a flat gradients for parameters updated by this process
+                # If we are last partition, ensure we have same size grads and partition size, if not pad with zero tensors
+                if partition_id == dist.get_world_size(group=self.real_zp_process_group[i]) - 1:
+                    single_grad_partition = self.flatten_dense_tensors_aligned(
+                        self.averaged_gradients[i],
+                        int(self.partition_size[i])).to(self.single_partition_of_fp32_groups[i].dtype)
+                else:
+                    single_grad_partition = self.flatten(self.averaged_gradients[i]).to(
+                        self.single_partition_of_fp32_groups[i].dtype)
+                assert single_grad_partition.numel() == self.partition_size[i], \
+                    "averaged gradients have different number of elements that partition size {} {} {} {}".format(
+                        single_grad_partition.numel(), self.partition_size[i], i, partition_id)
+
+                self.single_partition_of_fp32_groups[i].grad = single_grad_partition
+                # release all the gradient since we have already created a necessary copy in dp_grad_partition(ZeRO stage2)
+                self.free_grad_in_param_list(self.params_in_partition[i])
+
+                self.averaged_gradients[i] = None
+
+                self.unscale_and_clip_grads([single_grad_partition], scaled_global_grad_norm)
+
+                self.timers(OPTIMIZER_GRADIENTS_TIMER).stop()
+
+                # Step 3:- run the optimizer if no offloading
+                self.timers(OPTIMIZER_STEP_TIMER).start()
+                self._optimizer_step(i)
+                # Step 4:- get rid of the fp32 gradients. Not needed anymore
+                self.single_partition_of_fp32_groups[i].grad = None
+                del single_grad_partition
+                bit16_partitions = self.parallel_partitioned_bit16_groups[i]
+                fp32_partition = self.single_partition_of_fp32_groups[i]
+                bit16_partitions[partition_id].data.copy_(fp32_partition.data)
+                self.timers(OPTIMIZER_STEP_TIMER).stop()
+
+        see_memory_usage('After optimizer before all-gather')
+        if self.cpu_offload:
+            self.reset_cpu_buffers()
+
+        self.timers(OPTIMIZER_ALLGATHER_TIMER).start()
+
+        # if dist.get_rank(group=self.dp_process_group) == 0:
+        #     pdb.set_trace()  # 或者使用其他调试工具
+
+        # Gather the updated weights from everyone.
+        # Then all partitions of the model parameters are updated and ready for next round forward.
+        all_gather_into_tensor_dp_groups(groups_flat=self.bit16_groups_flat,
+                                         partitioned_param_groups=self.parallel_partitioned_bit16_groups,
+                                         zp_process_group=self.real_zp_process_group)
+        self.timers(OPTIMIZER_ALLGATHER_TIMER).stop()
+
+        # TODO: we probably don't need this? just to be safe
+        for i in range(len(self.bit16_groups)):
+            self._update_model_bit16_weights(i)
+
+        self.timers.log(OPTIMIZER_TIMERS)
+        see_memory_usage('After zero_optimizer step')
+
+        return
+
+    @torch.no_grad()
+    def update_lp_params(self):
+        for i, (bit16_partitions, fp32_partition) in enumerate(
+                zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)):
+            partition_id = dist.get_rank(group=self.real_zp_process_group[i])
+            bit16_partitions[partition_id].data.copy_(fp32_partition.data)
+            # print_rank_0(f'update_lp_params {i=} {partition_id=}', force=True)
+            # if i == 0:
+            #     print_rank_0(f'{fp32_partition[:10]=}', force=True)
+        all_gather_into_tensor_dp_groups(groups_flat=self.bit16_groups_flat,
+                                         partitioned_param_groups=self.parallel_partitioned_bit16_groups,
+                                         zp_process_group=self.real_zp_process_group)
+
+    def _average_expert_grad_norms(self, norm_groups):
+        for i, norm in enumerate(norm_groups):
+            if self.is_moe_param_group[i]:
+                scaled_norm = norm * 1.0 / float(dist.get_world_size(group=self.dp_process_group))
+                scaled_norm_tensor = torch.tensor(scaled_norm,
+                                                  device=get_accelerator().device_name(),
+                                                  dtype=torch.float)
+                dist.all_reduce(scaled_norm_tensor, group=self.dp_process_group)
+                norm_groups[i] = scaled_norm_tensor.item()
+
+    def unscale_and_clip_grads(self, grad_groups_flat, total_norm):
+        # compute combined scale factor for this group
+        combined_scale = self.loss_scale
+        if self.clip_grad > 0.:
+            # norm is in fact norm*scale
+            clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad
+            if clip > 1:
+                combined_scale = clip * self.loss_scale
+
+        for grad in grad_groups_flat:
+            if isinstance(grad, list):
+                sub_partitions = grad
+                for g in sub_partitions:
+                    g.data.mul_(1. / combined_scale)
+            else:
+                grad.data.mul_(1. / combined_scale)
+
+    def _check_overflow(self, partition_gradients=True):
+        self.overflow = self.has_overflow(partition_gradients)
+
+    # `params` is a list / generator of torch.Variable
+    def has_overflow_serial(self, params, is_grad_list=False):
+        for p in params:
+            if p.grad is not None and self._has_inf_or_nan(p.grad.data):
+                return True
+
+        return False
+
+    def has_overflow_partitioned_grads_serial(self):
+        for i in range(len(self.bit16_groups)):
+            for j, grad in enumerate(self.averaged_gradients[i]):
+                if grad is not None and self._has_inf_or_nan(grad.data, j):
+                    return True
+        return False
+
+    def has_overflow(self, partition_gradients=True):
+        if partition_gradients:
+            overflow = self.local_overflow if self.cpu_offload else self.has_overflow_partitioned_grads_serial()
+            overflow_gpu = get_accelerator().ByteTensor([overflow])
+            '''This will capture overflow across all data parallel and expert parallel process
+            Since expert parallel process are a subset of data parallel process'''
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.dp_process_group)
+
+        else:
+            params = []
+            for group in self.bit16_groups:
+                for param in group:
+                    params.append(param)
+
+            overflow = self.has_overflow_serial(params, is_grad_list=partition_gradients)
+            overflow_gpu = get_accelerator().ByteTensor([overflow])
+
+        # Since each model parallel GPU carries only part of the model,
+        # make sure overflow flag is synced across all the model parallel GPUs
+        self._model_parallel_all_reduce(tensor=overflow_gpu, op=dist.ReduceOp.MAX)
+
+        overflow = overflow_gpu[0].item()
+        return bool(overflow)
+
+    # `x` is a torch.Tensor
+    @staticmethod
+    def _has_inf_or_nan(x, j=None):
+        try:
+            # if x is half, the .float() incurs an additional deep copy, but it's necessary if
+            # Pytorch's .sum() creates a one-element tensor of the same type as x
+            # (which is true for some recent version of pytorch).
+            cpu_sum = float(x.float().sum())
+            # More efficient version that can be used if .sum() returns a Python scalar
+            # cpu_sum = float(x.sum())
+        except RuntimeError as instance:
+            # We want to check if inst is actually an overflow exception.
+            # RuntimeError could come from a different error.
+            # If so, we still want the exception to propagate.
+            if "value cannot be converted" not in instance.args[0]:
+                raise
+            return True
+        else:
+            if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+                return True
+            return False
+
+    def backward(self, loss, retain_graph=False):
+        """
+        :attr:`backward` performs the following steps:
+
+        1. fp32_loss = loss.float()
+        2. scaled_loss = fp32_loss*loss_scale
+        3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves
+        """
+        self.micro_step_id += 1
+
+        if self.contiguous_gradients:
+            self.ipg_buffer = []
+            buf_0 = torch.empty(int(self.reduce_bucket_size),
+                                dtype=self.dtype,
+                                device=get_accelerator().current_device_name())
+            self.ipg_buffer.append(buf_0)
+
+            # Use double buffers to avoid data access conflict when overlap_comm is enabled.
+            if self.overlap_comm:
+                buf_1 = torch.empty(int(self.reduce_bucket_size),
+                                    dtype=self.dtype,
+                                    device=get_accelerator().current_device_name())
+                self.ipg_buffer.append(buf_1)
+            self.ipg_index = 0
+
+        if self.custom_loss_scaler:
+            scaled_loss = self.external_loss_scale * loss
+            scaled_loss.backward()
+        else:
+            self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
+
+        # Only for Stage 1, Mode 2
+        if self.use_grad_accum_attribute:
+            self.fill_grad_accum_attribute()
+
+    def check_overflow(self, partition_gradients=True):
+        self._check_overflow(partition_gradients)
+
+    def _update_scale(self, has_overflow=False):
+        self.loss_scaler.update_scale(has_overflow)
+
+    # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state"
+    def _get_state(self):
+        return self.optimizer.state
+
+    def _set_state(self, value):
+        self.optimizer.state = value
+
+    state = property(_get_state, _set_state)
+
+    # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups"
+    # (for example, to adjust the learning rate)
+    def _get_param_groups(self):
+        return self.optimizer.param_groups
+
+    def _set_param_groups(self, value):
+        self.optimizer.param_groups = value
+
+    param_groups = property(_get_param_groups, _set_param_groups)
+
+    # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale"
+    def _get_loss_scale(self):
+        if self.custom_loss_scaler:
+            return self.external_loss_scale
+        else:
+            return self.loss_scaler.cur_scale
+
+    def _set_loss_scale(self, value):
+        self.loss_scaler.cur_scale = value
+
+    loss_scale = property(_get_loss_scale, _set_loss_scale)
+    cur_scale = property(_get_loss_scale, _set_loss_scale)
+
+    # Return group tensor after removing paddings that are added for alignment to DP world size.
+    # This method works on the assumption that each group contains a single flattened tensor.
+    def _get_groups_without_padding(self, groups_with_padding):
+        groups_without_padding = []
+        for i, group in enumerate(groups_with_padding):
+            lean_length = group.numel() - self.groups_padding[i]
+            groups_without_padding.append(group[:lean_length])
+
+        return groups_without_padding
+
+    # Return optimizer state after removing paddings that are added for alignment.
+    def _get_state_without_padding(self, state_with_padding, padding):
+        lean_state = {}
+        for key, value in state_with_padding.items():
+            if torch.is_tensor(value):
+                lean_length = value.numel() - padding
+                lean_state[key] = value[:lean_length]
+            else:
+                lean_state[key] = value
+
+        return lean_state
+
+    # Return base optimizer states.
+    # This method assumes that each param group contains a single flattened tensor.
+    def _get_base_optimizer_state(self):
+        optimizer_groups_state = []
+        for i, group in enumerate(self.optimizer.param_groups):
+            p = group['params'][0]
+            lean_optimizer_state = self._get_state_without_padding(self.optimizer.state[p], self.groups_padding[i])
+            optimizer_groups_state.append(lean_optimizer_state)
+
+        return optimizer_groups_state
+
+    def state_dict(self):
+        """
+        Returns a dict containing the current state of this :class:`FP16_Optimizer` instance.
+        This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict
+        of the contained Pytorch optimizer.
+        Example::
+            checkpoint = {}
+            checkpoint['model'] = model.state_dict()
+            checkpoint['optimizer'] = optimizer.state_dict()
+            torch.save(checkpoint, "saved.pth")
+        """
+        state_dict = {}
+        state_dict[LOSS_SCALER] = self.loss_scaler
+        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
+        state_dict['overflow'] = self.overflow
+        state_dict[CLIP_GRAD] = self.clip_grad
+
+        if self.elastic_checkpoint:
+            state_dict[BASE_OPTIMIZER_STATE] = self._get_base_optimizer_state()
+
+            if "step" in self.optimizer.param_groups[0]:
+                # Assuming "step" is the only item that changes through training iterations
+                assert all(group["step"] == self.optimizer.param_groups[0]["step"]
+                           for group in self.optimizer.param_groups), "All param groups must have the same step value"
+                state_dict[BASE_OPTIMIZER_STATE_STEP] = self.optimizer.param_groups[0]["step"]
+        else:
+            state_dict[BASE_OPTIMIZER_STATE] = self.optimizer.state_dict()
+
+        # Remove paddings for DP alignment to enable loading for other alignment values
+        fp32_groups_without_padding = self._get_groups_without_padding(self.single_partition_of_fp32_groups)
+        state_dict[SINGLE_PARTITION_OF_FP32_GROUPS] = fp32_groups_without_padding
+
+        state_dict[
+            ZERO_STAGE] = ZeroStageEnum.gradients if self.partition_gradients else ZeroStageEnum.optimizer_states
+        state_dict[GROUP_PADDINGS] = self.groups_padding
+        state_dict[PARTITION_COUNT] = self.partition_count
+
+        state_dict[DS_VERSION] = version
+        state_dict[PARAM_SLICE_MAPPINGS] = self._param_slice_mappings
+
+        return state_dict
+
+    # Restore base optimizer fp32 weights from elastic checkpoint by:
+    # 1) Merging fp32 weights from checkpoints of all partitions
+    # 2) Extracting fp32 weights for current partition from merged weights
+    # 3) Using extracted weights to update base optimizer weights directly.
+    def _restore_from_elastic_fp32_weights(self, all_state_dict):
+        merged_single_partition_of_fp32_groups = []
+
+        for i in range(len(self.single_partition_of_fp32_groups)):
+            partition_id = dist.get_rank(group=self.real_zp_process_group[i])
+            merged_partitions = [sd[SINGLE_PARTITION_OF_FP32_GROUPS][i] for sd in all_state_dict]
+            if self.is_moe_group(self.optimizer.param_groups[i]):
+                ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name'])
+                merged_partitions = [merged_partitions[i] for i in ranks]
+            flat_merged_partitions = self.flatten_dense_tensors_aligned(
+                merged_partitions,
+                self.nccl_start_alignment_factor * dist.get_world_size(group=self.real_zp_process_group[i]))
+            dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, i)
+            merged_single_partition_of_fp32_groups.append(dp_partitions[partition_id])
+
+        for current, saved in zip(self.single_partition_of_fp32_groups, merged_single_partition_of_fp32_groups):
+            current.data.copy_(saved.data)
+
+    # Restore base optimizer fp32 weights from ZeRO fp16 or bfloat16 weights
+    def _restore_from_bit16_weights(self):
+        for group_id, (bit16_partitions, fp32_partition) in enumerate(
+                zip(self.parallel_partitioned_bit16_groups, self.single_partition_of_fp32_groups)):
+            partition_id = dist.get_rank(group=self.real_zp_process_group[group_id])
+            fp32_partition.data.copy_(bit16_partitions[partition_id].data)
+
+    # Refresh the fp32 master params from the fp16 or bfloat16 copies.
+    def refresh_fp32_params(self):
+        self._restore_from_bit16_weights()
+
+    # Extract optimizer state for current partition from merged states of all partitions
+    def _partition_base_optimizer_state(self, state_key, all_partition_states, group_id):
+        partition_id = dist.get_rank(group=self.real_zp_process_group[group_id])
+        alignment = dist.get_world_size(group=self.real_zp_process_group[group_id])
+        if torch.is_tensor(all_partition_states[0]):
+            flat_merged_partitions = self.flatten_dense_tensors_aligned(all_partition_states, alignment)
+            dp_partitions = self.get_data_parallel_partitions(flat_merged_partitions, group_id)
+            return dp_partitions[partition_id]
+        else:
+            # Assume non-tensor states are not partitioned and equal across ranks, so return first one
+            return all_partition_states[0]
+
+    def _restore_base_optimizer_state(self, base_optimizer_group_states):
+        if type(base_optimizer_group_states) == dict:
+            base_optimizer_group_states = base_optimizer_group_states['state']
+        for i, group in enumerate(self.optimizer.param_groups):
+            p = group['params'][0]
+            for key, saved in base_optimizer_group_states[i].items():
+                if torch.is_tensor(self.optimizer.state[p][key]):
+                    dst_tensor = self.optimizer.state[p][key]
+                    src_tensor = _get_padded_tensor(saved, dst_tensor.numel())
+                    self.optimizer.state[p][key].data.copy_(src_tensor.data)
+                else:
+                    self.optimizer.state[p][key] = saved
+
+    def get_ep_ranks(self, rank=0, group_name=None):
+        from deepspeed.utils import groups
+        expert_parallel_size_ = groups._get_expert_parallel_world_size(group_name)
+        world_size = groups._get_data_parallel_world_size()
+        rank = groups._get_expert_parallel_rank(group_name)
+        ranks = range(rank, world_size, expert_parallel_size_)
+        return list(ranks)
+
+    # Restore base optimizer state from elastic checkpoint by
+    # 1) Merging optimizer state from checkpoints of all partitions
+    # 2) Extracting optimizer state for current partition from the merged state
+    # 3) Using the extracted value to directly update the base optimizer.
+    def _restore_elastic_base_optimizer_state(self, all_state_dict):
+        base_optimizer_group_states = []
+        for i in range(len(self.optimizer.param_groups)):
+            partition_states = {}
+            all_partition_group_states = [sd[BASE_OPTIMIZER_STATE][i] for sd in all_state_dict]
+
+            if self.is_moe_group(self.optimizer.param_groups[i]):
+                ranks = self.get_ep_ranks(group_name=self.optimizer.param_groups[i]['name'])
+                all_partition_group_states = [all_partition_group_states[i] for i in ranks]
+
+            for key in all_partition_group_states[0].keys():
+                all_partition_states = [all_states[key] for all_states in all_partition_group_states]
+                partition_states[key] = self._partition_base_optimizer_state(key, all_partition_states, i)
+            base_optimizer_group_states.append(partition_states)
+
+        self._restore_base_optimizer_state(base_optimizer_group_states)
+
+        # Restore step
+        if BASE_OPTIMIZER_STATE_STEP in all_state_dict[0]:
+            assert all(sd[BASE_OPTIMIZER_STATE_STEP] == all_state_dict[0][BASE_OPTIMIZER_STATE_STEP]
+                       for sd in all_state_dict), "State dicts of all partitions must have the same step value"
+            loaded_param_groups_step = all_state_dict[0][BASE_OPTIMIZER_STATE_STEP]
+            for param_group in self.optimizer.param_groups:
+                param_group['step'] = loaded_param_groups_step
+
+    def load_state_dict(self,
+                        state_dict_list,
+                        load_optimizer_states=True,
+                        load_from_fp32_weights=False,
+                        checkpoint_folder=None,
+                        load_serial=None):
+        if checkpoint_folder:
+            self._load_universal_checkpoint(checkpoint_folder, load_optimizer_states, load_from_fp32_weights)
+        else:
+            self._load_legacy_checkpoint(state_dict_list, load_optimizer_states, load_from_fp32_weights)
+
+    def _load_universal_checkpoint(self, checkpoint_folder, load_optimizer_states, load_from_fp32_weights):
+        self._load_hp_checkpoint_state(checkpoint_folder)
+
+    @property
+    def param_groups(self):
+        """Forward the wrapped optimizer's parameters."""
+        return self.optimizer.param_groups
+
+    def _load_hp_checkpoint_state(self, checkpoint_dir):
+        checkpoint_dir = os.path.join(checkpoint_dir, "zero")
+        optim_state_path = os.path.join(checkpoint_dir, "optimizer_state.pt")
+        assert os.path.isfile(
+            optim_state_path), f'{optim_state_path} containing optimizer global state is missing! Cannot proceed.'
+        optim_sd = torch.load(optim_state_path)
+        self._load_global_state(optim_sd)
+
+        tp_rank = bwc_tensor_model_parallel_rank(mpu=self.mpu)
+        tp_world_size = self.mpu.get_slice_parallel_world_size() if hasattr(self.mpu, "get_slice_parallel_world_size") \
+            else self.mpu.get_tensor_model_parallel_world_size()
+
+        for i, _ in enumerate(self.optimizer.param_groups):
+            for lp in self.bit16_groups[i]:
+                if lp._hp_mapping is not None:
+                    # print(f"Loading {self.param_names[lp]} {tp_rank=} {tp_world_size=}")
+                    lp.load_hp_checkpoint_state(os.path.join(checkpoint_dir, self.param_names[lp]), tp_rank,
+                                                tp_world_size)
+
+    def _load_global_state(self, sd):
+        self.loss_scaler = sd.get(LOSS_SCALER, self.loss_scaler)
+        self.dynamic_loss_scale = sd.get('dynamic_loss_scale', self.dynamic_loss_scale)
+        self.overflow = sd.get('overflow', self.overflow)
+        self.clip_grad = sd.get(CLIP_GRAD, self.clip_grad)
+
+        ckpt_version = sd.get(DS_VERSION, False)
+        assert ckpt_version, f"Empty ds_version in checkpoint, not clear how to proceed"
+        ckpt_version = pkg_version.parse(ckpt_version)
+
+        # zero stage 1 mode
+        if not self.partition_gradients:
+            required_version = pkg_version.parse("0.3.17")
+            error_str = f"ZeRO stage 1 changed in {required_version} and is not backwards compatible " \
+                        "with older stage 1 checkpoints. If you'd like to load an old ZeRO-1 checkpoint " \
+                        "please use an older version of DeepSpeed (<= 0.5.8) and set 'legacy_stage1': true in your zero config json."
+            assert required_version <= ckpt_version, f"Old version: {ckpt_version} {error_str}"
+
+    def _load_legacy_checkpoint(self, state_dict_list, load_optimizer_states=True, load_from_fp32_weights=False):
+        r"""Loading ZeRO checkpoint
+
+        Arguments:
+            state_dict_list: List of all saved ZeRO checkpoints, one for each saved partition.
+                Note that the number of saved partitions may differ from number of loading partitions to support
+                changing GPU count, specifically DP world size, between saving and loading checkpoints.
+            load_optimizer_states: Boolean indicating whether or not to load base optimizer states
+            load_from_fp32_weights: Boolean indicating whether to initialize fp32 master weights from fp32
+            copies in checkpoints (no precision loss) or from model's fp16 copies (with precision loss).
+        """
+        """
+        Loads a state_dict created by an earlier call to state_dict().
+        If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``,
+        whose parameters in turn came from ``model``, it is expected that the user
+        will call ``model.load_state_dict()`` before
+        ``fp16_optimizer_instance.load_state_dict()`` is called.
+        Example::
+            model = torch.nn.Linear(D_in, D_out).to(get_accelerator().device_name()).half()
+            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
+            ...
+            checkpoint = torch.load("saved.pth")
+            model.load_state_dict(checkpoint['model'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+        """
+
+        # I think it should actually be ok to reload the optimizer before the model.
+        dp_rank = dist.get_rank(group=self.zp_process_group)
+        current_rank_sd = state_dict_list[dp_rank]
+        self._load_global_state(current_rank_sd)
+
+        ckpt_is_rigid = isinstance(current_rank_sd[BASE_OPTIMIZER_STATE], dict)
+
+        # padding is always at the last rank/partition
+        # if DP=1024 and param-group elems=16 -> padding will be 1024-16 across all but one rank
+        # scenario-1 (shrink): saving w. 4 gpus -> loading w. 2 gpus
+        # scenario-2 (expand): saving w. 2 gpus -> loading w. 4 gpus
+        # if load_optimizer_states:
+        #     if new_dp_size:
+        #         self.strip_padding()
+        #         self.add_padding_w_new_dp_size()
+        #     self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE])
+
+        if load_optimizer_states:
+            if ckpt_is_rigid:
+                # loading rigid ckpt into either rigid or elastic exec
+                self.optimizer.load_state_dict(current_rank_sd[BASE_OPTIMIZER_STATE])
+            else:
+                if self.elastic_checkpoint:
+                    # loading elastic into elastic exec
+                    self._restore_elastic_base_optimizer_state(state_dict_list)
+                else:
+                    # loading an elastic checkpoint into rigid exec
+                    self._restore_base_optimizer_state(current_rank_sd[BASE_OPTIMIZER_STATE])
+
+        # At this point, the optimizer's references to the model's fp32 parameters are up to date.
+        # The optimizer's hyperparameters and internal buffers are also up to date.
+        # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
+        # out of date.  There are two options.
+        # 1:  Refresh the master params from the model's fp16 params.
+        # This requires less storage but incurs precision loss.
+        # 2:  Save and restore the fp32 master copies separately.
+        # We choose option 1 if changing DP degree and option 2 otherwise.
+        #
+        # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device
+        # of their associated parameters, because it's possible those buffers might not exist yet in
+        # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been
+        # constructed in the same way as the one whose state_dict we are loading, the same master params
+        # are guaranteed to exist, so we can just copy_() from the saved master params.
+
+        if load_from_fp32_weights:
+            # option 2 from above
+            if self.elastic_checkpoint and not ckpt_is_rigid:
+                self._restore_from_elastic_fp32_weights(state_dict_list)
+            else:
+                # For non-elastic checkpoint, simply copying from saved weights of current rank is sufficient.
+                for current, saved in zip(self.single_partition_of_fp32_groups,
+                                          current_rank_sd[SINGLE_PARTITION_OF_FP32_GROUPS]):
+                    src_tensor = _get_padded_tensor(saved, current.numel())
+                    current.data.copy_(src_tensor.data)
+        else:
+            # option 1 from above
+            self._restore_from_bit16_weights()
+
+        if load_optimizer_states:
+            self._link_all_hp_params()
+
+
+def _handle_overflow(cpu_sum, x, i):
+    import math
+    rank = dist.get_rank()
+    if rank == 0:
+        t_i = -1
+        for v_i, v in enumerate(x.data.contiguous().view(-1)):
+            if not math.isfinite(float(v)):
+                t_i = v_i
+                break
+        logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}")
+
+
+def estimate_zero2_model_states_mem_needs(total_params,
+                                          num_gpus_per_node=1,
+                                          num_nodes=1,
+                                          cpu_offload=True,
+                                          additional_buffer_factor=1.5):
+    total_gpus = num_nodes * num_gpus_per_node
+
+    if cpu_offload:
+        gpu_mem = 2 * total_params
+        cpu_mem = total_params * max(4 * total_gpus, 16) * additional_buffer_factor
+    else:
+        gpu_mem = 4 * total_params + int(16 * total_params / total_gpus)
+        cpu_mem = total_params * 4 * num_gpus_per_node * additional_buffer_factor
+
+    return int(cpu_mem), int(gpu_mem)
+
+
+def model_to_params(model):
+    # shared params calculated only once
+    total_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
+    return total_params
+
+
+def estimate_zero2_model_states_mem_needs_all_live(model,
+                                                   num_gpus_per_node=1,
+                                                   num_nodes=1,
+                                                   additional_buffer_factor=1.5):
+    """
+    Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients
+    for a given ``model`` and hardware setup.
+
+    If you have an actual model object, use this function and everything will be derived
+    automatically.
+
+    If it's a hypothetical model, use ``estimate_zero2_model_states_mem_needs_all_cold`` where you have to pass
+    the ``total_params`` explicitly.
+
+    Args:
+        - ``model``: ``nn.Module`` object
+        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
+        - ``num_nodes``: how many nodes (defaults to 1),
+        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
+
+    """
+
+    total_params = model_to_params(model)
+
+    estimate_zero2_model_states_mem_needs_all_cold(total_params=total_params,
+                                                   num_gpus_per_node=num_gpus_per_node,
+                                                   num_nodes=num_nodes,
+                                                   additional_buffer_factor=additional_buffer_factor)
+
+
+def estimate_zero2_model_states_mem_needs_all_cold(total_params,
+                                                   num_gpus_per_node=1,
+                                                   num_nodes=1,
+                                                   additional_buffer_factor=1.5):
+    """
+    Print out estimates on memory usage requirements for ZeRO 2 params, optim states and gradients
+    for a given ``model`` and hardware setup.
+
+    If it's a hypothetical model, use this function where you have to pass
+    the ``total_params`` and ``largest_layer_params`` explicitly.
+
+    If you have an actual model object, use ``estimate_zero2_model_states_mem_needs_all_live`` and everything
+    will be derived automatically.
+
+    Args:
+        - ``total_params``: total  model params
+        - ``num_gpus_per_node``: how many gpus per node (defaults to 1)
+        - ``num_nodes``: how many nodes (defaults to 1),
+        - ``additional_buffer_factor``: estimation factor (defaults to 1.5):
+
+    """
+
+    def format_options(cpu_offload):
+        enabled = []
+        device = f'{OffloadDeviceEnum.cpu:4}' if cpu_offload else "none"
+        enabled.append(f"offload_optimizer={device}")
+        return ", ".join(enabled)
+
+    nodes_str = "nodes" if num_nodes > 1 else "node"
+    gpus_str = "GPUs" if num_gpus_per_node > 1 else "GPU"
+    print("Estimated memory needed for params, optim states and gradients for a:\n"
+          f"HW: Setup with {num_nodes} {nodes_str}, {num_gpus_per_node} {gpus_str} per node.\n"
+          f"SW: Model with {int(total_params / 1e6)}M total params.")
+    print("  per CPU  |  per GPU |   Options")
+    for cpu_offload in [True, False]:
+        cpu_mem, gpu_mem = estimate_zero2_model_states_mem_needs(total_params=total_params,
+                                                                 num_gpus_per_node=num_gpus_per_node,
+                                                                 num_nodes=num_nodes,
+                                                                 cpu_offload=cpu_offload,
+                                                                 additional_buffer_factor=additional_buffer_factor)
+
+        options_str = format_options(cpu_offload=cpu_offload)
+        print(f" {cpu_mem / 2 ** 30:7.2f}GB | {gpu_mem / 2 ** 30:6.2f}GB | {options_str}")
diff --git a/opensora/adaptor/utils.py b/opensora/adaptor/utils.py
new file mode 100644
index 000000000..c32ff6efc
--- /dev/null
+++ b/opensora/adaptor/utils.py
@@ -0,0 +1,1042 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""
+Copyright NVIDIA/Megatron
+
+Helper functions and classes from multiple sources.
+"""
+
+from collections.abc import Iterable
+from deepspeed.moe.utils import is_moe_param
+import os
+import psutil
+import gc
+from math import sqrt
+from packaging import version as pkg_version
+
+import torch
+from deepspeed import comm as dist
+
+try:
+    from torch._six import inf
+except ModuleNotFoundError:
+    from torch import inf
+
+from deepspeed.utils import groups, logger
+from deepspeed.runtime.constants import PIPE_REPLICATED
+from numpy import prod
+from deepspeed.accelerator import get_accelerator
+
+from deepspeed.module_inject.policy import transpose
+from torch.nn import functional as F
+
+torch_memory_reserved = get_accelerator().memory_reserved
+torch_max_memory_reserved = get_accelerator().max_memory_reserved
+
+
+class DummyOptim():
+    """
+    Dummy optimizer presents model parameters as a param group, this is
+    primarily used to allow ZeRO-3 without an optimizer
+    """
+
+    def __init__(self, params):
+        self.param_groups = []
+        self.param_groups.append({'params': params})
+
+
+graph_cache = {}
+
+
+def graph_process(replay_first_step, func, *args, **kwargs):
+    # `func` should only contain operations on the GPU
+    # Please ensure that the memory address of the data required by 'func' remains constant
+    if func.__name__ not in graph_cache:
+        cuda_stream = get_accelerator().Stream()
+        cuda_stream.wait_stream(get_accelerator().current_stream())
+        with get_accelerator().stream(cuda_stream):
+            func(*args, **kwargs)
+        get_accelerator().current_stream().wait_stream(cuda_stream)
+        graph_cache[func.__name__] = get_accelerator().create_graph()
+        with get_accelerator().capture_to_graph(graph_cache[func.__name__]):
+            func(*args, **kwargs)
+        if replay_first_step:
+            get_accelerator().replay_graph(graph_cache[func.__name__])
+    else:
+        get_accelerator().replay_graph(graph_cache[func.__name__])
+
+
+def noop_decorator(func):
+    return func
+
+
+class noop_context(object):
+
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+def ensure_directory_exists(filename):
+    """Create the directory path to ``filename`` if it does not already exist.
+
+    Args:
+        filename (str): A file path.
+    """
+    dirname = os.path.dirname(filename)
+    os.makedirs(dirname, exist_ok=True)
+
+
+def set_random_seed(seed):
+    """Set the random seed for common PRNGs used during training: random, numpy, and torch.
+
+    Args:
+        seed (int): the seed to use
+    """
+    import numpy
+    import random
+    random.seed(seed)
+    numpy.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def is_model_parallel_parameter(p) -> bool:
+    if hasattr(p, 'model_parallel') and p.model_parallel:
+        return True
+
+    if hasattr(p, 'tensor_model_parallel') and p.tensor_model_parallel:
+        return True
+
+    return False
+
+
+def bwc_tensor_model_parallel_rank(mpu=None):
+    """Backwards-compatible way of querying the tensor model parallel rank from
+    an ``mpu`` object.
+
+    *Tensor* model parallelism means that tensors are physically split across
+    processes. This contrasts with *pipeline* model parallelism, in which the
+    layers are partitioned but tensors left intact.
+
+    The API for tensor model parallelism has changed across versions and this
+    helper provides a best-effort implementation across versions of ``mpu``
+    objects.  The preferred mechanism is
+    ``mpu.get_tensor_model_parallel_rank()``.
+
+    This should "just work" with both Megatron-LM and DeepSpeed's pipeline
+    parallelism.
+
+    Args:
+        mpu (model parallel unit, optional): The tensor model parallel rank.
+            If ``mpu=None``, returns 0. Defaults to ``None``.
+
+    Returns:
+        int: the rank
+    """
+    if mpu is None:
+        # No model parallelism in easy :)
+        return 0
+
+    if hasattr(mpu, 'get_tensor_model_parallel_rank'):
+        # New Megatron and DeepSpeed convention (post pipeline-parallelism release)
+        return mpu.get_tensor_model_parallel_rank()
+    elif hasattr(mpu, 'get_slice_parallel_rank'):
+        # Some DeepSpeed + pipeline parallelism versions
+        return mpu.get_slice_parallel_rank()
+    else:
+        # Deprecated Megatron and DeepSpeed convention
+        return mpu.get_model_parallel_rank()
+
+
+def copy_to_device(item, device, criterion_func):
+    """
+    Return a copy of tensor on specified device.
+    Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.
+    Parameters:
+        item: tensor to copy or (possibly nested) container of tensors to copy.
+        device: target device
+        criterion_func: Function to restrict copy operation to items meet criterion
+
+    Returns:
+        None
+    """
+    if criterion_func(item):
+        return item.to(device)
+    elif isinstance(item, list):
+        return [copy_to_device(v, device, criterion_func) for v in item]
+    elif isinstance(item, tuple):
+        return tuple([copy_to_device(v, device, criterion_func) for v in item])
+    elif isinstance(item, dict):
+        return {k: copy_to_device(v, device, criterion_func) for k, v in item.items()}
+    else:
+        return item
+
+
+def move_to_device(item, device, criterion_func):
+    """
+    Move tensor on to specified device by changing the storage.
+    Works on individual tensors, and tensors contained/nested in lists, tuples, and dicts.
+    Parameters:
+        item: tensor to move or (possibly nested) container of tensors to move.
+        device: target device
+        criterion_func: Function to restrict move operation to items meet criterion
+
+    Returns:
+        None
+    """
+    if criterion_func(item):
+        device_copy = item.to(device)
+        item.data = device_copy.data
+        return item
+    elif isinstance(item, list):
+        return [move_to_device(v, device, criterion_func) for v in item]
+    elif isinstance(item, tuple):
+        return tuple([move_to_device(v, device, criterion_func) for v in item])
+    elif isinstance(item, dict):
+        return {k: move_to_device(v, device, criterion_func) for k, v in item.items()}
+    else:
+        return item
+
+
+class CheckOverflow(object):
+    '''Checks for overflow in gradient across parallel process'''
+
+    def __init__(self, param_groups=None, mpu=None, zero_reduce_scatter=False, deepspeed=None):
+        self.mpu = mpu
+        self.params = [] if param_groups else None
+        self.zero_reduce_scatter = zero_reduce_scatter
+        self.deepspeed = deepspeed
+        self.has_moe_params = False
+        if param_groups:
+            for group in param_groups:
+                for param in group:
+                    self.params.append(param)
+                    if is_moe_param(param):
+                        self.has_moe_params = True
+
+    def check_using_norm(self, norm_group, reduce_overflow=True):
+        # TODO: I don't think reduce_overflow is needed if mpu is None
+        overflow = -1 in norm_group
+        overflow_gpu = get_accelerator().FloatTensor([overflow])
+        if self.has_moe_params:
+            # In this case, we need to do an all_reduce across
+            # the expert_parallel_group, so that if there was
+            # an overflow due to expert weights, we detect it
+
+            # Only need to check groups.get_largest_expert_parallel_group()
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group())
+        if self.mpu is not None:
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group())
+        elif reduce_overflow:
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX)
+            dist.barrier()
+        overflow = overflow_gpu[0].item()
+        return bool(overflow)
+
+    def check(self, param_groups=None):
+        params = []
+        has_moe_params = False
+        if param_groups is None:
+            params = self.params
+            has_moe_params = self.has_moe_params
+        else:
+            assert param_groups is not None, \
+                "self.params and param_groups both cannot be none"
+
+            for group in param_groups:
+                for param in group:
+                    params.append(param)
+                    if is_moe_param(param):
+                        has_moe_params = True
+
+        return self.has_overflow(params, has_moe_params=has_moe_params)
+
+    # `params` is a list / generator of torch.Variable
+    def has_overflow_serial(self, params):
+        for i, p in enumerate(params):
+            if p.grad is not None and self._has_inf_or_nan(p.grad.data, i):
+                return True
+        return False
+
+    def has_overflow(self, params, has_moe_params=None):
+        if has_moe_params is None:
+            has_moe_params = self.has_moe_params
+        overflow = self.has_overflow_serial(params)
+        # Since each model parallel GPU carries only part of the model,
+        # make sure overflow flag is synced across all the model parallel GPUs
+        overflow_gpu = get_accelerator().ByteTensor([overflow])
+        # deepspeed.comm.all_reduce(overflow_gpu,
+        #                             op=deepspeed.comm.ReduceOp.MAX,
+        #                             group=mpu.get_model_parallel_group())
+        if has_moe_params:
+            # All reduce this across expert_parallel_group, so that if an expert
+            # overflows, we detect it here
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group())
+        if self.zero_reduce_scatter:
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group())
+        elif self.mpu is not None:
+            if self.deepspeed is not None:
+                using_pipeline = hasattr(self.deepspeed, 'pipeline_enable_backward_allreduce')
+                if (using_pipeline and self.deepspeed.pipeline_enable_backward_allreduce is False) or (
+                        not using_pipeline and self.deepspeed.enable_backward_allreduce is False):
+                    dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_data_parallel_group())
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group())
+        elif self.deepspeed is not None and self.deepspeed.enable_backward_allreduce is False:
+            dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=dist.get_world_group())
+
+        overflow = overflow_gpu[0].item()
+        return bool(overflow)
+
+    # `x` is a torch.Tensor
+    @staticmethod
+    def _has_inf_or_nan(x, i):
+        try:
+            # if x is half, the .float() incurs an additional deep copy, but it's necessary if
+            # Pytorch's .sum() creates a one-element tensor of the same type as x
+            # (which is true for some recent version of pytorch).
+            cpu_sum = float(x.float().sum())
+            # More efficient version that can be used if .sum() returns a Python scalar
+            # cpu_sum = float(x.sum())
+        except RuntimeError as instance:
+            # We want to check if inst is actually an overflow exception.
+            # RuntimeError could come from a different error.
+            # If so, we still want the exception to propagate.
+            if "value cannot be converted" not in instance.args[0]:
+                raise
+            return True
+        else:
+            if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+                return True
+            return False
+
+
+def _handle_overflow(cpu_sum, x, i):
+    import math
+    rank = dist.get_rank()
+    if rank == 0:
+        t_i = -1
+        for v_i, v in enumerate(x.data.contiguous().view(-1)):
+            if not math.isfinite(float(v)):
+                t_i = v_i
+                break
+        logger.info(f"rank {rank} detected overflow {cpu_sum} in tensor {i}:{t_i} shape {x.shape}")
+
+
+def get_global_norm(norm_list):
+    """ Compute total from a list of norms
+    """
+    total_norm = 0.0
+    for norm in norm_list:
+        total_norm += norm**2.0
+    # logger.info(f'norm_list = {norm_list} global = {sqrt(total_norm)}')
+    return sqrt(total_norm)
+
+
+def clip_grad_norm_(parameters, max_norm, norm_type=2, mpu=None):
+    """Clips gradient norm of an iterable of parameters.
+
+    This has been adapted from Nvidia megatron. We add norm averaging
+    to consider MoE params when calculating norm as they will result
+    in different norms across different ranks.
+
+    This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
+    added functionality to handle model parallel parameters. Note that
+    the gradients are modified in place.
+
+    Arguments:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        max_norm (float or int): max norm of the gradients
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+
+    Returns:
+        Total norm of the parameters (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    max_norm = float(max_norm)
+    norm_type = float(norm_type)
+    if norm_type == inf:
+        total_norm = max(p.grad.data.abs().max() for p in parameters)
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        # Take max across all GPUs.
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()
+    else:
+        total_norm = 0
+        for p in parameters:
+            if mpu is not None:
+                if (mpu.get_model_parallel_rank() == 0) or is_model_parallel_parameter(p):
+                    param_norm = p.grad.data.norm(norm_type)
+                    total_norm += param_norm.item()**norm_type
+            else:
+                param_norm = p.grad.data.float().norm(norm_type)
+                total_norm += param_norm.item()**norm_type
+
+        # Sum across all model parallel GPUs.
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()**(1. / norm_type)
+
+    # Need to average total_norm across different GPUs due to the presence of moe params
+    pg = groups._get_data_parallel_group()
+    scaled_norm = total_norm * 1.0 / float(dist.get_world_size(group=pg))
+
+    scaled_norm_tensor = get_accelerator().FloatTensor([float(scaled_norm)])
+    dist.all_reduce(scaled_norm_tensor, group=pg)
+    total_norm = scaled_norm_tensor.item()
+
+    clip_coef = max_norm / (total_norm + 1e-6)
+    if clip_coef < 1:
+        for p in parameters:
+            p.grad.data.mul_(clip_coef)
+    return total_norm
+
+
+def get_grad_norm(parameters, norm_type=2, mpu=None):
+    """Get grad norm of an iterable of parameters.
+
+    This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
+    added functionality to handle model parallel parameters. Note that
+    the gradients are modified in place. Taken from Nvidia Megatron.
+
+    Arguments:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+
+    Returns:
+        Total norm of the parameters (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+
+    norm_type = float(norm_type)
+    if norm_type == inf:
+        total_norm = max(p.grad.data.abs().max() for p in parameters)
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        # Take max across all GPUs.
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()
+    else:
+        total_norm = 0.
+        tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
+        for p in parameters:
+            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+            if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
+                continue
+
+            # Filter to avoid over-counting replicated tensors from tensor
+            # model parallelism
+            if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p):
+                continue
+
+            param_norm = p.grad.data.float().norm(norm_type)
+            total_norm += param_norm.item()**norm_type
+
+        # Sum across all model parallel GPUs.
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()**(1. / norm_type)
+
+    if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
+        total_norm = -1
+
+    return total_norm
+
+
+def get_grad_zeros(parameters, mpu=None):
+    """Compute the number of grads with zero values.
+
+    This is adapted from get_grad_norm
+
+    Arguments:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+
+    Returns:
+        Total number of params with zero values (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+
+    total_zeros = 0.
+    tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
+    for p in parameters:
+        # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+        if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
+            continue
+
+        # Filter to avoid over-counting replicated tensors from tensor
+        # model parallelism
+        if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p):
+            continue
+
+        count_zeros = p.grad.numel() - torch.count_nonzero(p.grad)
+        total_zeros += count_zeros.item()
+
+    # Sum across all model parallel GPUs.
+    total_zeros_cuda = get_accelerator().FloatTensor([float(total_zeros)])
+    if mpu is not None:
+        dist.all_reduce(total_zeros_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+    total_zeros = total_zeros_cuda[0].item()
+
+    return total_zeros
+
+
+def get_weight_norm(parameters, norm_type=2, mpu=None):
+    """Get norm of an iterable of parameters.
+
+    This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
+    added functionality to handle model parallel parameters. Note that
+    the gradients are modified in place. Taken from Nvidia Megatron.
+
+    Arguments:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+
+    Returns:
+        Total norm of the parameters (viewed as a single vector).
+        -1 if the norm value is NaN or Inf.
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+
+    norm_type = float(norm_type)
+    if norm_type == inf:
+        total_norm = max(p.data.abs().max() for p in parameters)
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        # Take max across all GPUs.
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()
+    else:
+        total_norm = 0.
+        tensor_mp_rank = bwc_tensor_model_parallel_rank(mpu=mpu)
+        for p in parameters:
+            # Pipeline parallelism may replicate parameters. Avoid multi-counting.
+            if hasattr(p, PIPE_REPLICATED) and p.ds_pipe_replicated:
+                continue
+
+            # Filter to avoid over-counting replicated tensors from tensor
+            # model parallelism
+            if (tensor_mp_rank > 0) and not is_model_parallel_parameter(p):
+                continue
+
+            param_norm = p.data.float().norm(norm_type)
+            total_norm += param_norm**norm_type
+
+        # Sum across all model parallel GPUs.
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()**(1. / norm_type)
+
+    if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
+        total_norm = -1
+
+    return total_norm
+
+
+def prefix_sum_inc(weights):
+    """ Compute an inclusive prefix sum.
+
+    Example:
+        >>> prefix_sum_inc([3,4,5])
+        [3, 7, 12]
+    """
+    weights_ = [w for w in weights]
+    for x in range(1, len(weights_)):
+        weights_[x] += weights_[x - 1]
+    return weights_
+
+
+def partition_uniform(num_items, num_parts):
+    import numpy
+    parts = [0] * (num_parts + 1)
+    # First check for the trivial edge case
+    if num_items <= num_parts:
+        for p in range(num_parts + 1):
+            parts[p] = min(p, num_items)
+        return parts
+
+    chunksize = num_items // num_parts
+    residual = num_items - (chunksize * num_parts)
+
+    parts = numpy.arange(0, (num_parts + 1) * chunksize, chunksize)
+
+    for i in range(residual):
+        parts[i + 1:] += 1
+    parts = parts.tolist()
+
+    return parts
+
+
+def partition_balanced(weights, num_parts):
+    """
+    use dynamic programming solve `The Linear Partition Problem`.
+    see https://www8.cs.umu.se/kurser/TDBAfl/VT06/algorithms/BOOK/BOOK2/NODE45.HTM
+    """
+    import numpy as np
+    n = len(weights)
+    m = num_parts
+
+    if n <= m:
+        return partition_uniform(n, m)
+
+    dp_max = np.full((n + 1, m + 1), np.inf)
+    dp_min = np.full((n + 1, m + 1), np.inf)
+    dp_cost = np.full((n + 1, m + 1), np.inf)
+    position = np.zeros((n + 1, m + 1), dtype=int)
+    prefix_sum = np.zeros((n + 1))
+    prefix_sum[1:] = np.cumsum(weights)
+
+    dp_max[0, 0] = 0
+    dp_cost[0, 0] = 0
+    for i in range(1, n + 1):
+        for j in range(1, min(i, m) + 1):
+            for k in range(i):
+                max_sum = max(dp_max[k, j - 1], prefix_sum[i] - prefix_sum[k])
+                min_sum = min(dp_min[k, j - 1], prefix_sum[i] - prefix_sum[k])
+                cost = max_sum - min_sum
+                if dp_cost[i, j] >= cost:
+                    dp_cost[i, j] = cost
+                    dp_max[i, j] = max_sum
+                    dp_min[i, j] = min_sum
+                    position[i, j] = k
+
+    parts = [n]
+    for i in reversed(range(1, m + 1)):
+        parts.append(position[parts[-1], i])
+    parts.reverse()
+
+    return parts
+
+
+class PartitionedTensor:
+
+    def __init__(self, tensor, group, partition_meta=None):
+        super().__init__()
+
+        self.group = group
+        self.num_parts = dist.get_world_size(group=self.group)
+        self.rank = dist.get_rank(group=self.group)
+        self.orig_size = list(tensor.size())
+        self.orig_device = tensor.device
+        self.local_data, self.partition = self._partition_tensor(tensor)
+        self.even_split = tensor.numel() % self.num_parts == 0
+
+    @classmethod
+    def from_meta(cls, meta, local_part, group, device=get_accelerator().device_name()):
+        assert meta.dtype == torch.long
+        dummy = torch.ones(dist.get_world_size(group=group))
+        part_obj = cls(tensor=dummy, group=group)
+
+        meta = meta.tolist()
+
+        # [N, list0, ..., listN-1]
+        part_obj.orig_size = meta[1:(1 + meta[0])]
+        meta = meta[1 + meta[0]:]
+
+        part_obj.orig_device = device
+        part_obj.local_data = local_part.detach()
+
+        part_obj.group = group
+
+        # Partition is encoded like the rowptr of a CSR matrix:
+        # [num_parts, rank, 0, part_1, ..., part_num_parts]
+        # TODO: support shuffle between different partition granularities
+        assert part_obj.num_parts == meta[0]
+        assert part_obj.rank == meta[1]
+        part_obj.partition = meta[2:]  # length num_parts+1
+
+        return part_obj
+
+    def _partition_tensor(self, tensor):
+        partition = partition_uniform(num_items=tensor.numel(), num_parts=self.num_parts)
+        start = partition[self.rank]
+        length = partition[self.rank + 1] - start
+        tensor_part = tensor.detach().contiguous().view(-1).narrow(0, start=start, length=length).clone()
+
+        return tensor_part, partition
+
+    def full(self, device=None):
+        if device is None:
+            device = self.orig_device
+
+        # Allocate the full tensor as a flat buffer.
+        full_numel = prod(self.full_size())
+        flat_tensor = torch.zeros([full_numel], dtype=self.local_data.dtype, device=device)
+        if self.even_split:
+            # Collect the full tensor
+            dist.all_gather_into_tensor(flat_tensor, self.local_data, group=self.group)
+        else:
+            for part_id in range(self.num_parts):
+                part_size = self.partition[part_id + 1] - self.partition[part_id]
+                buf = flat_tensor.narrow(0, start=self.partition[part_id], length=part_size)
+                if part_id == self.rank:
+                    buf.copy_(self.local_data)
+                dist.broadcast(buf, part_id, self.group)
+        return flat_tensor.view(self.full_size()).clone().detach()
+
+    def to_meta(self):
+        """Returns a torch.LongTensor that encodes partitioning information.
+
+        Can be used along with ``data()`` to serialize a ``PartitionedTensor`` for
+        communication.
+
+        Returns:
+            torch.LongTensor: a tensor encoding the meta-information for the partitioning
+        """
+        meta = []
+        meta.append(len(self.orig_size))
+        meta += list(self.orig_size)
+        meta.append(self.num_parts)
+        meta.append(self.rank)
+        meta += self.partition
+        return torch.LongTensor(data=meta).to(self.orig_device)
+
+    def data(self):
+        return self.local_data
+
+    def local_size(self):
+        return self.local_data.size()
+
+    def full_size(self):
+        return self.orig_size
+
+
+mem_alloced = 0
+mem_cached = 0
+
+
+def memory_status(msg, print_rank=-1, reset_max=False):
+    global mem_alloced, mem_cached
+
+    rank = dist.get_rank()
+    if print_rank != -1 and rank != print_rank:
+        return
+
+    get_accelerator().synchronize()
+
+    if reset_max:
+        get_accelerator().reset_max_memory_cached()
+        get_accelerator().reset_max_memory_allocated()
+
+    new_alloced = get_accelerator().memory_allocated()
+    new_cached = get_accelerator().memory_cached()
+
+    delta_alloced = new_alloced - mem_alloced
+    delta_cached = new_cached - mem_cached
+
+    mem_cached = new_cached
+    mem_alloced = new_alloced
+
+    max_alloced = get_accelerator().max_memory_allocated()
+    max_cached = get_accelerator().max_memory_cached()
+
+    # convert to GB for printing
+    new_alloced /= 1024**3
+    new_cached /= 1024**3
+    delta_alloced /= 1024**3
+    delta_cached /= 1024**3
+    max_alloced /= 1024**3
+    max_cached /= 1024**3
+
+    print(
+        f'RANK={rank} MEMSTATS', msg, f'device={get_accelerator().current_device_name()} '
+        f'current alloc={new_alloced:0.4f}GB (delta={delta_alloced:0.4f}GB max={max_alloced:0.4f}GB) '
+        f'current cache={new_cached:0.4f}GB (delta={delta_cached:0.4f}GB max={max_cached:0.4f}GB)')
+
+
+def get_ma_status():
+    if dist.is_initialized() and not dist.get_rank() == 0:
+        return 0
+    return get_accelerator().memory_allocated()
+
+
+def empty_cache():
+    get_accelerator().empty_cache()
+    get_accelerator().reset_peak_memory_stats()
+
+
+def see_memory_usage(message, force=False):
+    if not force:
+        return
+    if dist.is_initialized() and not dist.get_rank() == 0:
+        return
+
+    # python doesn't do real-time garbage collection so do it explicitly to get the correct RAM reports
+    gc.collect()
+
+    # Print message except when distributed but not rank 0
+    logger.info(message)
+    logger.info(f"MA {round(get_accelerator().memory_allocated() / (1024 * 1024 * 1024),2 )} GB \
+        Max_MA {round(get_accelerator().max_memory_allocated() / (1024 * 1024 * 1024),2)} GB \
+        CA {round(torch_memory_reserved() / (1024 * 1024 * 1024),2)} GB \
+        Max_CA {round(torch_max_memory_reserved() / (1024 * 1024 * 1024))} GB ")
+
+    vm_stats = psutil.virtual_memory()
+    used_GB = round(((vm_stats.total - vm_stats.available) / (1024**3)), 2)
+    logger.info(f'CPU Virtual Memory:  used = {used_GB} GB, percent = {vm_stats.percent}%')
+
+    # get the peak memory to report correct data, so reset the counter for the next call
+    get_accelerator().reset_peak_memory_stats()
+
+
+def call_to_str(base, *args, **kwargs):
+    """Construct a string representation of a call.
+
+    Args:
+        base (str): name of the call
+        args (tuple, optional): args to ``base``
+        kwargs (dict, optional): kwargs supplied to ``base``
+
+    Returns:
+        str: A string representation of base(*args, **kwargs)
+    """
+    name = f'{base}('
+    if args:
+        name += ', '.join(repr(arg) for arg in args)
+        if kwargs:
+            name += ', '
+    if kwargs:
+        name += ', '.join(f'{key}={repr(arg)}' for key, arg in kwargs.items())
+    name += ')'
+    return name
+
+
+def get_only_unique_item(items):
+    item_set = set(items)
+    if len(item_set) != 1:
+        raise RuntimeError(f"expected there to be only one unique element in {items}")
+    unique_item, = item_set
+
+    return unique_item
+
+
+def clip_gradients(parameters, max_norm=1.0, global_grad_norm=None, mpu=None, eps=1e-6):
+    """Clip the gradient of a list of parameters.
+    Args:
+        parameters: List of parameters whose .grad will be clipped.
+        global_grad_norm (float, optional): Precomputed gradient norm. Defaults to None.
+        mpu (optional): model parallelism unit. Defaults to None.
+        eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6
+    Returns:
+        float: the global gradient norm
+    """
+    if global_grad_norm is None:
+        global_grad_norm = get_grad_norm(parameters, mpu=mpu)
+    clip_coef = max_norm / (global_grad_norm + eps)
+    if clip_coef < 1:
+        for p in parameters:
+            p.grad.detach().mul_(clip_coef)
+    return global_grad_norm
+
+
+def get_global_norm_of_tensors(input_tensors, norm_type=2, mpu=None, use_graph=False):
+    """Get norm of an iterable of tensors.
+
+    This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_ and
+    added functionality to handle model parallel parameters. Taken from Nvidia Megatron.
+
+    Arguments:
+        input_tensors (Iterable[Tensor]): an iterable of Tensors will have norm computed
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+
+    Returns:
+        Total norm of the tensors (viewed as a single vector).
+    """
+    assert isinstance(input_tensors, Iterable), f'expected Iterable type not {type(input_tensors)}'
+    assert all([torch.is_tensor(t) for t in input_tensors]), f'expected list of only tensors'
+
+    norm_type = float(norm_type)
+    if norm_type == inf:
+        total_norm = max(t.data.abs().max() for t in input_tensors)
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)])
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.MAX, group=mpu.get_model_parallel_group())
+            total_norm = total_norm_cuda[0].item()
+    else:
+        if use_graph:
+            if 'norm_tensors_compute_buffer' not in graph_cache:
+                graph_cache['norm_tensors_compute_buffer'] = [t.data.float().norm(norm_type) for t in input_tensors]
+            compute_buffer = graph_cache['norm_tensors_compute_buffer']
+
+            def _norm_tensors(tensor_list, _compute_buffer, _norm_type):
+                for i, t in enumerate(tensor_list):
+                    _compute_buffer[i].data.copy_(t.data.float().norm(_norm_type)**_norm_type)
+                    if i != 0:
+                        _compute_buffer[0].data.add_(_compute_buffer[i].data)
+
+            graph_process(False, _norm_tensors, input_tensors, compute_buffer, norm_type)
+
+            total_norm = compute_buffer[0]
+        else:
+            total_norm = sum([t.data.float().norm(norm_type).item()**norm_type for t in input_tensors])
+
+        total_norm_cuda = get_accelerator().FloatTensor([float(total_norm)]).detach()
+        if mpu is not None:
+            dist.all_reduce(total_norm_cuda, op=dist.ReduceOp.SUM, group=mpu.get_model_parallel_group())
+        total_norm = total_norm_cuda[0].item()**(1. / norm_type)
+
+    if total_norm == float('inf') or total_norm == -float('inf') or total_norm != total_norm:
+        total_norm = -1
+
+    return total_norm
+
+
+def clip_tensors_by_global_norm(input_tensors, max_norm=1.0, global_norm=None, mpu=None, eps=1e-6, use_graph=False):
+    """Clip list of tensors by global norm.
+    Args:
+        input_tensors: List of tensors to be clipped
+        global_norm (float, optional): Precomputed norm. Defaults to None.
+        mpu (optional): model parallelism unit. Defaults to None.
+        eps (float, optional): epsilon value added to grad norm. Defaults to 1e-6
+    Returns:
+        float: the global norm
+    """
+    if global_norm is None:
+        global_norm = get_global_norm_of_tensors(input_tensors, mpu=mpu, use_graph=use_graph)
+    clip_coef = max_norm / (global_norm + eps)
+    if clip_coef < 1:
+        if use_graph:
+
+            def clip_tensors(_tensor_list, _clip_coef_tensor):
+                for t in _tensor_list:
+                    t.detach().mul_(_clip_coef_tensor)
+
+            if 'clip_coef_tensor' not in graph_cache:
+                # Alloc memory
+                graph_cache['clip_coef_tensor'] = torch.tensor(clip_coef,
+                                                               dtype=torch.float32).to(get_accelerator().device_name())
+            clip_coef_tensor = graph_cache['clip_coef_tensor']
+            clip_coef_tensor.copy_(torch.tensor(clip_coef, dtype=torch.float32))
+            graph_process(False, clip_tensors, input_tensors, clip_coef_tensor)
+
+        else:
+            for t in input_tensors:
+                t.detach().mul_(clip_coef)
+    return global_norm
+
+
+def align_dense_tensors(tensor_list, alignment):
+    num_elements = sum(t.numel() for t in tensor_list)
+    remaining = num_elements % alignment
+
+    if remaining:
+        elements_to_add = alignment - remaining
+        pad_tensor = torch.zeros(elements_to_add, device=tensor_list[0].device, dtype=tensor_list[0].dtype)
+        padded_tensor_list = tensor_list + [pad_tensor]
+    else:
+        padded_tensor_list = tensor_list
+
+    return padded_tensor_list
+
+
+def all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, zp_process_group, dp_process_group=None):
+
+    for group_id, (group_flat, partitioned_params) in enumerate(zip(groups_flat, partitioned_param_groups)):
+        partition_id = dist.get_rank(group=zp_process_group[group_id])
+        dp_world_size = dist.get_world_size(group=dp_process_group)
+        if dp_world_size == 1:
+            # no groups share optimizer states
+            # pipeline parallel with bf16 will default call this even if dp size = 1.
+            continue
+        # print("call contiguous for all_gather_into_tensor_dp_groups")
+        dist.all_gather_into_tensor(group_flat, partitioned_params[partition_id].contiguous(), dp_process_group)
+
+
+def all_gather_dp_groups(groups_flat, partitioned_param_groups, zp_process_group, start_alignment_factor,
+                         allgather_bucket_size, dp_process_group=None):
+    # if dist.has_all_gather_into_tensor():
+    return all_gather_into_tensor_dp_groups(groups_flat, partitioned_param_groups, zp_process_group, dp_process_group)
+
+    # for group_id, partitioned_params in enumerate(partitioned_param_groups):
+    #     # Sequential AllGather Best of both worlds
+    #     partition_id = dist.get_rank(group=dp_process_group[group_id])
+    #     dp_world_size = dist.get_world_size(group=dp_process_group[group_id])
+    #
+    #     if dp_world_size == 1:
+    #         # no groups share optimizer states
+    #         # pipeline parallel with bf16 will default call this even if dp size = 1.
+    #         continue
+    #     num_shards = max(1, partitioned_params[partition_id].numel() * dp_world_size // allgather_bucket_size)
+    #
+    #     shard_size = partitioned_params[partition_id].numel() // num_shards
+    #
+    #     # Enforce nccl/rccl alignment of start location of each shard
+    #     shard_size = shard_size - (shard_size % start_alignment_factor)
+    #
+    #     num_elements = shard_size
+    #
+    #     assert shard_size * num_shards <= partitioned_params[partition_id].numel()
+    #
+    #     for shard_id in range(num_shards):
+    #
+    #         if shard_id == (num_shards - 1):
+    #             num_elements = partitioned_params[partition_id].numel() - shard_id * shard_size
+    #
+    #         shard_list = []
+    #         for dp_id in range(dp_world_size):
+    #             curr_shard = partitioned_params[dp_id].narrow(0, shard_id * shard_size, num_elements).detach()
+    #             shard_list.append(curr_shard)
+    #         dist.all_gather(shard_list, shard_list[partition_id].contiguous(), dp_process_group[group_id])
+
+
+class TLinear(torch.nn.Linear):
+
+    def __init__(self, orig_layer, name=""):
+        self.name = name
+        super().__init__(orig_layer.weight.shape[1], orig_layer.weight.shape[0], bias=(orig_layer.bias is not None))
+        self.weight.data = transpose(orig_layer.weight.data)
+        self.bias = orig_layer.bias
+        self._fwd_func = self._fwd_bias_add if self.bias is not None else self._fwd
+
+    def _fwd(self, input):
+        return F.linear(input, self.weight)
+
+    def _fwd_bias_add(self, input):
+        return F.linear(input, self.weight, bias=self.bias)
+
+    def forward(self, input):
+        return self._fwd_func(input)
+
+
+def get_inactive_params(param_list):
+    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+    return [param for param in param_list if (hasattr(param, 'ds_id') and \
+                            param.ds_status == ZeroParamStatus.NOT_AVAILABLE)]
+
+
+def required_torch_version(min_version=None, max_version=None):
+    assert min_version or max_version, "Must provide a min_version or max_version argument"
+
+    torch_version = pkg_version.parse(torch.__version__)
+
+    if min_version and pkg_version.parse(str(min_version)) > torch_version:
+        return False
+
+    if max_version and pkg_version.parse(str(max_version)) < torch_version:
+        return False
+
+    return True
diff --git a/opensora/adaptor/zp_manager.py b/opensora/adaptor/zp_manager.py
new file mode 100644
index 000000000..e05bfdbcd
--- /dev/null
+++ b/opensora/adaptor/zp_manager.py
@@ -0,0 +1,31 @@
+import torch
+import os
+import torch.distributed as dist
+
+
+class ZPManager(object):
+    def __init__(self, zp_size=8):
+        self.rank = int(os.getenv('RANK', '0'))
+        self.world_size = int(os.getenv("WORLD_SIZE", '1'))
+        self.zp_size = zp_size
+        self.zp_group = None
+        self.zp_rank = None
+        self.is_initialized = False
+
+    def init_group(self):
+        if self.is_initialized:
+            return
+
+        self.is_initialized = True
+
+        """Initialize the sequence parallel group."""
+        num_zp_groups: int = self.world_size // self.zp_size
+        for i in range(num_zp_groups):
+            ranks = range(i * self.zp_size, (i + 1) * self.zp_size)
+            group = dist.new_group(ranks)
+            if self.rank in ranks:
+                self.zp_group = group
+                self.zp_rank = self.rank % self.zp_size
+
+
+zp_manager = ZPManager()
diff --git a/opensora/dataset/__init__.py b/opensora/dataset/__init__.py
new file mode 100644
index 000000000..b4ba76f52
--- /dev/null
+++ b/opensora/dataset/__init__.py
@@ -0,0 +1,92 @@
+from torchvision.transforms import Compose
+from transformers import AutoTokenizer, AutoImageProcessor
+
+from torchvision import transforms
+from torchvision.transforms import Lambda
+
+from opensora.dataset.t2v_datasets import T2V_dataset
+from opensora.dataset.inpaint_datasets import Inpaint_dataset
+from opensora.dataset.transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVideo, CenterCropResizeVideo, LongSideResizeVideo, SpatialStrideCropVideo
+from opensora.models.causalvideovae import ae_norm, ae_denorm
+
+def getdataset(args):
+    temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
+    norm_fun = ae_norm[args.ae]
+    resize_topcrop = [CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True), ]
+    resize = [CenterCropResizeVideo((args.max_height, args.max_width)), ]
+    transform = transforms.Compose([
+        ToTensorVideo(),
+        *resize, 
+        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+        norm_fun
+    ])
+    transform_topcrop = transforms.Compose([
+        ToTensorVideo(),
+        *resize_topcrop, 
+        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+        norm_fun
+    ])
+    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+    if args.dataset == 't2v':
+        return T2V_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer, 
+                           transform_topcrop=transform_topcrop)
+    elif args.dataset == 'inpaint' or args.dataset == 'i2v':
+        return Inpaint_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer, 
+                            transform_topcrop=transform_topcrop)
+    
+    raise NotImplementedError(args.dataset)
+
+
+if __name__ == "__main__":
+    from accelerate import Accelerator
+    from opensora.dataset.t2v_datasets import dataset_prog
+    import random
+    from tqdm import tqdm
+    args = type('args', (), 
+    {
+        'ae': 'CausalVAEModel_4x8x8', 
+        'dataset': 't2v', 
+        'attention_mode': 'xformers', 
+        'use_rope': True, 
+        'model_max_length': 300, 
+        'max_height': 320,
+        'max_width': 240,
+        'num_frames': 1,
+        'use_image_num': 0, 
+        'compress_kv_factor': 1, 
+        'interpolation_scale_t': 1,
+        'interpolation_scale_h': 1,
+        'interpolation_scale_w': 1,
+        'cache_dir': '../cache_dir', 
+        'image_data': '/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt', 
+        'video_data': '1',
+        'train_fps': 24, 
+        'drop_short_ratio': 1.0, 
+        'use_img_from_vid': False, 
+        'speed_factor': 1.0, 
+        'cfg': 0.1, 
+        'text_encoder_name': 'google/mt5-xxl', 
+        'dataloader_num_workers': 10,
+
+    }
+    )
+    accelerator = Accelerator()
+    dataset = getdataset(args)
+    num = len(dataset_prog.img_cap_list)
+    zero = 0
+    for idx in tqdm(range(num)):
+        image_data = dataset_prog.img_cap_list[idx]
+        caps = [i['cap'] if isinstance(i['cap'], list) else [i['cap']] for i in image_data]
+        try:
+            caps = [[random.choice(i)] for i in caps]
+        except Exception as e:
+            print(e)
+            # import ipdb;ipdb.set_trace()
+            print(image_data)
+            zero += 1
+            continue
+        assert caps[0] is not None and len(caps[0]) > 0
+    print(num, zero)
+    import ipdb;ipdb.set_trace()
+    print('end')
\ No newline at end of file
diff --git a/opensora/dataset/inpaint_datasets.py b/opensora/dataset/inpaint_datasets.py
new file mode 100644
index 000000000..7f22baa23
--- /dev/null
+++ b/opensora/dataset/inpaint_datasets.py
@@ -0,0 +1,116 @@
+
+import os, random
+import numpy as np
+
+import torch
+from accelerate.logging import get_logger
+
+from opensora.utils.utils import text_preprocessing
+
+from opensora.dataset.t2v_datasets import SingletonMeta, DataSetProg
+from opensora.dataset.t2v_datasets import T2V_dataset
+
+logger = get_logger(__name__)
+
+dataset_prog = DataSetProg()
+
+
+class Meta_dataset(T2V_dataset):
+    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
+        super().__init__(args, transform, temporal_sample, tokenizer, transform_topcrop)
+
+        if self.num_frames != 1:
+            # inpaint
+            # The proportion of executing the i2v task.
+            self.i2v_ratio = args.i2v_ratio
+            self.transition_ratio = args.transition_ratio
+            self.v2v_ratio = args.v2v_ratio
+            self.clear_video_ratio = args.clear_video_ratio
+            assert self.i2v_ratio + self.transition_ratio + self.v2v_ratio + self.clear_video_ratio < 1, 'The sum of i2v_ratio, transition_ratio, v2v_ratio and clear video ratio should be less than 1.'
+        
+        self.default_text_ratio = args.default_text_ratio
+        self.default_text = f"The {'video' if self.num_frames != 1 else 'image'} showcases a scene with coherent and clear visuals."
+    
+    def get_mask_masked_video(self, video):
+        # video shape (T, C, H, W)
+        # 1 means masked, 0 means not masked
+        t, c, h, w = video.shape
+        mask = torch.ones_like(video, device=video.device, dtype=video.dtype)
+        
+        rand_num = random.random()
+        # i2v
+        if rand_num < self.i2v_ratio:
+            mask[0] = 0
+        # transition
+        elif rand_num < self.i2v_ratio + self.transition_ratio:
+            mask[0] = 0
+            mask[-1] = 0
+        # video continuation
+        elif rand_num < self.i2v_ratio + self.transition_ratio + self.v2v_ratio:
+            end_idx = random.randint(1, t)
+            mask[:end_idx] = 0
+        # clear video
+        elif rand_num < self.i2v_ratio + self.transition_ratio + self.v2v_ratio + self.clear_video_ratio:
+            mask[:] = 0
+        # random mask
+        else:
+            idx_to_select = random.randint(0, t - 1)
+            selected_indices = random.sample(range(0, t), idx_to_select)
+            mask[selected_indices] = 0
+        masked_video = video * (mask < 0.5)
+
+        # save_video(masked_video.permute(0, 2, 3, 1).cpu().numpy(), 'masked_video.mp4')
+        return dict(mask=mask, masked_video=masked_video)
+
+    def drop(self, text):
+        rand_num = random.random()
+        rand_num_text = random.random()
+
+        if rand_num < self.cfg:
+            text = self.default_text if rand_num_text < self.default_text_ratio else ''
+
+        return dict(text=text)
+
+    
+class Inpaint_dataset(Meta_dataset):
+    def get_video(self, idx):
+        video_path = dataset_prog.cap_list[idx]['path']
+        assert os.path.exists(video_path), f"file {video_path} do not exist!"
+        frame_indice = dataset_prog.cap_list[idx]['sample_frame_index']
+        video = self.decord_read(video_path, predefine_num_frames=len(frame_indice))
+
+        h, w = video.shape[-2:]
+        assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}'
+        t = video.shape[0]
+        video = self.transform(video)  # T C H W -> T C H W
+
+        # inpaint
+        inpaint_cond_data = self.get_mask_masked_video(video)
+        mask, masked_video = inpaint_cond_data['mask'], inpaint_cond_data['masked_video']
+        video = torch.cat([video, masked_video, mask], dim=1) # T 3*C H W
+
+        # video = torch.rand(221, 3, 480, 640)
+
+        video = video.transpose(0, 1)  # T C H W -> C T H W
+        text = dataset_prog.cap_list[idx]['cap']
+        if not isinstance(text, list):
+            text = [text]
+        text = [random.choice(text)]
+
+        text = text_preprocessing(text, support_Chinese=self.support_Chinese)
+
+        text = self.drop(text)['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']
+        cond_mask = text_tokens_and_mask['attention_mask']
+        return dict(pixel_values=video, input_ids=input_ids, cond_mask=cond_mask)
+    
\ No newline at end of file
diff --git a/opensora/dataset/t2v_datasets.py b/opensora/dataset/t2v_datasets.py
new file mode 100644
index 000000000..4ed8444a2
--- /dev/null
+++ b/opensora/dataset/t2v_datasets.py
@@ -0,0 +1,422 @@
+import time
+import traceback
+
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+import glob
+import json
+import os, io, csv, math, random
+import numpy as np
+import torchvision
+from einops import rearrange
+from decord import VideoReader
+from os.path import join as opj
+from collections import Counter
+
+import torch
+import torchvision.transforms as transforms
+from torch.utils.data.dataset import Dataset
+from torch.utils.data import DataLoader, Dataset, get_worker_info
+from tqdm import tqdm
+from PIL import Image
+from accelerate.logging import get_logger
+
+from opensora.utils.dataset_utils import DecordInit
+from opensora.utils.utils import text_preprocessing
+logger = get_logger(__name__)
+
+def filter_json_by_existed_files(directory, data, postfixes=[".mp4", ".jpg"]):
+    # 构建搜索模式，以匹配指定后缀的文件
+    matching_files = []
+    for postfix in postfixes:
+        pattern = os.path.join(directory, '**', f'*{postfix}')
+        matching_files.extend(glob.glob(pattern, recursive=True))
+
+    # 使用文件的绝对路径构建集合
+    mp4_files_set = set(os.path.abspath(path) for path in matching_files)
+
+    # 过滤数据条目，只保留路径在mp4文件集合中的条目
+    filtered_items = [item for item in data if item['path'] in mp4_files_set]
+
+    return filtered_items
+
+
+def random_video_noise(t, c, h, w):
+    vid = torch.rand(t, c, h, w) * 255.0
+    vid = vid.to(torch.uint8)
+    return vid
+
+
+class SingletonMeta(type):
+    """
+    这是一个元类，用于创建单例类。
+    """
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class DataSetProg(metaclass=SingletonMeta):
+    def __init__(self):
+        self.cap_list = []
+        self.elements = []
+        self.num_workers = 1
+        self.n_elements = 0
+        self.worker_elements = dict()
+        self.n_used_elements = dict()
+
+    def set_cap_list(self, num_workers, cap_list, n_elements):
+        self.num_workers = num_workers
+        self.cap_list = cap_list
+        self.n_elements = n_elements
+        self.elements = list(range(n_elements))
+        random.shuffle(self.elements)
+        print(f"n_elements: {len(self.elements)}", flush=True)
+
+        for i in range(self.num_workers):
+            self.n_used_elements[i] = 0
+            per_worker = int(math.ceil(len(self.elements) / float(self.num_workers)))
+            start = i * per_worker
+            end = min(start + per_worker, len(self.elements))
+            self.worker_elements[i] = self.elements[start: end]
+
+    def get_item(self, work_info):
+        if work_info is None:
+            worker_id = 0
+        else:
+            worker_id = work_info.id
+
+        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])]
+        self.n_used_elements[worker_id] += 1
+        return idx
+
+
+dataset_prog = DataSetProg()
+
+def find_closest_y(x, vae_stride_t=4, model_ds_t=4):
+    if x < 29:
+        return -1  
+    for y in range(x, 12, -1):
+        if (y - 1) % vae_stride_t == 0 and ((y - 1) // vae_stride_t + 1) % model_ds_t == 0:
+            # 4, 8: y in [29, 61, 93, 125, 157, 189, 221, 253, 285, 317, 349, 381, 413, 445, 477, 509, ...]
+            # 4, 4: y in [29, 45, 61, 77, 93, 109, 125, 141, 157, 173, 189, 205, 221, 237, 253, 269, 285, 301, 317, 333, 349, 365, 381, 397, 413, 429, 445, 461, 477, 493, 509, ...]
+            return y
+    return -1 
+
+def filter_resolution(h, w, max_h_div_w_ratio=17/16, min_h_div_w_ratio=8 / 16):
+    if h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio:
+        return True
+    return False
+        
+
+
+class T2V_dataset(Dataset):
+    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
+        self.data = args.data
+        self.num_frames = args.num_frames
+        self.train_fps = args.train_fps
+        self.use_image_num = args.use_image_num
+        self.use_img_from_vid = args.use_img_from_vid
+        self.transform = transform
+        self.transform_topcrop = transform_topcrop
+        self.temporal_sample = temporal_sample
+        self.tokenizer = tokenizer
+        self.model_max_length = args.model_max_length
+        self.cfg = args.cfg
+        self.speed_factor = args.speed_factor
+        self.max_height = args.max_height
+        self.max_width = args.max_width
+        self.drop_short_ratio = args.drop_short_ratio
+        assert self.speed_factor >= 1
+        self.v_decoder = DecordInit()
+
+        self.support_Chinese = True
+        if not ('mt5' in args.text_encoder_name):
+            self.support_Chinese = False
+
+        cap_list = self.get_cap_list()
+        
+        assert len(cap_list) > 0
+        cap_list, self.sample_num_frames = self.define_frame_index(cap_list)
+        self.lengths = self.sample_num_frames
+
+        n_elements = len(cap_list)
+        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list, n_elements)
+
+        print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
+
+    def set_checkpoint(self, n_used_elements):
+        for i in range(len(dataset_prog.n_used_elements)):
+            dataset_prog.n_used_elements[i] = n_used_elements
+
+    def __len__(self):
+        return dataset_prog.n_elements
+
+    def __getitem__(self, idx):
+        try:
+            data = self.get_data(idx)
+            return data
+        except Exception as e:
+            logger.info(f'Error with {e}')
+            # 打印异常堆栈
+            if idx in dataset_prog.cap_list:
+                logger.info(f"Caught an exception! {dataset_prog.cap_list[idx]}")
+            # traceback.print_exc()
+            # traceback.print_stack()
+            return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+    def get_data(self, idx):
+        path = dataset_prog.cap_list[idx]['path']
+        if path.endswith('.mp4'):
+            return self.get_video(idx)
+        else:
+            return self.get_image(idx)
+    
+    def get_video(self, idx):
+        # npu_config.print_msg(f"current idx is {idx}")
+        # video = random.choice([random_video_noise(65, 3, 336, 448), random_video_noise(65, 3, 1024, 1024), random_video_noise(65, 3, 360, 480)])
+        # # print('random shape', video.shape)
+        # 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)
+
+        video_path = dataset_prog.cap_list[idx]['path']
+        assert os.path.exists(video_path), f"file {video_path} do not exist!"
+        frame_indice = dataset_prog.cap_list[idx]['sample_frame_index']
+        video = self.decord_read(video_path, predefine_num_frames=len(frame_indice))
+
+        h, w = video.shape[-2:]
+        assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}'
+        t = video.shape[0]
+        video = self.transform(video)  # T C H W -> T C H W
+
+        # video = torch.rand(221, 3, 480, 640)
+
+        video = video.transpose(0, 1)  # T C H W -> C T H W
+        text = dataset_prog.cap_list[idx]['cap']
+        if not isinstance(text, list):
+            text = [text]
+        text = [random.choice(text)]
+
+        text = text_preprocessing(text, support_Chinese=self.support_Chinese) if random.random() > self.cfg else ""
+        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']
+        cond_mask = text_tokens_and_mask['attention_mask']
+        return dict(pixel_values=video, input_ids=input_ids, cond_mask=cond_mask)
+
+    def get_image(self, idx):
+        image_data = dataset_prog.cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
+
+        # import ipdb;ipdb.set_trace()
+        image = Image.open(image_data['path']).convert('RGB')  # [h, w, c]
+        image = torch.from_numpy(np.array(image))  # [h, w, c]
+        image = rearrange(image, 'h w c -> c h w').unsqueeze(0)  #  [1 c h w]
+        # for i in image:
+        #     h, w = i.shape[-2:]
+        #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
+        
+        image = self.transform_topcrop(image) if 'human_images' in image_data['path'] else self.transform(image) #  [1 C H W] -> num_img [1 C H W]
+
+        # image = [torch.rand(1, 3, 480, 640) for i in image_data]
+        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
+
+        caps = image_data['cap'] if isinstance(image_data['cap'], list) else [image_data['cap']]
+        caps = [random.choice(caps)]
+        text = text_preprocessing(caps, support_Chinese=self.support_Chinese)
+        input_ids, cond_mask = [], []
+        text = text if random.random() > self.cfg else ""
+        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']  # 1, l
+        cond_mask = text_tokens_and_mask['attention_mask']  # 1, l
+        return dict(pixel_values=image, input_ids=input_ids, cond_mask=cond_mask)
+
+    def define_frame_index(self, cap_list):
+        
+        new_cap_list = []
+        sample_num_frames = []
+        cnt_too_long = 0
+        cnt_too_short = 0
+        cnt_no_cap = 0
+        cnt_no_resolution = 0
+        cnt_resolution_mismatch = 0
+        cnt_movie = 0
+        cnt_img = 0
+        for i in cap_list:
+            path = i['path']
+            cap = i.get('cap', None)
+            # ======no caption=====
+            if cap is None:
+                cnt_no_cap += 1
+                continue
+            if path.endswith('.mp4'):
+                # ======no fps and duration=====
+                duration = i.get('duration', None)
+                fps = i.get('fps', None)
+                if fps is None or duration is None:
+                    continue
+
+                # ======resolution mismatch=====
+                resolution = i.get('resolution', None)
+                if resolution is None:
+                    cnt_no_resolution += 1
+                    continue
+                else:
+                    if resolution.get('height', None) is None or resolution.get('width', None) is None:
+                        cnt_no_resolution += 1
+                        continue
+                    height, width = i['resolution']['height'], i['resolution']['width']
+                    aspect = self.max_height / self.max_width
+                    hw_aspect_thr = 1.5
+                    is_pick = filter_resolution(height, width, max_h_div_w_ratio=hw_aspect_thr*aspect, 
+                                                min_h_div_w_ratio=1/hw_aspect_thr*aspect)
+                    if not is_pick:
+                        cnt_resolution_mismatch += 1
+                        continue
+
+                    # # ignore image resolution mismatch
+                    # if self.max_height > resolution['height'] or self.max_width > resolution['width']:
+                    #     cnt_resolution_mismatch += 1
+                    #     continue
+
+                # import ipdb;ipdb.set_trace()
+                i['num_frames'] = int(fps * duration)
+                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration. 
+                if i['num_frames'] > 2.0 * (self.num_frames * fps / self.train_fps * self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
+                    cnt_too_long += 1
+                    continue
+                # if i['num_frames'] < 1.0/1 * (self.num_frames * fps / self.train_fps * self.speed_factor):  # too short video is not suitable for this training stage
+                #     cnt_too_short += 1
+                #     continue 
+
+                # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
+                frame_interval = fps / self.train_fps
+                start_frame_idx = 0 
+                frame_indices = np.arange(start_frame_idx, i['num_frames'], frame_interval).astype(int)
+                frame_indices = frame_indices[frame_indices < i['num_frames']]
+
+                # comment out it to enable dynamic frames training
+                if len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio:
+                    cnt_too_short += 1
+                    continue
+
+                #  too long video will be temporal-crop randomly
+                if len(frame_indices) > self.num_frames:
+                    begin_index, end_index = self.temporal_sample(len(frame_indices))
+                    frame_indices = frame_indices[begin_index: end_index]
+                    # frame_indices = frame_indices[:self.num_frames]  # head crop
+                # to find a suitable end_frame_idx, to ensure we do not need pad video
+                end_frame_idx = find_closest_y(len(frame_indices), vae_stride_t=4, model_ds_t=4)
+                if end_frame_idx == -1:  # too short that can not be encoded exactly by videovae
+                    cnt_too_short += 1
+                    continue
+                frame_indices = frame_indices[:end_frame_idx]
+
+                i['sample_frame_index'] = frame_indices.tolist()
+                new_cap_list.append(i)
+                i['sample_num_frames'] = len(i['sample_frame_index'])  # will use in dataloader(group sampler)
+                sample_num_frames.append(i['sample_num_frames'])
+            elif path.endswith('.jpg'):  # image
+                cnt_img += 1
+                new_cap_list.append(i)
+                i['sample_num_frames'] = 1
+                sample_num_frames.append(i['sample_num_frames'])
+            else:
+                raise NameError(f"Unknown file extention {path.split('.')[-1]}, only support .mp4 for video and .jpg for image")
+        # import ipdb;ipdb.set_trace()
+        logger.info(f'no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, '
+                f'no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, '
+                f'Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, '
+                f'before filter: {len(cap_list)}, after filter: {len(new_cap_list)}')
+        return new_cap_list, sample_num_frames
+    
+    def decord_read(self, path, predefine_num_frames):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        fps = decord_vr.get_avg_fps() if decord_vr.get_avg_fps() > 0 else 30.0
+        # import ipdb;ipdb.set_trace()
+        # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
+        frame_interval = 1.0 if abs(fps - self.train_fps) < 0.1 else fps / self.train_fps
+        start_frame_idx = 0  
+        frame_indices = np.arange(start_frame_idx, total_frames, frame_interval).astype(int)
+        frame_indices = frame_indices[frame_indices < total_frames]
+        #import ipdb;ipdb.set_trace()
+        # speed up
+        max_speed_factor = len(frame_indices) / self.num_frames
+        if self.speed_factor > 1 and max_speed_factor > 1:
+            # speed_factor = random.uniform(1.0, min(self.speed_factor, max_speed_factor))
+            speed_factor = min(self.speed_factor, max_speed_factor)
+            target_frame_count = int(len(frame_indices) / speed_factor)
+            speed_frame_idx = np.linspace(0, len(frame_indices) - 1, target_frame_count, dtype=int)
+            frame_indices = frame_indices[speed_frame_idx]
+
+        #  too long video will be temporal-crop randomly
+        if len(frame_indices) > self.num_frames:
+            begin_index, end_index = self.temporal_sample(len(frame_indices))
+            frame_indices = frame_indices[begin_index: end_index]
+            # frame_indices = frame_indices[:self.num_frames]  # head crop
+
+        # to find a suitable end_frame_idx, to ensure we do not need pad video
+        end_frame_idx = find_closest_y(len(frame_indices), vae_stride_t=4, model_ds_t=4)
+        if end_frame_idx == -1:  # too short that can not be encoded exactly by videovae
+            raise IndexError(f'video ({path}) has {total_frames} frames, but need to sample {len(frame_indices)} frames ({frame_indices})')
+        frame_indices = frame_indices[:end_frame_idx]
+        if predefine_num_frames != len(frame_indices):
+            raise ValueError(f'predefine_num_frames ({predefine_num_frames}) is not equal with frame_indices ({len(frame_indices)})')
+        if len(frame_indices) < self.num_frames and self.drop_short_ratio >= 1:
+            raise IndexError(f'video ({path}) has {total_frames} frames, but need to sample {len(frame_indices)} frames ({frame_indices})')
+        video_data = decord_vr.get_batch(frame_indices).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 read_jsons(self, data):
+        cap_lists = []
+        with open(data, 'r') as f:
+            folder_anno = [i.strip().split(',') for i in f.readlines() if len(i.strip()) > 0]
+        for folder, anno in folder_anno:
+            with open(anno, 'r') as f:
+                sub_list = json.load(f)
+            logger.info(f'Building {anno}...')
+            for i in range(len(sub_list)):
+                sub_list[i]['path'] = opj(folder, sub_list[i]['path'])
+            cap_lists += sub_list
+        return cap_lists
+
+    # def get_img_cap_list(self):
+    #     use_image_num = self.use_image_num if self.use_image_num != 0 else 1
+    #     if npu_config is None:
+    #         img_cap_lists = self.read_jsons(self.image_data, postfix=".jpg")
+    #         img_cap_lists = [img_cap_lists[i: i + use_image_num] for i in range(0, len(img_cap_lists), use_image_num)]
+    #     else:
+    #         img_cap_lists = npu_config.try_load_pickle("img_cap_lists_all",
+    #                                                    lambda: self.read_jsons(self.image_data, postfix=".jpg"))
+    #         img_cap_lists = [img_cap_lists[i: i + use_image_num] for i in range(0, len(img_cap_lists), use_image_num)]
+    #         img_cap_lists = img_cap_lists[npu_config.get_local_rank()::npu_config.N_NPU_PER_NODE]
+    #     return img_cap_lists[:-1]  # drop last to avoid error length
+
+    def get_cap_list(self):
+        cap_lists = self.read_jsons(self.data)
+        return cap_lists
diff --git a/opensora/dataset/transform.py b/opensora/dataset/transform.py
new file mode 100644
index 000000000..96a1b696d
--- /dev/null
+++ b/opensora/dataset/transform.py
@@ -0,0 +1,592 @@
+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 center_crop_th_tw(clip, th, tw, top_crop):
+    if not _is_tensor_video_clip(clip):
+        raise ValueError("clip should be a 4D torch.tensor")
+    
+    # import ipdb;ipdb.set_trace()
+    h, w = clip.size(-2), clip.size(-1)
+    tr = th / tw
+    if h / w > tr:
+        new_h = int(w * tr)
+        new_w = w
+    else:
+        new_h = h
+        new_w = int(h / tr)
+    
+    i = 0 if top_crop else int(round((h - new_h) / 2.0))
+    j = int(round((w - new_w) / 2.0))
+    return crop(clip, i, j, new_h, new_w)
+
+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 SpatialStrideCropVideo:
+    def __init__(self, stride):
+            self.stride = stride
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: cropped video clip by stride.
+                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 = h // self.stride * self.stride, w // self.stride * self.stride
+
+        return 0, 0, th, tw  # from top-left
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+class LongSideResizeVideo:
+    '''
+    First use the long side,
+    then resize to the specified size
+    '''
+
+    def __init__(
+            self,
+            size,
+            skip_low_resolution=False, 
+            interpolation_mode="bilinear",
+    ):
+        self.size = size
+        self.skip_low_resolution = skip_low_resolution
+        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 video clip.
+                size is (T, C, 512, *) or (T, C, *, 512)
+        """
+        _, _, h, w = clip.shape
+        if self.skip_low_resolution and max(h, w) <= self.size:
+            return clip
+        if h > w:
+            w = int(w * self.size / h)
+            h = self.size
+        else:
+            h = int(h * self.size / w)
+            w = self.size
+        resize_clip = resize(clip, target_size=(h, w),
+                                         interpolation_mode=self.interpolation_mode)
+        return resize_clip
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+class CenterCropResizeVideo:
+    '''
+    First use the short side for cropping length,
+    center crop video, then resize to the specified size
+    '''
+
+    def __init__(
+            self,
+            size,
+            top_crop=False, 
+            interpolation_mode="bilinear",
+    ):
+        if len(size) != 2:
+            raise ValueError(f"size should be tuple (height, width), instead got {size}")
+        self.size = size
+        self.top_crop = top_crop
+        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 = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop)
+        # import ipdb;ipdb.set_trace()
+        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
+
+class DynamicSampleDuration(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, t_stride, extra_1):
+        self.t_stride = t_stride
+        self.extra_1 = extra_1
+
+    def __call__(self, t, h, w):
+        if self.extra_1:
+            t = t - 1
+        truncate_t_list = list(range(t+1))[t//2:][::self.t_stride]  # need half at least
+        truncate_t = random.choice(truncate_t_list)
+        if self.extra_1:
+            truncate_t = truncate_t + 1
+        return 0, truncate_t
+
+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))
diff --git a/opensora/eval/cal_flolpips.py b/opensora/eval/cal_flolpips.py
new file mode 100644
index 000000000..eaf7ca9e1
--- /dev/null
+++ b/opensora/eval/cal_flolpips.py
@@ -0,0 +1,83 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+from einops import rearrange
+import sys
+sys.path.append(".")
+from opensora.eval.flolpips.pwcnet import Network as PWCNet
+from opensora.eval.flolpips.flolpips import FloLPIPS
+
+loss_fn = FloLPIPS(net='alex', version='0.1').eval().requires_grad_(False)
+flownet = PWCNet().eval().requires_grad_(False)
+
+def trans(x):
+    return x
+
+
+def calculate_flolpips(videos1, videos2, device):
+    global loss_fn, flownet
+    
+    print("calculate_flowlpips...")
+    loss_fn = loss_fn.to(device)
+    flownet = flownet.to(device)
+    
+    if videos1.shape != videos2.shape:
+        print("Warning: the shape of videos are not equal.")
+        min_frames = min(videos1.shape[1], videos2.shape[1])
+        videos1 = videos1[:, :min_frames]
+        videos2 = videos2[:, :min_frames]
+        
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    flolpips_results = []
+    for video_num in tqdm(range(videos1.shape[0])):
+        video1 = videos1[video_num].to(device)
+        video2 = videos2[video_num].to(device)
+        frames_rec = video1[:-1]
+        frames_rec_next = video1[1:]
+        frames_gt = video2[:-1]
+        frames_gt_next = video2[1:]
+        t, c, h, w = frames_gt.shape
+        flow_gt = flownet(frames_gt, frames_gt_next)
+        flow_dis = flownet(frames_rec, frames_rec_next)
+        flow_diff = flow_gt - flow_dis
+        flolpips = loss_fn.forward(frames_gt, frames_rec, flow_diff, normalize=True)
+        flolpips_results.append(flolpips.cpu().numpy().tolist())
+        
+    flolpips_results = np.array(flolpips_results) # [batch_size, num_frames]
+    flolpips = {}
+    flolpips_std = {}
+
+    for clip_timestamp in range(flolpips_results.shape[1]):
+        flolpips[clip_timestamp] = np.mean(flolpips_results[:,clip_timestamp], axis=-1)
+        flolpips_std[clip_timestamp] = np.std(flolpips_results[:,clip_timestamp], axis=-1)
+
+    result = {
+        "value": flolpips,
+        "value_std": flolpips_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+        "result": flolpips_results,
+        "details": flolpips_results.tolist()
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+
+    import json
+    result = calculate_flolpips(videos1, videos2, "cuda:0")
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/eval/cal_fvd.py b/opensora/eval/cal_fvd.py
new file mode 100644
index 000000000..1f1a9806e
--- /dev/null
+++ b/opensora/eval/cal_fvd.py
@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+def trans(x):
+    # if greyscale images add channel
+    if x.shape[-3] == 1:
+        x = x.repeat(1, 1, 3, 1, 1)
+
+    # permute BTCHW -> BCTHW
+    x = x.permute(0, 2, 1, 3, 4) 
+
+    return x
+
+def calculate_fvd(videos1, videos2, device, method='styleganv'):
+
+    if method == 'styleganv':
+        from fvd.styleganv.fvd import get_fvd_feats, frechet_distance, load_i3d_pretrained
+    elif method == 'videogpt':
+        from fvd.videogpt.fvd import load_i3d_pretrained
+        from fvd.videogpt.fvd import get_fvd_logits as get_fvd_feats
+        from fvd.videogpt.fvd import frechet_distance
+
+    print("calculate_fvd...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    i3d = load_i3d_pretrained(device=device)
+    fvd_results = []
+
+    # support grayscale input, if grayscale -> channel*3
+    # BTCHW -> BCTHW
+    # videos -> [batch_size, channel, timestamps, h, w]
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    fvd_results = {}
+
+    # for calculate FVD, each clip_timestamp must >= 10
+    for clip_timestamp in tqdm(range(10, videos1.shape[-3]+1)):
+       
+        # get a video clip
+        # videos_clip [batch_size, channel, timestamps[:clip], h, w]
+        videos_clip1 = videos1[:, :, : clip_timestamp]
+        videos_clip2 = videos2[:, :, : clip_timestamp]
+
+        # get FVD features
+        feats1 = get_fvd_feats(videos_clip1, i3d=i3d, device=device)
+        feats2 = get_fvd_feats(videos_clip2, i3d=i3d, device=device)
+      
+        # calculate FVD when timestamps[:clip]
+        fvd_results[clip_timestamp] = frechet_distance(feats1, feats2)
+
+    result = {
+        "value": fvd_results,
+        "video_setting": videos1.shape,
+        "video_setting_name": "batch_size, channel, time, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+    # device = torch.device("cpu")
+
+    import json
+    result = calculate_fvd(videos1, videos2, device, method='videogpt')
+    print(json.dumps(result, indent=4))
+
+    result = calculate_fvd(videos1, videos2, device, method='styleganv')
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
diff --git a/opensora/eval/cal_lpips.py b/opensora/eval/cal_lpips.py
new file mode 100644
index 000000000..10edc610b
--- /dev/null
+++ b/opensora/eval/cal_lpips.py
@@ -0,0 +1,97 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+
+import torch
+import lpips
+
+spatial = True         # Return a spatial map of perceptual distance.
+
+# Linearly calibrated models (LPIPS)
+loss_fn = lpips.LPIPS(net='alex', spatial=spatial) # Can also set net = 'squeeze' or 'vgg'
+# loss_fn = lpips.LPIPS(net='alex', spatial=spatial, lpips=False) # Can also set net = 'squeeze' or 'vgg'
+
+def trans(x):
+    # if greyscale images add channel
+    if x.shape[-3] == 1:
+        x = x.repeat(1, 1, 3, 1, 1)
+
+    # value range [0, 1] -> [-1, 1]
+    x = x * 2 - 1
+
+    return x
+
+def calculate_lpips(videos1, videos2, device):
+    # image should be RGB, IMPORTANT: normalized to [-1,1]
+    print("calculate_lpips...")
+
+    assert videos1.shape == videos2.shape
+
+    # videos [batch_size, timestamps, channel, h, w]
+
+    # support grayscale input, if grayscale -> channel*3
+    # value range [0, 1] -> [-1, 1]
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    lpips_results = []
+
+    for video_num in tqdm(range(videos1.shape[0])):
+        # get a video
+        # video [timestamps, channel, h, w]
+        video1 = videos1[video_num]
+        video2 = videos2[video_num]
+
+        lpips_results_of_a_video = []
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] tensor
+
+            img1 = video1[clip_timestamp].unsqueeze(0).to(device)
+            img2 = video2[clip_timestamp].unsqueeze(0).to(device)
+            
+            loss_fn.to(device)
+
+            # calculate lpips of a video
+            lpips_results_of_a_video.append(loss_fn.forward(img1, img2).mean().detach().cpu().tolist())
+        lpips_results.append(lpips_results_of_a_video)
+    
+    lpips_results = np.array(lpips_results)
+    
+    lpips = {}
+    lpips_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        lpips[clip_timestamp] = np.mean(lpips_results[:,clip_timestamp])
+        lpips_std[clip_timestamp] = np.std(lpips_results[:,clip_timestamp])
+
+
+    result = {
+        "value": lpips,
+        "value_std": lpips_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+    # device = torch.device("cpu")
+
+    import json
+    result = calculate_lpips(videos1, videos2, device)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/eval/cal_psnr.py b/opensora/eval/cal_psnr.py
new file mode 100644
index 000000000..b325106c6
--- /dev/null
+++ b/opensora/eval/cal_psnr.py
@@ -0,0 +1,84 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+
+def img_psnr(img1, img2):
+    # [0,1]
+    # compute mse
+    # mse = np.mean((img1-img2)**2)
+    mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2)
+    # compute psnr
+    if mse < 1e-10:
+        return 100
+    psnr = 20 * math.log10(1 / math.sqrt(mse))
+    return psnr
+
+def trans(x):
+    return x
+
+def calculate_psnr(videos1, videos2):
+    print("calculate_psnr...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    psnr_results = []
+    
+    for video_num in tqdm(range(videos1.shape[0])):
+        # get a video
+        # video [timestamps, channel, h, w]
+        video1 = videos1[video_num]
+        video2 = videos2[video_num]
+
+        psnr_results_of_a_video = []
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] numpy
+
+            img1 = video1[clip_timestamp].numpy()
+            img2 = video2[clip_timestamp].numpy()
+            
+            # calculate psnr of a video
+            psnr_results_of_a_video.append(img_psnr(img1, img2))
+
+        psnr_results.append(psnr_results_of_a_video)
+    
+    psnr_results = np.array(psnr_results) # [batch_size, num_frames]
+    psnr = {}
+    psnr_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        psnr[clip_timestamp] = np.mean(psnr_results[:,clip_timestamp])
+        psnr_std[clip_timestamp] = np.std(psnr_results[:,clip_timestamp])
+
+    result = {
+        "value": psnr,
+        "value_std": psnr_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+
+    import json
+    result = calculate_psnr(videos1, videos2)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/eval/cal_ssim.py b/opensora/eval/cal_ssim.py
new file mode 100644
index 000000000..b2de811b6
--- /dev/null
+++ b/opensora/eval/cal_ssim.py
@@ -0,0 +1,113 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import cv2
+ 
+def ssim(img1, img2):
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1 ** 2
+    mu2_sq = mu2 ** 2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()
+ 
+ 
+def calculate_ssim_function(img1, img2):
+    # [0,1]
+    # ssim is the only metric extremely sensitive to gray being compared to b/w 
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[0] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[i], img2[i]))
+            return np.array(ssims).mean()                   
+        elif img1.shape[0] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+def trans(x):
+    return x
+
+def calculate_ssim(videos1, videos2):
+    print("calculate_ssim...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    ssim_results = []
+    
+    for video_num in tqdm(range(videos1.shape[0])):
+        # get a video
+        # video [timestamps, channel, h, w]
+        video1 = videos1[video_num]
+        video2 = videos2[video_num]
+
+        ssim_results_of_a_video = []
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] numpy
+
+            img1 = video1[clip_timestamp].numpy()
+            img2 = video2[clip_timestamp].numpy()
+            
+            # calculate ssim of a video
+            ssim_results_of_a_video.append(calculate_ssim_function(img1, img2))
+
+        ssim_results.append(ssim_results_of_a_video)
+
+    ssim_results = np.array(ssim_results)
+
+    ssim = {}
+    ssim_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        ssim[clip_timestamp] = np.mean(ssim_results[:,clip_timestamp])
+        ssim_std[clip_timestamp] = np.std(ssim_results[:,clip_timestamp])
+
+    result = {
+        "value": ssim,
+        "value_std": ssim_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+
+    import json
+    result = calculate_ssim(videos1, videos2)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/eval/eval_clip_score.py b/opensora/eval/eval_clip_score.py
new file mode 100755
index 000000000..ed2f51082
--- /dev/null
+++ b/opensora/eval/eval_clip_score.py
@@ -0,0 +1,225 @@
+"""Calculates the CLIP Scores
+
+The CLIP model is a contrasitively learned language-image model. There is
+an image encoder and a text encoder. It is believed that the CLIP model could 
+measure the similarity of cross modalities. Please find more information from 
+https://github.com/openai/CLIP.
+
+The CLIP Score measures the Cosine Similarity between two embedded features.
+This repository utilizes the pretrained CLIP Model to calculate 
+the mean average of cosine similarities. 
+
+See --help to see further details.
+
+Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.
+
+Copyright 2023 The Hong Kong Polytechnic University
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+import os.path as osp
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+
+import clip
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+
+IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
+                    'tif', 'tiff', 'webp'}
+
+TEXT_EXTENSIONS = {'txt'}
+
+
+class DummyDataset(Dataset):
+    
+    FLAGS = ['img', 'txt']
+    def __init__(self, real_path, generated_path,
+                 real_flag: str = 'img',
+                 generated_flag: str = 'img',
+                 transform = None,
+                 tokenizer = None) -> None:
+        super().__init__()
+        assert real_flag in self.FLAGS and generated_flag in self.FLAGS, \
+            'CLIP Score only support modality of {}. However, get {} and {}'.format(
+                self.FLAGS, real_flag, generated_flag
+            )
+        self.real_folder = self._combine_without_prefix(real_path)
+        self.real_flag = real_flag
+        self.fake_foler = self._combine_without_prefix(generated_path)
+        self.generated_flag = generated_flag
+        self.transform = transform
+        self.tokenizer = tokenizer
+        # assert self._check()
+
+    def __len__(self):
+        return len(self.real_folder)
+
+    def __getitem__(self, index):
+        if index >= len(self):
+            raise IndexError
+        real_path = self.real_folder[index]
+        generated_path = self.fake_foler[index]
+        real_data = self._load_modality(real_path, self.real_flag)
+        fake_data = self._load_modality(generated_path, self.generated_flag)
+
+        sample = dict(real=real_data, fake=fake_data)
+        return sample
+    
+    def _load_modality(self, path, modality):
+        if modality == 'img':
+            data = self._load_img(path)
+        elif modality == 'txt':
+            data = self._load_txt(path)
+        else:
+            raise TypeError("Got unexpected modality: {}".format(modality))
+        return data
+
+    def _load_img(self, path):
+        img = Image.open(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img
+    
+    def _load_txt(self, path):
+        with open(path, 'r') as fp:
+            data = fp.read()
+            fp.close()
+        if self.tokenizer is not None:
+            data = self.tokenizer(data).squeeze()
+        return data
+
+    def _check(self):
+        for idx in range(len(self)):
+            real_name = self.real_folder[idx].split('.')
+            fake_name = self.fake_folder[idx].split('.')
+            if fake_name != real_name:
+                return False
+        return True
+    
+    def _combine_without_prefix(self, folder_path, prefix='.'):
+        folder = []
+        for name in os.listdir(folder_path):
+            if name[0] == prefix:
+                continue
+            folder.append(osp.join(folder_path, name))
+        folder.sort()
+        return folder
+
+
+@torch.no_grad()
+def calculate_clip_score(dataloader, model, real_flag, generated_flag):
+    score_acc = 0.
+    sample_num = 0.
+    logit_scale = model.logit_scale.exp()
+    for batch_data in tqdm(dataloader):
+        real = batch_data['real']
+        real_features = forward_modality(model, real, real_flag)
+        fake = batch_data['fake']
+        fake_features = forward_modality(model, fake, generated_flag)
+        
+        # normalize features
+        real_features = real_features / real_features.norm(dim=1, keepdim=True).to(torch.float32)
+        fake_features = fake_features / fake_features.norm(dim=1, keepdim=True).to(torch.float32)
+        
+        # calculate scores
+        # score = logit_scale * real_features @ fake_features.t()
+        # score_acc += torch.diag(score).sum()
+        score = logit_scale * (fake_features * real_features).sum()
+        score_acc += score
+        sample_num += real.shape[0]
+    
+    return score_acc / sample_num
+
+        
+def forward_modality(model, data, flag):
+    device = next(model.parameters()).device
+    if flag == 'img':
+        features = model.encode_image(data.to(device))
+    elif flag == 'txt':
+        features = model.encode_text(data.to(device))
+    else:
+        raise TypeError
+    return features
+
+
+def main():
+    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--batch-size', type=int, default=50,
+                    help='Batch size to use')
+    parser.add_argument('--clip-model', type=str, default='ViT-B/32',
+                    help='CLIP model to use')
+    parser.add_argument('--num-workers', type=int, default=8,
+                    help=('Number of processes to use for data loading. '
+                          'Defaults to `min(8, num_cpus)`'))
+    parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+    parser.add_argument('--real_flag', type=str, default='img',
+                    help=('The modality of real path. '
+                          'Default to img'))
+    parser.add_argument('--generated_flag', type=str, default='txt',
+                    help=('The modality of generated path. '
+                          'Default to txt'))
+    parser.add_argument('--real_path', type=str, 
+                    help=('Paths to the real images or '
+                          'to .npz statistic files'))
+    parser.add_argument('--generated_path', type=str,
+                    help=('Paths to the generated images or '
+                          'to .npz statistic files'))
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    if args.num_workers is None:
+        try:
+            num_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            # os.sched_getaffinity is not available under Windows, use
+            # os.cpu_count instead (which may not return the *available* number
+            # of CPUs).
+            num_cpus = os.cpu_count()
+
+        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
+    else:
+        num_workers = args.num_workers
+
+    print('Loading CLIP model: {}'.format(args.clip_model))
+    model, preprocess = clip.load(args.clip_model, device=device)
+    
+    dataset = DummyDataset(args.real_path, args.generated_path,
+                           args.real_flag, args.generated_flag,
+                           transform=preprocess, tokenizer=clip.tokenize)
+    dataloader = DataLoader(dataset, args.batch_size, 
+                            num_workers=num_workers, pin_memory=True)
+    
+    print('Calculating CLIP Score:')
+    clip_score = calculate_clip_score(dataloader, model,
+                                      args.real_flag, args.generated_flag)
+    clip_score = clip_score.cpu().item()
+    print('CLIP Score: ', clip_score)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/opensora/eval/eval_common_metric.py b/opensora/eval/eval_common_metric.py
new file mode 100644
index 000000000..452c03209
--- /dev/null
+++ b/opensora/eval/eval_common_metric.py
@@ -0,0 +1,224 @@
+"""Calculates the CLIP Scores
+
+The CLIP model is a contrasitively learned language-image model. There is
+an image encoder and a text encoder. It is believed that the CLIP model could 
+measure the similarity of cross modalities. Please find more information from 
+https://github.com/openai/CLIP.
+
+The CLIP Score measures the Cosine Similarity between two embedded features.
+This repository utilizes the pretrained CLIP Model to calculate 
+the mean average of cosine similarities. 
+
+See --help to see further details.
+
+Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.
+
+Copyright 2023 The Hong Kong Polytechnic University
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+import os.path as osp
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader, Subset
+from decord import VideoReader, cpu
+import random
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.io import read_video
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms._transforms_video import CenterCropVideo
+import sys
+sys.path.append(".")
+from opensora.eval.cal_lpips import calculate_lpips
+from opensora.eval.cal_fvd import calculate_fvd
+from opensora.eval.cal_psnr import calculate_psnr
+from opensora.eval.cal_flolpips import calculate_flolpips
+from opensora.eval.cal_ssim import calculate_ssim
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+class VideoDataset(Dataset):
+    def __init__(self, 
+                 real_video_dir,
+                 generated_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.generated_video_files = self._combine_without_prefix(generated_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]
+        generated_video_file = self.generated_video_files[index]
+        print(real_video_file, generated_video_file)
+        real_video_tensor  = self._load_video(real_video_file)
+        generated_video_tensor  = self._load_video(generated_video_file)
+        return {'real': real_video_tensor, 'generated':generated_video_tensor }
+
+
+    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(0, 3, 1, 2) # (T, H, W, C) -> (C, T, H, W)
+        return _preprocess(video_data, short_size=self.short_size, crop_size = self.crop_size)
+
+
+    def _combine_without_prefix(self, folder_path, prefix='.'):
+        folder = []
+        os.makedirs(folder_path, exist_ok=True)
+        for name in os.listdir(folder_path):
+            if name[0] == prefix:
+                continue
+            if osp.isfile(osp.join(folder_path, name)):
+                folder.append(osp.join(folder_path, name))
+        folder.sort()
+        return folder
+
+def _preprocess(video_data, short_size=128, crop_size=None):
+    transform = Compose(
+        [
+            Lambda(lambda x: x / 255.0),
+            ShortSideScale(size=short_size),
+            CenterCropVideo(crop_size=crop_size),
+        ]
+    )
+    video_outputs = transform(video_data)
+    # video_outputs = torch.unsqueeze(video_outputs, 0) # (bz,c,t,h,w)
+    return video_outputs
+
+
+def calculate_common_metric(args, dataloader, device):
+
+    score_list = []
+    for batch_data in tqdm(dataloader): # {'real': real_video_tensor, 'generated':generated_video_tensor }
+        real_videos = batch_data['real'] 
+        generated_videos = batch_data['generated']
+        assert real_videos.shape[2] == generated_videos.shape[2]
+        if args.metric == 'fvd':
+            tmp_list = list(calculate_fvd(real_videos, generated_videos, args.device, method=args.fvd_method)['value'].values())
+        elif args.metric == 'ssim':
+            tmp_list = list(calculate_ssim(real_videos, generated_videos)['value'].values())
+        elif args.metric == 'psnr':
+            tmp_list = list(calculate_psnr(real_videos, generated_videos)['value'].values())
+        elif args.metric == 'flolpips':
+            result = calculate_flolpips(real_videos, generated_videos, args.device)
+            tmp_list = list(result['value'].values())
+        else:
+            tmp_list  = list(calculate_lpips(real_videos, generated_videos, args.device)['value'].values())
+        score_list += tmp_list
+    return np.mean(score_list)
+        
+def main():
+    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--batch_size', type=int, default=2,
+                    help='Batch size to use')
+    parser.add_argument('--real_video_dir', type=str,
+                    help=('the path of real videos`'))
+    parser.add_argument('--generated_video_dir', type=str,
+                    help=('the path of generated videos`'))
+    parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+    parser.add_argument('--num_workers', type=int, default=8,
+                    help=('Number of processes to use for data loading. '
+                          'Defaults to `min(8, num_cpus)`'))
+    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('--subset_size', type=int, default=None)
+    parser.add_argument("--metric", type=str, default="fvd",choices=['fvd','psnr','ssim','lpips', 'flolpips'])
+    parser.add_argument("--fvd_method", type=str, default='styleganv',choices=['styleganv','videogpt'])
+
+
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    if args.num_workers is None:
+        try:
+            num_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            # os.sched_getaffinity is not available under Windows, use
+            # os.cpu_count instead (which may not return the *available* number
+            # of CPUs).
+            num_cpus = os.cpu_count()
+
+        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
+    else:
+        num_workers = args.num_workers
+
+    
+    dataset = VideoDataset(args.real_video_dir,
+                           args.generated_video_dir,
+                            num_frames = args.num_frames,
+                            sample_rate = args.sample_rate,
+                            crop_size=args.crop_size,
+                            resolution=args.resolution) 
+
+    if args.subset_size:
+        indices = range(args.subset_size)
+        dataset = Subset(dataset, indices=indices)
+    
+    dataloader = DataLoader(dataset, args.batch_size, 
+                            num_workers=num_workers, pin_memory=True)
+    
+
+    metric_score = calculate_common_metric(args, dataloader,device)
+    print('metric: ', args.metric, " ",metric_score)
+
+if __name__ == '__main__':
+    main()
diff --git a/opensora/eval/flolpips/correlation/correlation.py b/opensora/eval/flolpips/correlation/correlation.py
new file mode 100644
index 000000000..7c91055ba
--- /dev/null
+++ b/opensora/eval/flolpips/correlation/correlation.py
@@ -0,0 +1,397 @@
+#!/usr/bin/env python
+
+import torch
+
+import cupy
+import re
+
+kernel_Correlation_rearrange = '''
+	extern "C" __global__ void kernel_Correlation_rearrange(
+		const int n,
+		const float* input,
+		float* output
+	) {
+	  int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x;
+
+	  if (intIndex >= n) {
+	    return;
+	  }
+
+	  int intSample = blockIdx.z;
+	  int intChannel = blockIdx.y;
+
+	  float fltValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex];
+
+	  __syncthreads();
+
+	  int intPaddedY = (intIndex / SIZE_3(input)) + 4;
+	  int intPaddedX = (intIndex % SIZE_3(input)) + 4;
+	  int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX;
+
+	  output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = fltValue;
+	}
+'''
+
+kernel_Correlation_updateOutput = '''
+	extern "C" __global__ void kernel_Correlation_updateOutput(
+	  const int n,
+	  const float* rbot0,
+	  const float* rbot1,
+	  float* top
+	) {
+	  extern __shared__ char patch_data_char[];
+	  
+	  float *patch_data = (float *)patch_data_char;
+	  
+	  // First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1
+	  int x1 = blockIdx.x + 4;
+	  int y1 = blockIdx.y + 4;
+	  int item = blockIdx.z;
+	  int ch_off = threadIdx.x;
+	  
+	  // Load 3D patch into shared shared memory
+	  for (int j = 0; j < 1; j++) { // HEIGHT
+	    for (int i = 0; i < 1; i++) { // WIDTH
+	      int ji_off = (j + i) * SIZE_3(rbot0);
+	      for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
+	        int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch;
+	        int idxPatchData = ji_off + ch;
+	        patch_data[idxPatchData] = rbot0[idx1];
+	      }
+	    }
+	  }
+	  
+	  __syncthreads();
+	  
+	  __shared__ float sum[32];
+	  
+	  // Compute correlation
+	  for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) {
+	    sum[ch_off] = 0;
+	  
+	    int s2o = top_channel % 9 - 4;
+	    int s2p = top_channel / 9 - 4;
+	    
+	    for (int j = 0; j < 1; j++) { // HEIGHT
+	      for (int i = 0; i < 1; i++) { // WIDTH
+	        int ji_off = (j + i) * SIZE_3(rbot0);
+	        for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
+	          int x2 = x1 + s2o;
+	          int y2 = y1 + s2p;
+	          
+	          int idxPatchData = ji_off + ch;
+	          int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch;
+	          
+	          sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2];
+	        }
+	      }
+	    }
+	    
+	    __syncthreads();
+	    
+	    if (ch_off == 0) {
+	      float total_sum = 0;
+	      for (int idx = 0; idx < 32; idx++) {
+	        total_sum += sum[idx];
+	      }
+	      const int sumelems = SIZE_3(rbot0);
+	      const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x;
+	      top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems;
+	    }
+	  }
+	}
+'''
+
+kernel_Correlation_updateGradFirst = '''
+	#define ROUND_OFF 50000
+
+	extern "C" __global__ void kernel_Correlation_updateGradFirst(
+	  const int n,
+	  const int intSample,
+	  const float* rbot0,
+	  const float* rbot1,
+	  const float* gradOutput,
+	  float* gradFirst,
+	  float* gradSecond
+	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
+	  int n = intIndex % SIZE_1(gradFirst); // channels
+	  int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos
+	  int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos
+	  
+	  // round_off is a trick to enable integer division with ceil, even for negative numbers
+	  // We use a large offset, for the inner part not to become negative.
+	  const int round_off = ROUND_OFF;
+	  const int round_off_s1 = round_off;
+	  
+	  // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
+	  int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
+	  int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
+	  
+	  // Same here:
+	  int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4)
+	  int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4)
+	  
+	  float sum = 0;
+	  if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
+	    xmin = max(0,xmin);
+	    xmax = min(SIZE_3(gradOutput)-1,xmax);
+	    
+	    ymin = max(0,ymin);
+	    ymax = min(SIZE_2(gradOutput)-1,ymax);
+	    
+	    for (int p = -4; p <= 4; p++) {
+	      for (int o = -4; o <= 4; o++) {
+	        // Get rbot1 data:
+	        int s2o = o;
+	        int s2p = p;
+	        int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n;
+	        float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n]
+	        
+	        // Index offset for gradOutput in following loops:
+	        int op = (p+4) * 9 + (o+4); // index[o,p]
+	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
+	        
+	        for (int y = ymin; y <= ymax; y++) {
+	          for (int x = xmin; x <= xmax; x++) {
+	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
+	            sum += gradOutput[idxgradOutput] * bot1tmp;
+	          }
+	        }
+	      }
+	    }
+	  }
+	  const int sumelems = SIZE_1(gradFirst);
+	  const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4);
+	  gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems;
+	} }
+'''
+
+kernel_Correlation_updateGradSecond = '''
+	#define ROUND_OFF 50000
+
+	extern "C" __global__ void kernel_Correlation_updateGradSecond(
+	  const int n,
+	  const int intSample,
+	  const float* rbot0,
+	  const float* rbot1,
+	  const float* gradOutput,
+	  float* gradFirst,
+	  float* gradSecond
+	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
+	  int n = intIndex % SIZE_1(gradSecond); // channels
+	  int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos
+	  int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos
+	  
+	  // round_off is a trick to enable integer division with ceil, even for negative numbers
+	  // We use a large offset, for the inner part not to become negative.
+	  const int round_off = ROUND_OFF;
+	  const int round_off_s1 = round_off;
+	  
+	  float sum = 0;
+	  for (int p = -4; p <= 4; p++) {
+	    for (int o = -4; o <= 4; o++) {
+	      int s2o = o;
+	      int s2p = p;
+	      
+	      //Get X,Y ranges and clamp
+	      // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
+	      int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
+	      int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
+	      
+	      // Same here:
+	      int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o)
+	      int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p)
+          
+	      if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
+	        xmin = max(0,xmin);
+	        xmax = min(SIZE_3(gradOutput)-1,xmax);
+	        
+	        ymin = max(0,ymin);
+	        ymax = min(SIZE_2(gradOutput)-1,ymax);
+	        
+	        // Get rbot0 data:
+	        int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n;
+	        float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n]
+	        
+	        // Index offset for gradOutput in following loops:
+	        int op = (p+4) * 9 + (o+4); // index[o,p]
+	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
+	        
+	        for (int y = ymin; y <= ymax; y++) {
+	          for (int x = xmin; x <= xmax; x++) {
+	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
+	            sum += gradOutput[idxgradOutput] * bot0tmp;
+	          }
+	        }
+	      }
+	    }
+	  }
+	  const int sumelems = SIZE_1(gradSecond);
+	  const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4);
+	  gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems;
+	} }
+'''
+
+def cupy_kernel(strFunction, objVariables):
+	strKernel = globals()[strFunction]
+
+	while True:
+		objMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel)
+
+		if objMatch is None:
+			break
+		# end
+
+		intArg = int(objMatch.group(2))
+
+		strTensor = objMatch.group(4)
+		intSizes = objVariables[strTensor].size()
+
+		strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg]))
+	# end
+
+	while True:
+		objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)
+
+		if objMatch is None:
+			break
+		# end
+
+		intArgs = int(objMatch.group(2))
+		strArgs = objMatch.group(4).split(',')
+
+		strTensor = strArgs[0]
+		intStrides = objVariables[strTensor].stride()
+		strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]
+
+		strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')
+	# end
+
+	return strKernel
+# end
+
+@cupy.memoize(for_each_device=True)
+def cupy_launch(strFunction, strKernel):
+    return cupy.RawKernel(strKernel, strFunction)
+# end
+
+class _FunctionCorrelation(torch.autograd.Function):
+	@staticmethod
+	def forward(self, first, second):
+		rbot0 = first.new_zeros([ first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1] ])
+		rbot1 = first.new_zeros([ first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1] ])
+
+		self.save_for_backward(first, second, rbot0, rbot1)
+
+		first = first.contiguous();	assert(first.is_cuda == True)
+		second = second.contiguous(); assert(second.is_cuda == True)
+
+		output = first.new_zeros([ first.shape[0], 81, first.shape[2], first.shape[3] ])
+
+		if first.is_cuda == True:
+			n = first.shape[2] * first.shape[3]
+			cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
+				'input': first,
+				'output': rbot0
+			}))(
+				grid=tuple([ int((n + 16 - 1) / 16), first.shape[1], first.shape[0] ]),
+				block=tuple([ 16, 1, 1 ]),
+				args=[ n, first.data_ptr(), rbot0.data_ptr() ]
+			)
+
+			n = second.shape[2] * second.shape[3]
+			cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
+				'input': second,
+				'output': rbot1
+			}))(
+				grid=tuple([ int((n + 16 - 1) / 16), second.shape[1], second.shape[0] ]),
+				block=tuple([ 16, 1, 1 ]),
+				args=[ n, second.data_ptr(), rbot1.data_ptr() ]
+			)
+
+			n = output.shape[1] * output.shape[2] * output.shape[3]
+			cupy_launch('kernel_Correlation_updateOutput', cupy_kernel('kernel_Correlation_updateOutput', {
+				'rbot0': rbot0,
+				'rbot1': rbot1,
+				'top': output
+			}))(
+				grid=tuple([ output.shape[3], output.shape[2], output.shape[0] ]),
+				block=tuple([ 32, 1, 1 ]),
+				shared_mem=first.shape[1] * 4,
+				args=[ n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr() ]
+			)
+
+		elif first.is_cuda == False:
+			raise NotImplementedError()
+
+		# end
+
+		return output
+	# end
+
+	@staticmethod
+	def backward(self, gradOutput):
+		first, second, rbot0, rbot1 = self.saved_tensors
+
+		gradOutput = gradOutput.contiguous(); assert(gradOutput.is_cuda == True)
+
+		gradFirst = first.new_zeros([ first.shape[0], first.shape[1], first.shape[2], first.shape[3] ]) if self.needs_input_grad[0] == True else None
+		gradSecond = first.new_zeros([ first.shape[0], first.shape[1], first.shape[2], first.shape[3] ]) if self.needs_input_grad[1] == True else None
+
+		if first.is_cuda == True:
+			if gradFirst is not None:
+				for intSample in range(first.shape[0]):
+					n = first.shape[1] * first.shape[2] * first.shape[3]
+					cupy_launch('kernel_Correlation_updateGradFirst', cupy_kernel('kernel_Correlation_updateGradFirst', {
+						'rbot0': rbot0,
+						'rbot1': rbot1,
+						'gradOutput': gradOutput,
+						'gradFirst': gradFirst,
+						'gradSecond': None
+					}))(
+						grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
+						block=tuple([ 512, 1, 1 ]),
+						args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), gradFirst.data_ptr(), None ]
+					)
+				# end
+			# end
+
+			if gradSecond is not None:
+				for intSample in range(first.shape[0]):
+					n = first.shape[1] * first.shape[2] * first.shape[3]
+					cupy_launch('kernel_Correlation_updateGradSecond', cupy_kernel('kernel_Correlation_updateGradSecond', {
+						'rbot0': rbot0,
+						'rbot1': rbot1,
+						'gradOutput': gradOutput,
+						'gradFirst': None,
+						'gradSecond': gradSecond
+					}))(
+						grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
+						block=tuple([ 512, 1, 1 ]),
+						args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), None, gradSecond.data_ptr() ]
+					)
+				# end
+			# end
+
+		elif first.is_cuda == False:
+			raise NotImplementedError()
+
+		# end
+
+		return gradFirst, gradSecond
+	# end
+# end
+
+def FunctionCorrelation(tenFirst, tenSecond):
+	return _FunctionCorrelation.apply(tenFirst, tenSecond)
+# end
+
+class ModuleCorrelation(torch.nn.Module):
+	def __init__(self):
+		super(ModuleCorrelation, self).__init__()
+	# end
+
+	def forward(self, tenFirst, tenSecond):
+		return _FunctionCorrelation.apply(tenFirst, tenSecond)
+	# end
+# end
\ No newline at end of file
diff --git a/opensora/eval/flolpips/flolpips.py b/opensora/eval/flolpips/flolpips.py
new file mode 100644
index 000000000..66f473149
--- /dev/null
+++ b/opensora/eval/flolpips/flolpips.py
@@ -0,0 +1,308 @@
+
+from __future__ import absolute_import
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .pretrained_networks import vgg16, alexnet, squeezenet
+import torch.nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+import cv2
+
+from .pwcnet import Network as PWCNet
+from .utils import *
+
+def spatial_average(in_tens, keepdim=True):
+    return in_tens.mean([2,3],keepdim=keepdim)
+
+def mw_spatial_average(in_tens, flow, keepdim=True):
+    _,_,h,w = in_tens.shape
+    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
+    flow_mag = torch.sqrt(flow[:,0:1]**2 + flow[:,1:2]**2)
+    flow_mag = flow_mag / torch.sum(flow_mag, dim=[1,2,3], keepdim=True)
+    return torch.sum(in_tens*flow_mag, dim=[2,3],keepdim=keepdim)
+
+
+def mtw_spatial_average(in_tens, flow, texture, keepdim=True):
+    _,_,h,w = in_tens.shape
+    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
+    texture = F.interpolate(texture, (h,w), align_corners=False, mode='bilinear')
+    flow_mag = torch.sqrt(flow[:,0:1]**2 + flow[:,1:2]**2)
+    flow_mag = (flow_mag - flow_mag.min()) / (flow_mag.max() - flow_mag.min()) + 1e-6
+    texture = (texture - texture.min()) / (texture.max() - texture.min()) + 1e-6
+    weight = flow_mag / texture
+    weight /= torch.sum(weight)
+    return torch.sum(in_tens*weight, dim=[2,3],keepdim=keepdim)
+
+
+
+def m2w_spatial_average(in_tens, flow, keepdim=True):
+    _,_,h,w = in_tens.shape
+    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
+    flow_mag = flow[:,0:1]**2 + flow[:,1:2]**2 # B,1,H,W
+    flow_mag = flow_mag / torch.sum(flow_mag)
+    return torch.sum(in_tens*flow_mag, dim=[2,3],keepdim=keepdim)
+
+def upsample(in_tens, out_HW=(64,64)): # assumes scale factor is same for H and W
+    in_H, in_W = in_tens.shape[2], in_tens.shape[3]
+    return nn.Upsample(size=out_HW, mode='bilinear', align_corners=False)(in_tens)
+
+# Learned perceptual metric
+class LPIPS(nn.Module):
+    def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, 
+        pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=False):
+        # lpips - [True] means with linear calibration on top of base network
+        # pretrained - [True] means load linear weights
+
+        super(LPIPS, self).__init__()
+        if(verbose):
+            print('Setting up [%s] perceptual loss: trunk [%s], v[%s], spatial [%s]'%
+                ('LPIPS' if lpips else 'baseline', net, version, 'on' if spatial else 'off'))
+
+        self.pnet_type = net
+        self.pnet_tune = pnet_tune
+        self.pnet_rand = pnet_rand
+        self.spatial = spatial
+        self.lpips = lpips # false means baseline of just averaging all layers
+        self.version = version
+        self.scaling_layer = ScalingLayer()
+
+        if(self.pnet_type in ['vgg','vgg16']):
+            net_type = vgg16
+            self.chns = [64,128,256,512,512]
+        elif(self.pnet_type=='alex'):
+            net_type = alexnet
+            self.chns = [64,192,384,256,256]
+        elif(self.pnet_type=='squeeze'):
+            net_type = squeezenet
+            self.chns = [64,128,256,384,384,512,512]
+        self.L = len(self.chns)
+
+        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)
+
+        if(lpips):
+            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+            self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4]
+            if(self.pnet_type=='squeeze'): # 7 layers for squeezenet
+                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
+                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
+                self.lins+=[self.lin5,self.lin6]
+            self.lins = nn.ModuleList(self.lins)
+
+            if(pretrained):
+                if(model_path is None):
+                    import inspect
+                    import os
+                    model_path = os.path.abspath(os.path.join(inspect.getfile(self.__init__), '..', 'weights/v%s/%s.pth'%(version,net)))
+
+                if(verbose):
+                    print('Loading model from: %s'%model_path)
+                self.load_state_dict(torch.load(model_path, map_location='cpu'), strict=False)          
+
+        if(eval_mode):
+            self.eval()
+
+    def forward(self, in0, in1, retPerLayer=False, normalize=False):
+        if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
+            in0 = 2 * in0  - 1
+            in1 = 2 * in1  - 1
+
+        # v0.0 - original release had a bug, where input was not scaled
+        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
+        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+
+        for kk in range(self.L):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk]-feats1[kk])**2
+
+        if(self.lpips):
+            if(self.spatial):
+                res = [upsample(self.lins[kk](diffs[kk]), out_HW=in0.shape[2:]) for kk in range(self.L)]
+            else:
+                res = [spatial_average(self.lins[kk](diffs[kk]), keepdim=True) for kk in range(self.L)]
+        else:
+            if(self.spatial):
+                res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_HW=in0.shape[2:]) for kk in range(self.L)]
+            else:
+                res = [spatial_average(diffs[kk].sum(dim=1,keepdim=True), keepdim=True) for kk in range(self.L)]
+
+        # val = res[0]
+        # for l in range(1,self.L):
+        #     val += res[l]
+        #     print(val)
+
+        # a = spatial_average(self.lins[kk](diffs[kk]), keepdim=True)
+        # b = torch.max(self.lins[kk](feats0[kk]**2))
+        # for kk in range(self.L):
+        #     a += spatial_average(self.lins[kk](diffs[kk]), keepdim=True)
+        #     b = torch.max(b,torch.max(self.lins[kk](feats0[kk]**2)))
+        # a = a/self.L
+        # from IPython import embed
+        # embed()
+        # return 10*torch.log10(b/a)
+        
+        # if(retPerLayer):
+        #     return (val, res)
+        # else:
+        return torch.sum(torch.cat(res, 1), dim=(1,2,3), keepdims=False)
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030,-.088,-.188])[None,:,None,None])
+        self.register_buffer('scale', torch.Tensor([.458,.448,.450])[None,:,None,None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    ''' A single linear layer which does a 1x1 conv '''
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+
+        layers = [nn.Dropout(),] if(use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),]
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+class Dist2LogitLayer(nn.Module):
+    ''' takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) '''
+    def __init__(self, chn_mid=32, use_sigmoid=True):
+        super(Dist2LogitLayer, self).__init__()
+
+        layers = [nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),]
+        layers += [nn.LeakyReLU(0.2,True),]
+        layers += [nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),]
+        layers += [nn.LeakyReLU(0.2,True),]
+        layers += [nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),]
+        if(use_sigmoid):
+            layers += [nn.Sigmoid(),]
+        self.model = nn.Sequential(*layers)
+
+    def forward(self,d0,d1,eps=0.1):
+        return self.model.forward(torch.cat((d0,d1,d0-d1,d0/(d1+eps),d1/(d0+eps)),dim=1))
+
+class BCERankingLoss(nn.Module):
+    def __init__(self, chn_mid=32):
+        super(BCERankingLoss, self).__init__()
+        self.net = Dist2LogitLayer(chn_mid=chn_mid)
+        # self.parameters = list(self.net.parameters())
+        self.loss = torch.nn.BCELoss()
+
+    def forward(self, d0, d1, judge):
+        per = (judge+1.)/2.
+        self.logit = self.net.forward(d0,d1)
+        return self.loss(self.logit, per)
+
+# L2, DSSIM metrics
+class FakeNet(nn.Module):
+    def __init__(self, use_gpu=True, colorspace='Lab'):
+        super(FakeNet, self).__init__()
+        self.use_gpu = use_gpu
+        self.colorspace = colorspace
+
+class L2(FakeNet):
+    def forward(self, in0, in1, retPerLayer=None):
+        assert(in0.size()[0]==1) # currently only supports batchSize 1
+
+        if(self.colorspace=='RGB'):
+            (N,C,X,Y) = in0.size()
+            value = torch.mean(torch.mean(torch.mean((in0-in1)**2,dim=1).view(N,1,X,Y),dim=2).view(N,1,1,Y),dim=3).view(N)
+            return value
+        elif(self.colorspace=='Lab'):
+            value = l2(tensor2np(tensor2tensorlab(in0.data,to_norm=False)), 
+                tensor2np(tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
+            ret_var = Variable( torch.Tensor((value,) ) )
+            if(self.use_gpu):
+                ret_var = ret_var.cuda()
+            return ret_var
+
+class DSSIM(FakeNet):
+
+    def forward(self, in0, in1, retPerLayer=None):
+        assert(in0.size()[0]==1) # currently only supports batchSize 1
+
+        if(self.colorspace=='RGB'):
+            value = dssim(1.*tensor2im(in0.data), 1.*tensor2im(in1.data), range=255.).astype('float')
+        elif(self.colorspace=='Lab'):
+            value = dssim(tensor2np(tensor2tensorlab(in0.data,to_norm=False)), 
+                tensor2np(tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
+        ret_var = Variable( torch.Tensor((value,) ) )
+        if(self.use_gpu):
+            ret_var = ret_var.cuda()
+        return ret_var
+
+def print_network(net):
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print('Network',net)
+    print('Total number of parameters: %d' % num_params)
+
+
+class FloLPIPS(LPIPS):
+    def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=False):
+        super(FloLPIPS, self).__init__(pretrained, net, version, lpips, spatial, pnet_rand, pnet_tune, use_dropout, model_path, eval_mode, verbose)
+
+    def forward(self, in0, in1, flow, retPerLayer=False, normalize=False):
+        if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
+            in0 = 2 * in0  - 1
+            in1 = 2 * in1  - 1
+
+        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
+        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+
+        for kk in range(self.L):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk]-feats1[kk])**2
+
+        res = [mw_spatial_average(self.lins[kk](diffs[kk]), flow, keepdim=True) for kk in range(self.L)]
+
+        return torch.sum(torch.cat(res, 1), dim=(1,2,3), keepdims=False)
+
+
+
+
+
+class Flolpips(nn.Module):
+    def __init__(self):
+        super(Flolpips, self).__init__()
+        self.loss_fn = FloLPIPS(net='alex',version='0.1')
+        self.flownet = PWCNet()
+    
+    @torch.no_grad()
+    def forward(self, I0, I1, frame_dis, frame_ref):
+        """
+        args:
+            I0: first frame of the triplet, shape: [B, C, H, W]
+            I1: third frame of the triplet, shape: [B, C, H, W]
+            frame_dis: prediction of the intermediate frame, shape: [B, C, H, W]
+            frame_ref: ground-truth of the intermediate frame, shape: [B, C, H, W]
+        """
+        assert I0.size() == I1.size() == frame_dis.size() == frame_ref.size(), \
+                "the 4 input tensors should have same size"
+
+        flow_ref = self.flownet(frame_ref, I0)
+        flow_dis = self.flownet(frame_dis, I0)
+        flow_diff = flow_ref - flow_dis
+        flolpips_wrt_I0 = self.loss_fn.forward(frame_ref, frame_dis, flow_diff, normalize=True)
+
+        flow_ref = self.flownet(frame_ref, I1)
+        flow_dis = self.flownet(frame_dis, I1)
+        flow_diff = flow_ref - flow_dis
+        flolpips_wrt_I1 = self.loss_fn.forward(frame_ref, frame_dis, flow_diff, normalize=True)
+
+        flolpips = (flolpips_wrt_I0 + flolpips_wrt_I1) / 2
+        return flolpips
\ No newline at end of file
diff --git a/opensora/eval/flolpips/pretrained_networks.py b/opensora/eval/flolpips/pretrained_networks.py
new file mode 100644
index 000000000..82f14e206
--- /dev/null
+++ b/opensora/eval/flolpips/pretrained_networks.py
@@ -0,0 +1,180 @@
+from collections import namedtuple
+import torch
+from torchvision import models as tv
+
+class squeezenet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(squeezenet, self).__init__()
+        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.slice6 = torch.nn.Sequential()
+        self.slice7 = torch.nn.Sequential()
+        self.N_slices = 7
+        for x in range(2):
+            self.slice1.add_module(str(x), pretrained_features[x])
+        for x in range(2,5):
+            self.slice2.add_module(str(x), pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), pretrained_features[x])
+        for x in range(10, 11):
+            self.slice5.add_module(str(x), pretrained_features[x])
+        for x in range(11, 12):
+            self.slice6.add_module(str(x), pretrained_features[x])
+        for x in range(12, 13):
+            self.slice7.add_module(str(x), pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        h = self.slice6(h)
+        h_relu6 = h
+        h = self.slice7(h)
+        h_relu7 = h
+        vgg_outputs = namedtuple("SqueezeOutputs", ['relu1','relu2','relu3','relu4','relu5','relu6','relu7'])
+        out = vgg_outputs(h_relu1,h_relu2,h_relu3,h_relu4,h_relu5,h_relu6,h_relu7)
+
+        return out
+
+
+class alexnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(alexnet, self).__init__()
+        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(2):
+            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(2, 5):
+            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(10, 12):
+            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        alexnet_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5'])
+        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
+
+        return out
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+
+        return out
+
+
+
+class resnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True, num=18):
+        super(resnet, self).__init__()
+        if(num==18):
+            self.net = tv.resnet18(pretrained=pretrained)
+        elif(num==34):
+            self.net = tv.resnet34(pretrained=pretrained)
+        elif(num==50):
+            self.net = tv.resnet50(pretrained=pretrained)
+        elif(num==101):
+            self.net = tv.resnet101(pretrained=pretrained)
+        elif(num==152):
+            self.net = tv.resnet152(pretrained=pretrained)
+        self.N_slices = 5
+
+        self.conv1 = self.net.conv1
+        self.bn1 = self.net.bn1
+        self.relu = self.net.relu
+        self.maxpool = self.net.maxpool
+        self.layer1 = self.net.layer1
+        self.layer2 = self.net.layer2
+        self.layer3 = self.net.layer3
+        self.layer4 = self.net.layer4
+
+    def forward(self, X):
+        h = self.conv1(X)
+        h = self.bn1(h)
+        h = self.relu(h)
+        h_relu1 = h
+        h = self.maxpool(h)
+        h = self.layer1(h)
+        h_conv2 = h
+        h = self.layer2(h)
+        h_conv3 = h
+        h = self.layer3(h)
+        h_conv4 = h
+        h = self.layer4(h)
+        h_conv5 = h
+
+        outputs = namedtuple("Outputs", ['relu1','conv2','conv3','conv4','conv5'])
+        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
+
+        return out
diff --git a/opensora/eval/flolpips/pwcnet.py b/opensora/eval/flolpips/pwcnet.py
new file mode 100644
index 000000000..374f18e20
--- /dev/null
+++ b/opensora/eval/flolpips/pwcnet.py
@@ -0,0 +1,344 @@
+#!/usr/bin/env python
+
+import torch
+
+import getopt
+import math
+import numpy
+import os
+import PIL
+import PIL.Image
+import sys
+
+# try:
+from .correlation import correlation # the custom cost volume layer
+# except:
+# 	sys.path.insert(0, './correlation'); import correlation # you should consider upgrading python
+# end
+
+##########################################################
+
+# assert(int(str('').join(torch.__version__.split('.')[0:2])) >= 13) # requires at least pytorch version 1.3.0
+
+# torch.set_grad_enabled(False) # make sure to not compute gradients for computational performance
+
+# torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
+
+# ##########################################################
+
+# arguments_strModel = 'default' # 'default', or 'chairs-things'
+# arguments_strFirst = './images/first.png'
+# arguments_strSecond = './images/second.png'
+# arguments_strOut = './out.flo'
+
+# for strOption, strArgument in getopt.getopt(sys.argv[1:], '', [ strParameter[2:] + '=' for strParameter in sys.argv[1::2] ])[0]:
+# 	if strOption == '--model' and strArgument != '': arguments_strModel = strArgument # which model to use
+# 	if strOption == '--first' and strArgument != '': arguments_strFirst = strArgument # path to the first frame
+# 	if strOption == '--second' and strArgument != '': arguments_strSecond = strArgument # path to the second frame
+# 	if strOption == '--out' and strArgument != '': arguments_strOut = strArgument # path to where the output should be stored
+# end
+
+##########################################################
+
+
+
+def backwarp(tenInput, tenFlow):
+	backwarp_tenGrid = {}
+	backwarp_tenPartial = {}
+	if str(tenFlow.shape) not in backwarp_tenGrid:
+		tenHor = torch.linspace(-1.0 + (1.0 / tenFlow.shape[3]), 1.0 - (1.0 / tenFlow.shape[3]), tenFlow.shape[3]).view(1, 1, 1, -1).expand(-1, -1, tenFlow.shape[2], -1)
+		tenVer = torch.linspace(-1.0 + (1.0 / tenFlow.shape[2]), 1.0 - (1.0 / tenFlow.shape[2]), tenFlow.shape[2]).view(1, 1, -1, 1).expand(-1, -1, -1, tenFlow.shape[3])
+
+		backwarp_tenGrid[str(tenFlow.shape)] = torch.cat([ tenHor, tenVer ], 1).cuda()
+	# end
+
+	if str(tenFlow.shape) not in backwarp_tenPartial:
+		backwarp_tenPartial[str(tenFlow.shape)] = tenFlow.new_ones([ tenFlow.shape[0], 1, tenFlow.shape[2], tenFlow.shape[3] ])
+	# end
+
+	tenFlow = torch.cat([ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0) ], 1)
+	tenInput = torch.cat([ tenInput, backwarp_tenPartial[str(tenFlow.shape)] ], 1)
+
+	tenOutput = torch.nn.functional.grid_sample(input=tenInput, grid=(backwarp_tenGrid[str(tenFlow.shape)] + tenFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros', align_corners=False)
+
+	tenMask = tenOutput[:, -1:, :, :]; tenMask[tenMask > 0.999] = 1.0; tenMask[tenMask < 1.0] = 0.0
+
+	return tenOutput[:, :-1, :, :] * tenMask
+# end
+
+##########################################################
+
+class Network(torch.nn.Module):
+	def __init__(self):
+		super(Network, self).__init__()
+
+		class Extractor(torch.nn.Module):
+			def __init__(self):
+				super(Extractor, self).__init__()
+
+				self.netOne = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netTwo = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netThr = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFou = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFiv = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netSix = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=128, out_channels=196, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+			# end
+
+			def forward(self, tenInput):
+				tenOne = self.netOne(tenInput)
+				tenTwo = self.netTwo(tenOne)
+				tenThr = self.netThr(tenTwo)
+				tenFou = self.netFou(tenThr)
+				tenFiv = self.netFiv(tenFou)
+				tenSix = self.netSix(tenFiv)
+
+				return [ tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix ]
+			# end
+		# end
+
+		class Decoder(torch.nn.Module):
+			def __init__(self, intLevel):
+				super(Decoder, self).__init__()
+
+				intPrevious = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 1]
+				intCurrent = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 0]
+
+				if intLevel < 6: self.netUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1)
+				if intLevel < 6: self.netUpfeat = torch.nn.ConvTranspose2d(in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=4, stride=2, padding=1)
+				if intLevel < 6: self.fltBackwarp = [ None, None, None, 5.0, 2.5, 1.25, 0.625, None ][intLevel + 1]
+
+				self.netOne = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netTwo = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netThr = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128, out_channels=96, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFou = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96, out_channels=64, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFiv = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64, out_channels=32, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netSix = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=3, stride=1, padding=1)
+				)
+			# end
+
+			def forward(self, tenFirst, tenSecond, objPrevious):
+				tenFlow = None
+				tenFeat = None
+
+				if objPrevious is None:
+					tenFlow = None
+					tenFeat = None
+
+					tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=tenSecond), negative_slope=0.1, inplace=False)
+
+					tenFeat = torch.cat([ tenVolume ], 1)
+
+				elif objPrevious is not None:
+					tenFlow = self.netUpflow(objPrevious['tenFlow'])
+					tenFeat = self.netUpfeat(objPrevious['tenFeat'])
+
+					tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=backwarp(tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp)), negative_slope=0.1, inplace=False)
+
+					tenFeat = torch.cat([ tenVolume, tenFirst, tenFlow, tenFeat ], 1)
+
+				# end
+
+				tenFeat = torch.cat([ self.netOne(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netTwo(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netThr(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netFou(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netFiv(tenFeat), tenFeat ], 1)
+
+				tenFlow = self.netSix(tenFeat)
+
+				return {
+					'tenFlow': tenFlow,
+					'tenFeat': tenFeat
+				}
+			# end
+		# end
+
+		class Refiner(torch.nn.Module):
+			def __init__(self):
+				super(Refiner, self).__init__()
+
+				self.netMain = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=2, dilation=2),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=4, dilation=4),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=96, kernel_size=3, stride=1, padding=8, dilation=8),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, stride=1, padding=16, dilation=16),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, dilation=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1, dilation=1)
+				)
+			# end
+
+			def forward(self, tenInput):
+				return self.netMain(tenInput)
+			# end
+		# end
+
+		self.netExtractor = Extractor()
+
+		self.netTwo = Decoder(2)
+		self.netThr = Decoder(3)
+		self.netFou = Decoder(4)
+		self.netFiv = Decoder(5)
+		self.netSix = Decoder(6)
+
+		self.netRefiner = Refiner()
+
+		self.load_state_dict({ strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/github/pytorch-pwc/network-' + 'default' + '.pytorch').items() })
+	# end
+
+	def forward(self, tenFirst, tenSecond):
+		intWidth = tenFirst.shape[3]
+		intHeight = tenFirst.shape[2]
+
+		intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
+		intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
+
+		tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+		tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+		
+		tenFirst = self.netExtractor(tenPreprocessedFirst)
+		tenSecond = self.netExtractor(tenPreprocessedSecond)
+
+
+		objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None)
+		objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate)
+		objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate)
+		objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate)
+		objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate)
+
+		tenFlow = objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat'])
+		tenFlow = 20.0 * torch.nn.functional.interpolate(input=tenFlow, size=(intHeight, intWidth), mode='bilinear', align_corners=False)
+		tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
+		tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
+
+		return tenFlow
+	# end
+# end
+
+netNetwork = None
+
+##########################################################
+
+def estimate(tenFirst, tenSecond):
+	global netNetwork
+
+	if netNetwork is None:
+		netNetwork = Network().cuda().eval()
+	# end
+
+	assert(tenFirst.shape[1] == tenSecond.shape[1])
+	assert(tenFirst.shape[2] == tenSecond.shape[2])
+
+	intWidth = tenFirst.shape[2]
+	intHeight = tenFirst.shape[1]
+
+	assert(intWidth == 1024) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
+	assert(intHeight == 436) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
+
+	tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth)
+	tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth)
+
+	intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
+	intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
+
+	tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+	tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+
+	tenFlow = 20.0 * torch.nn.functional.interpolate(input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), size=(intHeight, intWidth), mode='bilinear', align_corners=False)
+
+	tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
+	tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
+
+	return tenFlow[0, :, :, :].cpu()
+# end
+
+##########################################################
+
+# if __name__ == '__main__':
+# 	tenFirst = torch.FloatTensor(numpy.ascontiguousarray(numpy.array(PIL.Image.open(arguments_strFirst))[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) * (1.0 / 255.0)))
+# 	tenSecond = torch.FloatTensor(numpy.ascontiguousarray(numpy.array(PIL.Image.open(arguments_strSecond))[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) * (1.0 / 255.0)))
+
+# 	tenOutput = estimate(tenFirst, tenSecond)
+
+# 	objOutput = open(arguments_strOut, 'wb')
+
+# 	numpy.array([ 80, 73, 69, 72 ], numpy.uint8).tofile(objOutput)
+# 	numpy.array([ tenOutput.shape[2], tenOutput.shape[1] ], numpy.int32).tofile(objOutput)
+# 	numpy.array(tenOutput.numpy().transpose(1, 2, 0), numpy.float32).tofile(objOutput)
+
+# 	objOutput.close()
+# end
\ No newline at end of file
diff --git a/opensora/eval/flolpips/utils.py b/opensora/eval/flolpips/utils.py
new file mode 100644
index 000000000..98d4502d5
--- /dev/null
+++ b/opensora/eval/flolpips/utils.py
@@ -0,0 +1,95 @@
+import numpy as np
+import cv2
+import torch
+
+
+def normalize_tensor(in_feat,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True))
+    return in_feat/(norm_factor+eps)
+
+def l2(p0, p1, range=255.):
+    return .5*np.mean((p0 / range - p1 / range)**2)
+
+def dssim(p0, p1, range=255.):
+    from skimage.measure import compare_ssim
+    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+def tensor2np(tensor_obj):
+    # change dimension of a tensor object into a numpy array
+    return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))
+
+def np2tensor(np_obj):
+     # change dimenion of np array into tensor array
+    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
+
+def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):
+    # image tensor to lab tensor
+    from skimage import color
+
+    img = tensor2im(image_tensor)
+    img_lab = color.rgb2lab(img)
+    if(mc_only):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+    if(to_norm and not mc_only):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+        img_lab = img_lab/100.
+
+    return np2tensor(img_lab)
+
+def read_frame_yuv2rgb(stream, width, height, iFrame, bit_depth, pix_fmt='420'):
+    if pix_fmt == '420':
+        multiplier = 1
+        uv_factor = 2
+    elif pix_fmt == '444':
+        multiplier = 2
+        uv_factor = 1
+    else:
+        print('Pixel format {} is not supported'.format(pix_fmt))
+        return
+
+    if bit_depth == 8:
+        datatype = np.uint8
+        stream.seek(iFrame*1.5*width*height*multiplier)
+        Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
+        
+        # read chroma samples and upsample since original is 4:2:0 sampling
+        U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+        V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+
+    else:
+        datatype = np.uint16
+        stream.seek(iFrame*3*width*height*multiplier)
+        Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
+                
+        U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+        V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+
+    if pix_fmt == '420':
+        yuv = np.empty((height*3//2, width), dtype=datatype)
+        yuv[0:height,:] = Y
+
+        yuv[height:height+height//4,:] = U.reshape(-1, width)
+        yuv[height+height//4:,:] = V.reshape(-1, width)
+
+        if bit_depth != 8:
+            yuv = (yuv/(2**bit_depth-1)*255).astype(np.uint8)
+
+        #convert to rgb
+        rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_I420)
+    
+    else:
+        yvu = np.stack([Y,V,U],axis=2)
+        if bit_depth != 8:
+            yvu = (yvu/(2**bit_depth-1)*255).astype(np.uint8)
+        rgb = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2RGB)
+
+    return rgb
diff --git a/opensora/eval/fvd/styleganv/fvd.py b/opensora/eval/fvd/styleganv/fvd.py
new file mode 100644
index 000000000..3043a2a4a
--- /dev/null
+++ b/opensora/eval/fvd/styleganv/fvd.py
@@ -0,0 +1,90 @@
+import torch
+import os
+import math
+import torch.nn.functional as F
+
+# https://github.com/universome/fvd-comparison
+
+
+def load_i3d_pretrained(device=torch.device('cpu')):
+    i3D_WEIGHTS_URL = "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt"
+    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_torchscript.pt')
+    print(filepath)
+    if not os.path.exists(filepath):
+        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
+        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+    i3d = torch.jit.load(filepath).eval().to(device)
+    i3d = torch.nn.DataParallel(i3d)
+    return i3d
+    
+
+def get_feats(videos, detector, device, bs=10):
+    # videos : torch.tensor BCTHW [0, 1]
+    detector_kwargs = dict(rescale=False, resize=False, return_features=True) # Return raw features before the softmax layer.
+    feats = np.empty((0, 400))
+    with torch.no_grad():
+        for i in range((len(videos)-1)//bs + 1):
+            feats = np.vstack([feats, detector(torch.stack([preprocess_single(video) for video in videos[i*bs:(i+1)*bs]]).to(device), **detector_kwargs).detach().cpu().numpy()])
+    return feats
+
+
+def get_fvd_feats(videos, i3d, device, bs=10):
+    # videos in [0, 1] as torch tensor BCTHW
+    # videos = [preprocess_single(video) for video in videos]
+    embeddings = get_feats(videos, i3d, device, bs)
+    return embeddings
+
+
+def preprocess_single(video, resolution=224, sequence_length=None):
+    # video: CTHW, [0, 1]
+    c, t, h, w = video.shape
+
+    # temporal crop
+    if sequence_length is not None:
+        assert sequence_length <= t
+        video = video[:, :sequence_length]
+
+    # scale shorter side to resolution
+    scale = resolution / min(h, w)
+    if h < w:
+        target_size = (resolution, math.ceil(w * scale))
+    else:
+        target_size = (math.ceil(h * scale), resolution)
+    video = F.interpolate(video, size=target_size, mode='bilinear', align_corners=False)
+
+    # center crop
+    c, t, h, w = video.shape
+    w_start = (w - resolution) // 2
+    h_start = (h - resolution) // 2
+    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
+
+    # [0, 1] -> [-1, 1]
+    video = (video - 0.5) * 2
+
+    return video.contiguous()
+
+
+"""
+Copy-pasted from https://github.com/cvpr2022-stylegan-v/stylegan-v/blob/main/src/metrics/frechet_video_distance.py
+"""
+from typing import Tuple
+from scipy.linalg import sqrtm
+import numpy as np
+
+
+def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    mu = feats.mean(axis=0) # [d]
+    sigma = np.cov(feats, rowvar=False) # [d, d]
+    return mu, sigma
+
+
+def frechet_distance(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
+    mu_gen, sigma_gen = compute_stats(feats_fake)
+    mu_real, sigma_real = compute_stats(feats_real)
+    m = np.square(mu_gen - mu_real).sum()
+    if feats_fake.shape[0]>1:
+        s, _ = sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
+        fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
+    else:
+        fid = np.real(m)
+    return float(fid)
\ No newline at end of file
diff --git a/opensora/eval/fvd/videogpt/fvd.py b/opensora/eval/fvd/videogpt/fvd.py
new file mode 100644
index 000000000..e81c82904
--- /dev/null
+++ b/opensora/eval/fvd/videogpt/fvd.py
@@ -0,0 +1,137 @@
+import torch
+import os
+import math
+import torch.nn.functional as F
+import numpy as np
+import einops
+
+def load_i3d_pretrained(device=torch.device('cpu')):
+    i3D_WEIGHTS_URL = "https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI"
+    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_pretrained_400.pt')
+    print(filepath)
+    if not os.path.exists(filepath):
+        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
+        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+    from .pytorch_i3d import InceptionI3d
+    i3d = InceptionI3d(400, in_channels=3).eval().to(device)
+    i3d.load_state_dict(torch.load(filepath, map_location=device))
+    i3d = torch.nn.DataParallel(i3d)
+    return i3d
+
+def preprocess_single(video, resolution, sequence_length=None):
+    # video: THWC, {0, ..., 255}
+    video = video.permute(0, 3, 1, 2).float() / 255. # TCHW
+    t, c, h, w = video.shape
+
+    # temporal crop
+    if sequence_length is not None:
+        assert sequence_length <= t
+        video = video[:sequence_length]
+
+    # scale shorter side to resolution
+    scale = resolution / min(h, w)
+    if h < w:
+        target_size = (resolution, math.ceil(w * scale))
+    else:
+        target_size = (math.ceil(h * scale), resolution)
+    video = F.interpolate(video, size=target_size, mode='bilinear',
+                          align_corners=False)
+
+    # center crop
+    t, c, h, w = video.shape
+    w_start = (w - resolution) // 2
+    h_start = (h - resolution) // 2
+    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
+    video = video.permute(1, 0, 2, 3).contiguous() # CTHW
+
+    video -= 0.5
+
+    return video
+
+def preprocess(videos, target_resolution=224):
+    # we should tras videos in [0-1] [b c t h w] as th.float 
+    # -> videos in {0, ..., 255} [b t h w c] as np.uint8 array
+    videos = einops.rearrange(videos, 'b c t h w -> b t h w c')
+    videos = (videos*255).numpy().astype(np.uint8)
+
+    b, t, h, w, c = videos.shape
+    videos = torch.from_numpy(videos)
+    videos = torch.stack([preprocess_single(video, target_resolution) for video in videos])
+    return videos * 2 # [-0.5, 0.5] -> [-1, 1]
+
+def get_fvd_logits(videos, i3d, device, bs=10):
+    videos = preprocess(videos)
+    embeddings = get_logits(i3d, videos, device, bs=10)
+    return embeddings
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L161
+def _symmetric_matrix_square_root(mat, eps=1e-10):
+    u, s, v = torch.svd(mat)
+    si = torch.where(s < eps, s, torch.sqrt(s))
+    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L400
+def trace_sqrt_product(sigma, sigma_v):
+    sqrt_sigma = _symmetric_matrix_square_root(sigma)
+    sqrt_a_sigmav_a = torch.matmul(sqrt_sigma, torch.matmul(sigma_v, sqrt_sigma))
+    return torch.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a))
+
+# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
+def cov(m, rowvar=False):
+    '''Estimate a covariance matrix given data.
+
+    Covariance indicates the level to which two variables vary together.
+    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
+    then the covariance matrix element `C_{ij}` is the covariance of
+    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.
+
+    Args:
+        m: A 1-D or 2-D array containing multiple variables and observations.
+            Each row of `m` represents a variable, and each column a single
+            observation of all those variables.
+        rowvar: If `rowvar` is True, then each row represents a
+            variable, with observations in the columns. Otherwise, the
+            relationship is transposed: each column represents a variable,
+            while the rows contain observations.
+
+    Returns:
+        The covariance matrix of the variables.
+    '''
+    if m.dim() > 2:
+        raise ValueError('m has more than 2 dimensions')
+    if m.dim() < 2:
+        m = m.view(1, -1)
+    if not rowvar and m.size(0) != 1:
+        m = m.t()
+
+    fact = 1.0 / (m.size(1) - 1) # unbiased estimate
+    m -= torch.mean(m, dim=1, keepdim=True)
+    mt = m.t()  # if complex: mt = m.t().conj()
+    return fact * m.matmul(mt).squeeze()
+
+
+def frechet_distance(x1, x2):
+    x1 = x1.flatten(start_dim=1)
+    x2 = x2.flatten(start_dim=1)
+    m, m_w = x1.mean(dim=0), x2.mean(dim=0)
+    sigma, sigma_w = cov(x1, rowvar=False), cov(x2, rowvar=False)
+    mean = torch.sum((m - m_w) ** 2)
+    if x1.shape[0]>1:
+        sqrt_trace_component = trace_sqrt_product(sigma, sigma_w)
+        trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component
+        fd = trace + mean
+    else:
+        fd = np.real(mean)
+    return float(fd)
+
+
+def get_logits(i3d, videos, device, bs=10):
+    # assert videos.shape[0] % 16 == 0
+    with torch.no_grad():
+        logits = []
+        for i in range(0, videos.shape[0], bs):
+            batch = videos[i:i + bs].to(device)
+            # logits.append(i3d.module.extract_features(batch)) # wrong
+            logits.append(i3d(batch)) # right
+        logits = torch.cat(logits, dim=0)
+        return logits
diff --git a/opensora/eval/fvd/videogpt/pytorch_i3d.py b/opensora/eval/fvd/videogpt/pytorch_i3d.py
new file mode 100644
index 000000000..58a16cdd7
--- /dev/null
+++ b/opensora/eval/fvd/videogpt/pytorch_i3d.py
@@ -0,0 +1,322 @@
+# Original code from https://github.com/piergiaj/pytorch-i3d
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class MaxPool3dSamePadding(nn.MaxPool3d):
+
+    def compute_pad(self, dim, s):
+        if s % self.stride[dim] == 0:
+            return max(self.kernel_size[dim] - self.stride[dim], 0)
+        else:
+            return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)
+
+    def forward(self, x):
+        # compute 'same' padding
+        (batch, channel, t, h, w) = x.size()
+        out_t = np.ceil(float(t) / float(self.stride[0]))
+        out_h = np.ceil(float(h) / float(self.stride[1]))
+        out_w = np.ceil(float(w) / float(self.stride[2]))
+        pad_t = self.compute_pad(0, t)
+        pad_h = self.compute_pad(1, h)
+        pad_w = self.compute_pad(2, w)
+
+        pad_t_f = pad_t // 2
+        pad_t_b = pad_t - pad_t_f
+        pad_h_f = pad_h // 2
+        pad_h_b = pad_h - pad_h_f
+        pad_w_f = pad_w // 2
+        pad_w_b = pad_w - pad_w_f
+
+        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
+        x = F.pad(x, pad)
+        return super(MaxPool3dSamePadding, self).forward(x)
+
+
+class Unit3D(nn.Module):
+
+    def __init__(self, in_channels,
+                 output_channels,
+                 kernel_shape=(1, 1, 1),
+                 stride=(1, 1, 1),
+                 padding=0,
+                 activation_fn=F.relu,
+                 use_batch_norm=True,
+                 use_bias=False,
+                 name='unit_3d'):
+
+        """Initializes Unit3D module."""
+        super(Unit3D, self).__init__()
+
+        self._output_channels = output_channels
+        self._kernel_shape = kernel_shape
+        self._stride = stride
+        self._use_batch_norm = use_batch_norm
+        self._activation_fn = activation_fn
+        self._use_bias = use_bias
+        self.name = name
+        self.padding = padding
+
+        self.conv3d = nn.Conv3d(in_channels=in_channels,
+                                out_channels=self._output_channels,
+                                kernel_size=self._kernel_shape,
+                                stride=self._stride,
+                                padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function
+                                bias=self._use_bias)
+
+        if self._use_batch_norm:
+            self.bn = nn.BatchNorm3d(self._output_channels, eps=1e-5, momentum=0.001)
+
+    def compute_pad(self, dim, s):
+        if s % self._stride[dim] == 0:
+            return max(self._kernel_shape[dim] - self._stride[dim], 0)
+        else:
+            return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)
+
+
+    def forward(self, x):
+        # compute 'same' padding
+        (batch, channel, t, h, w) = x.size()
+        out_t = np.ceil(float(t) / float(self._stride[0]))
+        out_h = np.ceil(float(h) / float(self._stride[1]))
+        out_w = np.ceil(float(w) / float(self._stride[2]))
+        pad_t = self.compute_pad(0, t)
+        pad_h = self.compute_pad(1, h)
+        pad_w = self.compute_pad(2, w)
+
+        pad_t_f = pad_t // 2
+        pad_t_b = pad_t - pad_t_f
+        pad_h_f = pad_h // 2
+        pad_h_b = pad_h - pad_h_f
+        pad_w_f = pad_w // 2
+        pad_w_b = pad_w - pad_w_f
+
+        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
+        x = F.pad(x, pad)
+
+        x = self.conv3d(x)
+        if self._use_batch_norm:
+            x = self.bn(x)
+        if self._activation_fn is not None:
+            x = self._activation_fn(x)
+        return x
+
+
+
+class InceptionModule(nn.Module):
+    def __init__(self, in_channels, out_channels, name):
+        super(InceptionModule, self).__init__()
+
+        self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0,
+                         name=name+'/Branch_0/Conv3d_0a_1x1')
+        self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_1/Conv3d_0a_1x1')
+        self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3],
+                          name=name+'/Branch_1/Conv3d_0b_3x3')
+        self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_2/Conv3d_0a_1x1')
+        self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3],
+                          name=name+'/Branch_2/Conv3d_0b_3x3')
+        self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],
+                                stride=(1, 1, 1), padding=0)
+        self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_3/Conv3d_0b_1x1')
+        self.name = name
+
+    def forward(self, x):
+        b0 = self.b0(x)
+        b1 = self.b1b(self.b1a(x))
+        b2 = self.b2b(self.b2a(x))
+        b3 = self.b3b(self.b3a(x))
+        return torch.cat([b0,b1,b2,b3], dim=1)
+
+
+class InceptionI3d(nn.Module):
+    """Inception-v1 I3D architecture.
+    The model is introduced in:
+        Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset
+        Joao Carreira, Andrew Zisserman
+        https://arxiv.org/pdf/1705.07750v1.pdf.
+    See also the Inception architecture, introduced in:
+        Going deeper with convolutions
+        Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,
+        Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.
+        http://arxiv.org/pdf/1409.4842v1.pdf.
+    """
+
+    # Endpoints of the model in order. During construction, all the endpoints up
+    # to a designated `final_endpoint` are returned in a dictionary as the
+    # second return value.
+    VALID_ENDPOINTS = (
+        'Conv3d_1a_7x7',
+        'MaxPool3d_2a_3x3',
+        'Conv3d_2b_1x1',
+        'Conv3d_2c_3x3',
+        'MaxPool3d_3a_3x3',
+        'Mixed_3b',
+        'Mixed_3c',
+        'MaxPool3d_4a_3x3',
+        'Mixed_4b',
+        'Mixed_4c',
+        'Mixed_4d',
+        'Mixed_4e',
+        'Mixed_4f',
+        'MaxPool3d_5a_2x2',
+        'Mixed_5b',
+        'Mixed_5c',
+        'Logits',
+        'Predictions',
+    )
+
+    def __init__(self, num_classes=400, spatial_squeeze=True,
+                 final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5):
+        """Initializes I3D model instance.
+        Args:
+          num_classes: The number of outputs in the logit layer (default 400, which
+              matches the Kinetics dataset).
+          spatial_squeeze: Whether to squeeze the spatial dimensions for the logits
+              before returning (default True).
+          final_endpoint: The model contains many possible endpoints.
+              `final_endpoint` specifies the last endpoint for the model to be built
+              up to. In addition to the output at `final_endpoint`, all the outputs
+              at endpoints up to `final_endpoint` will also be returned, in a
+              dictionary. `final_endpoint` must be one of
+              InceptionI3d.VALID_ENDPOINTS (default 'Logits').
+          name: A string (optional). The name of this module.
+        Raises:
+          ValueError: if `final_endpoint` is not recognized.
+        """
+
+        if final_endpoint not in self.VALID_ENDPOINTS:
+            raise ValueError('Unknown final endpoint %s' % final_endpoint)
+
+        super(InceptionI3d, self).__init__()
+        self._num_classes = num_classes
+        self._spatial_squeeze = spatial_squeeze
+        self._final_endpoint = final_endpoint
+        self.logits = None
+
+        if self._final_endpoint not in self.VALID_ENDPOINTS:
+            raise ValueError('Unknown final endpoint %s' % self._final_endpoint)
+
+        self.end_points = {}
+        end_point = 'Conv3d_1a_7x7'
+        self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7],
+                                            stride=(2, 2, 2), padding=(3,3,3),  name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_2a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Conv3d_2b_1x1'
+        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0,
+                                       name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Conv3d_2c_3x3'
+        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1,
+                                       name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_3a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_3b'
+        self.end_points[end_point] = InceptionModule(192, [64,96,128,16,32,32], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_3c'
+        self.end_points[end_point] = InceptionModule(256, [128,128,192,32,96,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_4a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4b'
+        self.end_points[end_point] = InceptionModule(128+192+96+64, [192,96,208,16,48,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4c'
+        self.end_points[end_point] = InceptionModule(192+208+48+64, [160,112,224,24,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4d'
+        self.end_points[end_point] = InceptionModule(160+224+64+64, [128,128,256,24,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4e'
+        self.end_points[end_point] = InceptionModule(128+256+64+64, [112,144,288,32,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4f'
+        self.end_points[end_point] = InceptionModule(112+288+64+64, [256,160,320,32,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_5a_2x2'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_5b'
+        self.end_points[end_point] = InceptionModule(256+320+128+128, [256,160,320,32,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_5c'
+        self.end_points[end_point] = InceptionModule(256+320+128+128, [384,192,384,48,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Logits'
+        self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7],
+                                     stride=(1, 1, 1))
+        self.dropout = nn.Dropout(dropout_keep_prob)
+        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
+                             kernel_shape=[1, 1, 1],
+                             padding=0,
+                             activation_fn=None,
+                             use_batch_norm=False,
+                             use_bias=True,
+                             name='logits')
+
+        self.build()
+
+
+    def replace_logits(self, num_classes):
+        self._num_classes = num_classes
+        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
+                             kernel_shape=[1, 1, 1],
+                             padding=0,
+                             activation_fn=None,
+                             use_batch_norm=False,
+                             use_bias=True,
+                             name='logits')
+
+
+    def build(self):
+        for k in self.end_points.keys():
+            self.add_module(k, self.end_points[k])
+
+    def forward(self, x):
+        for end_point in self.VALID_ENDPOINTS:
+            if end_point in self.end_points:
+                x = self._modules[end_point](x) # use _modules to work with dataparallel
+
+        x = self.logits(self.dropout(self.avg_pool(x)))
+        if self._spatial_squeeze:
+            logits = x.squeeze(3).squeeze(3)
+        logits = logits.mean(dim=2)
+        # logits is batch X time X classes, which is what we want to work with
+        return logits
+
+
+    def extract_features(self, x):
+        for end_point in self.VALID_ENDPOINTS:
+            if end_point in self.end_points:
+                x = self._modules[end_point](x)
+        return self.avg_pool(x)
\ No newline at end of file
diff --git a/opensora/eval/script/cal_clip_score.sh b/opensora/eval/script/cal_clip_score.sh
new file mode 100644
index 000000000..8bc241c15
--- /dev/null
+++ b/opensora/eval/script/cal_clip_score.sh
@@ -0,0 +1,23 @@
+# clip_score cross modality
+python eval_clip_score.py \
+    --real_path path/to/image \
+    --generated_path path/to/text \
+    --batch-size 50 \
+    --device "cuda"
+
+# clip_score within the same modality
+python eval_clip_score.py \
+    --real_path path/to/textA \
+    --generated_path path/to/textB \
+    --real_flag txt \
+    --generated_flag txt \
+    --batch-size 50 \
+    --device "cuda"
+
+python eval_clip_score.py \
+    --real_path path/to/imageA \
+    --generated_path path/to/imageB \
+    --real_flag img \
+    --generated_flag img \
+    --batch-size 50 \
+    --device "cuda"
diff --git a/opensora/eval/script/cal_fvd.sh b/opensora/eval/script/cal_fvd.sh
new file mode 100644
index 000000000..e344242a0
--- /dev/null
+++ b/opensora/eval/script/cal_fvd.sh
@@ -0,0 +1,9 @@
+python eval_common_metric.py \
+    --real_video_dir path/to/imageA\
+    --generated_video_dir path/to/imageB \
+    --batch_size 10 \
+    --crop_size 64 \
+    --num_frames 20 \
+    --device 'cuda' \
+    --metric 'fvd' \
+    --fvd_method 'styleganv'
diff --git a/opensora/eval/script/cal_lpips.sh b/opensora/eval/script/cal_lpips.sh
new file mode 100644
index 000000000..9ebaf26e8
--- /dev/null
+++ b/opensora/eval/script/cal_lpips.sh
@@ -0,0 +1,8 @@
+python eval_common_metric.py \
+    --real_video_dir path/to/imageA\
+    --generated_video_dir path/to/imageB \
+    --batch_size 10 \
+    --num_frames 20 \
+    --crop_size 64 \
+    --device 'cuda' \
+    --metric 'lpips'
\ No newline at end of file
diff --git a/opensora/eval/script/cal_psnr.sh b/opensora/eval/script/cal_psnr.sh
new file mode 100644
index 000000000..a60a5d1c1
--- /dev/null
+++ b/opensora/eval/script/cal_psnr.sh
@@ -0,0 +1,9 @@
+
+python eval_common_metric.py \
+    --real_video_dir /data/xiaogeng_liu/data/video1 \
+    --generated_video_dir /data/xiaogeng_liu/data/video2 \
+    --batch_size 10 \
+    --num_frames 20 \
+    --crop_size 64 \
+    --device 'cuda' \
+    --metric 'psnr'
\ No newline at end of file
diff --git a/opensora/eval/script/cal_ssim.sh b/opensora/eval/script/cal_ssim.sh
new file mode 100644
index 000000000..404d8da9a
--- /dev/null
+++ b/opensora/eval/script/cal_ssim.sh
@@ -0,0 +1,8 @@
+python eval_common_metric.py \
+    --real_video_dir /data/xiaogeng_liu/data/video1 \
+    --generated_video_dir /data/xiaogeng_liu/data/video2 \
+    --batch_size 10 \
+    --num_frames 20 \
+    --crop_size 64 \
+    --device 'cuda' \
+    --metric 'ssim'
\ No newline at end of file
diff --git a/opensora/models/__init__.py b/opensora/models/__init__.py
new file mode 100644
index 000000000..3eca32c73
--- /dev/null
+++ b/opensora/models/__init__.py
@@ -0,0 +1 @@
+from .causalvideovae.model import CausalVAEModelWrapper
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/__init__.py b/opensora/models/causalvideovae/__init__.py
new file mode 100644
index 000000000..bc1357393
--- /dev/null
+++ b/opensora/models/causalvideovae/__init__.py
@@ -0,0 +1,32 @@
+from torchvision.transforms import Lambda
+from .model.causal_vae import CausalVAEModelWrapper
+
+ae_stride_config = {
+    'CausalVAEModel_D4_2x8x8': [2, 8, 8],
+    'CausalVAEModel_D8_2x8x8': [2, 8, 8],
+    'CausalVAEModel_D4_4x8x8': [4, 8, 8],
+    'CausalVAEModel_D8_4x8x8': [4, 8, 8],
+}
+
+
+ae_channel_config = {
+    'CausalVAEModel_D4_2x8x8': 4,
+    'CausalVAEModel_D8_2x8x8': 8,
+    'CausalVAEModel_D4_4x8x8': 4,
+    'CausalVAEModel_D8_4x8x8': 8,
+}
+
+
+ae_denorm = {
+    'CausalVAEModel_D4_2x8x8': lambda x: (x + 1.) / 2.,
+    'CausalVAEModel_D8_2x8x8': lambda x: (x + 1.) / 2.,
+    'CausalVAEModel_D4_4x8x8': lambda x: (x + 1.) / 2.,
+    'CausalVAEModel_D8_4x8x8': lambda x: (x + 1.) / 2.,
+}
+
+ae_norm = {
+    'CausalVAEModel_D4_2x8x8': Lambda(lambda x: 2. * x - 1.),
+    'CausalVAEModel_D8_2x8x8': Lambda(lambda x: 2. * x - 1.),
+    'CausalVAEModel_D4_4x8x8': Lambda(lambda x: 2. * x - 1.),
+    'CausalVAEModel_D8_4x8x8': Lambda(lambda x: 2. * x - 1.),
+}
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/dataset/__init__.py b/opensora/models/causalvideovae/dataset/__init__.py
new file mode 100755
index 000000000..cb0d0f041
--- /dev/null
+++ b/opensora/models/causalvideovae/dataset/__init__.py
@@ -0,0 +1,65 @@
+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 .t2v_datasets import T2V_dataset
+from .transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVideo, CenterCropResizeVideo, LongSideResizeVideo, SpatialStrideCropVideo
+
+
+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 == 't2v':
+        if args.multi_scale:
+            resize = [
+                LongSideResizeVideo(args.max_image_size, skip_low_resolution=True),
+                SpatialStrideCropVideo(args.stride)
+                ]
+        else:
+            resize = [CenterCropResizeVideo(args.max_image_size), ]
+        transform = transforms.Compose([
+            ToTensorVideo(),
+            *resize, 
+            # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+            norm_fun
+        ])
+        tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+        return T2V_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer)
+    raise NotImplementedError(args.dataset)
diff --git a/opensora/models/causalvideovae/dataset/feature_datasets.py b/opensora/models/causalvideovae/dataset/feature_datasets.py
new file mode 100755
index 000000000..e8f85f2c1
--- /dev/null
+++ b/opensora/models/causalvideovae/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 .transform import center_crop, RandomCropVideo
+from ..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
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/dataset/t2v_datasets.py b/opensora/models/causalvideovae/dataset/t2v_datasets.py
new file mode 100755
index 000000000..776f15fff
--- /dev/null
+++ b/opensora/models/causalvideovae/dataset/t2v_datasets.py
@@ -0,0 +1,203 @@
+import json
+import os, io, csv, math, random
+import numpy as np
+import torchvision
+from einops import rearrange
+from decord import VideoReader
+from os.path import join as opj
+import gc
+import torch
+import torchvision.transforms as transforms
+from torch.utils.data.dataset import Dataset
+from tqdm import tqdm
+from PIL import Image
+
+from ..utils.dataset_utils import DecordInit
+from ..utils.utils import text_preprocessing
+
+
+def random_video_noise(t, c, h, w):
+    vid = torch.rand(t, c, h, w) * 255.0
+    vid = vid.to(torch.uint8)
+    return vid
+
+class T2V_dataset(Dataset):
+    def __init__(self, args, transform, temporal_sample, tokenizer):
+        self.image_data = args.image_data
+        self.video_data = args.video_data
+        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()
+
+        self.vid_cap_list = self.get_vid_cap_list()
+        
+        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.vid_cap_list)
+
+    def __getitem__(self, idx):
+        try:
+        # import ipdb;ipdb.set_trace()
+            video_data = self.get_video(idx)
+            image_data = {}
+            if self.use_image_num != 0 and self.use_img_from_vid:
+                image_data = self.get_image_from_video(video_data)
+            elif self.use_image_num != 0 and not self.use_img_from_vid:
+                image_data = self.get_image(idx)
+            else:
+                raise NotImplementedError
+            gc.collect()
+            return dict(video_data=video_data, image_data=image_data)
+        except Exception as e:
+            # print(f'Error with {e}, {self.vid_cap_list[idx]}')
+            if os.path.exists(self.vid_cap_list[idx]['path']) and '_resize_1080p' in self.vid_cap_list[idx]['path']:
+                os.remove(self.vid_cap_list[idx]['path'])
+                print('remove:', self.vid_cap_list[idx]['path'])
+            return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+    def get_video(self, idx):
+        # video = random.choice([random_video_noise(65, 3, 720, 360) * 255, random_video_noise(65, 3, 1024, 1024), random_video_noise(65, 3, 360, 720)])
+        # # print('random shape', video.shape)
+        # 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)
+        
+        video_path = self.vid_cap_list[idx]['path']
+        frame_idx = self.vid_cap_list[idx]['frame_idx']
+        #print('before decord')
+        video = self.decord_read(video_path, frame_idx)
+        # video = self.tv_read(video_path, frame_idx)
+        #print('after decord')
+        video = self.transform(video)  # T C H W -> T C H W
+        # del raw_video
+        # gc.collect()
+        # video = torch.rand(65, 3, 512, 512)
+        #print('after transform')
+        video = video.transpose(0, 1)  # T C H W -> C T H W
+        text = self.vid_cap_list[idx]['cap']
+
+        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']
+        cond_mask = text_tokens_and_mask['attention_mask']
+        return dict(video=video, input_ids=input_ids, cond_mask=cond_mask)
+
+    def get_image_from_video(self, video_data):
+        select_image_idx = np.linspace(0, self.num_frames-1, self.use_image_num, dtype=int)
+        assert self.num_frames >= self.use_image_num
+        image = [video_data['video'][:, i:i+1] for i in select_image_idx]  # num_img [c, 1, h, w]
+        input_ids = video_data['input_ids'].repeat(self.use_image_num, 1)  # self.use_image_num, l
+        cond_mask = video_data['cond_mask'].repeat(self.use_image_num, 1)  # self.use_image_num, l
+        return dict(image=image, input_ids=input_ids, cond_mask=cond_mask)
+
+    def get_image(self, idx):
+        idx = idx % len(self.img_cap_list)  # out of range
+        image_data = self.img_cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
+        
+        image = [Image.open(i['path']).convert('RGB') for i in image_data] # num_img [h, w, c]
+        image = [torch.from_numpy(np.array(i)) for i in image] # num_img [h, w, c]
+        image = [rearrange(i, 'h w c -> c h w').unsqueeze(0) for i in image] # num_img [1 c h w]
+        image = [self.transform(i) for i in image]  # num_img [1 C H W] -> num_img [1 C H W]
+        image = [i.transpose(0, 1) for i in image]  # num_img [1 C H W] -> num_img [C 1 H W]
+
+        caps = [i['cap'] for i in image_data]
+        text = [text_preprocessing(cap) for cap in caps]
+        input_ids, cond_mask = [], []
+        for t in text:
+            text_tokens_and_mask = self.tokenizer(
+                t,
+                max_length=self.model_max_length,
+                padding='max_length',
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors='pt'
+            )
+            input_ids.append(text_tokens_and_mask['input_ids'])
+            cond_mask.append(text_tokens_and_mask['attention_mask'])
+        input_ids = torch.cat(input_ids)  # self.use_image_num, l
+        cond_mask = torch.cat(cond_mask)  # self.use_image_num, l
+        return dict(image=image, input_ids=input_ids, cond_mask=cond_mask)
+
+    def tv_read(self, path, frame_idx=None):
+        vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+        total_frames = len(vframes)
+        if frame_idx is None:
+            start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        else:
+            start_frame_ind, end_frame_ind = frame_idx.split(':')
+            start_frame_ind, end_frame_ind = int(start_frame_ind), int(end_frame_ind)
+        # 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)
+        # frame_indice = np.linspace(0, 63, self.num_frames, dtype=int)
+
+        video = vframes[frame_indice]  # (T, C, H, W)
+
+        return video
+
+    def decord_read(self, path, frame_idx=None):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        # Sampling video frames
+        if frame_idx is None:
+            start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        else:
+            start_frame_ind, end_frame_ind = frame_idx.split(':')
+            start_frame_ind, end_frame_ind = int(start_frame_ind), int(end_frame_ind)
+            start_frame_ind, end_frame_ind = int(start_frame_ind), int(start_frame_ind) + self.num_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)
+        # frame_indice = np.linspace(0, 63, 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_vid_cap_list(self):
+        vid_cap_lists = []
+        with open(self.video_data, 'r') as f:
+            folder_anno = [i.strip().split(',') for i in f.readlines() if len(i.strip()) > 0]
+            # print(folder_anno)
+        for folder, anno in folder_anno:
+            with open(anno, 'r') as f:
+                vid_cap_list = json.load(f)
+            print(f'Building {anno}...')
+            for i in tqdm(range(len(vid_cap_list))):
+                path = opj(folder, vid_cap_list[i]['path'])
+                if os.path.exists(path.replace('.mp4', '_resize_1080p.mp4')):
+                    path = path.replace('.mp4', '_resize_1080p.mp4')
+                vid_cap_list[i]['path'] = path
+
+            vid_cap_lists += vid_cap_list
+        return vid_cap_lists
+
+    def get_img_cap_list(self):
+        img_cap_lists = []
+        with open(self.image_data, 'r') as f:
+            folder_anno = [i.strip().split(',') for i in f.readlines() if len(i.strip()) > 0]
+        for folder, anno in folder_anno:
+            with open(anno, 'r') as f:
+                img_cap_list = json.load(f)
+            print(f'Building {anno}...')
+            for i in tqdm(range(len(img_cap_list))):
+                img_cap_list[i]['path'] = opj(folder, img_cap_list[i]['path'])
+            img_cap_lists += img_cap_list
+        img_cap_lists = [img_cap_lists[i: i+self.use_image_num] for i in range(0, len(img_cap_lists), self.use_image_num)]
+        return img_cap_lists[:-1]  # drop last to avoid error length
diff --git a/opensora/models/causalvideovae/dataset/transform.py b/opensora/models/causalvideovae/dataset/transform.py
new file mode 100755
index 000000000..bb89c2c85
--- /dev/null
+++ b/opensora/models/causalvideovae/dataset/transform.py
@@ -0,0 +1,573 @@
+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 SpatialStrideCropVideo:
+    def __init__(self, stride):
+            self.stride = stride
+
+    def __call__(self, clip):
+        """
+        Args:
+            clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)
+        Returns:
+            torch.tensor: cropped video clip by stride.
+                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 = h // self.stride * self.stride, w // self.stride * self.stride
+
+        return 0, 0, th, tw  # from top-left
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size})"
+
+class LongSideResizeVideo:
+    '''
+    First use the long side,
+    then resize to the specified size
+    '''
+
+    def __init__(
+            self,
+            size,
+            skip_low_resolution=False, 
+            interpolation_mode="bilinear",
+    ):
+        self.size = size
+        self.skip_low_resolution = skip_low_resolution
+        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 video clip.
+                size is (T, C, 512, *) or (T, C, *, 512)
+        """
+        _, _, h, w = clip.shape
+        if self.skip_low_resolution and max(h, w) <= self.size:
+            return clip
+        if h > w:
+            w = int(w * self.size / h)
+            h = self.size
+        else:
+            h = int(h * self.size / w)
+            w = self.size
+        resize_clip = resize(clip, target_size=(h, w),
+                                         interpolation_mode=self.interpolation_mode)
+        return resize_clip
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
+
+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
+
+class DynamicSampleDuration(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, t_stride, extra_1):
+        self.t_stride = t_stride
+        self.extra_1 = extra_1
+
+    def __call__(self, t, h, w):
+        if self.extra_1:
+            t = t - 1
+        truncate_t_list = list(range(t+1))[t//2:][::self.t_stride]  # need half at least
+        truncate_t = random.choice(truncate_t_list)
+        if self.extra_1:
+            truncate_t = truncate_t + 1
+        return 0, truncate_t
+
+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))
diff --git a/opensora/models/causalvideovae/dataset/ucf101.py b/opensora/models/causalvideovae/dataset/ucf101.py
new file mode 100755
index 000000000..6e982a76c
--- /dev/null
+++ b/opensora/models/causalvideovae/dataset/ucf101.py
@@ -0,0 +1,80 @@
+import math
+import os
+
+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 ..utils.dataset_utils import DecordInit
+
+
+class UCF101(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.classes = sorted(os.listdir(self.data_path))
+        self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)}
+        self.samples = self._make_dataset()
+
+
+    def _make_dataset(self):
+        dataset = []
+        for class_name in self.classes:
+            class_path = os.path.join(self.data_path, class_name)
+            for fname in os.listdir(class_path):
+                if fname.endswith('.avi'):
+                    item = (os.path.join(class_path, fname), self.class_to_idx[class_name])
+                    dataset.append(item)
+        return dataset
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        video_path, label = 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
+            return video, label
+        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
+
diff --git a/opensora/models/causalvideovae/eval/cal_flolpips.py b/opensora/models/causalvideovae/eval/cal_flolpips.py
new file mode 100755
index 000000000..15cdd771e
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/cal_flolpips.py
@@ -0,0 +1,83 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+from einops import rearrange
+import sys
+sys.path.append(".")
+from flolpips.pwcnet import Network as PWCNet
+from flolpips.flolpips import FloLPIPS
+
+loss_fn = FloLPIPS(net='alex', version='0.1').eval().requires_grad_(False)
+flownet = PWCNet().eval().requires_grad_(False)
+
+def trans(x):
+    return x
+
+
+def calculate_flolpips(videos1, videos2, device):
+    global loss_fn, flownet
+    
+    print("calculate_flowlpips...")
+    loss_fn = loss_fn.to(device)
+    flownet = flownet.to(device)
+    
+    if videos1.shape != videos2.shape:
+        print("Warning: the shape of videos are not equal.")
+        min_frames = min(videos1.shape[1], videos2.shape[1])
+        videos1 = videos1[:, :min_frames]
+        videos2 = videos2[:, :min_frames]
+        
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    flolpips_results = []
+    for video_num in tqdm(range(videos1.shape[0])):
+        video1 = videos1[video_num].to(device)
+        video2 = videos2[video_num].to(device)
+        frames_rec = video1[:-1]
+        frames_rec_next = video1[1:]
+        frames_gt = video2[:-1]
+        frames_gt_next = video2[1:]
+        t, c, h, w = frames_gt.shape
+        flow_gt = flownet(frames_gt, frames_gt_next)
+        flow_dis = flownet(frames_rec, frames_rec_next)
+        flow_diff = flow_gt - flow_dis
+        flolpips = loss_fn.forward(frames_gt, frames_rec, flow_diff, normalize=True)
+        flolpips_results.append(flolpips.cpu().numpy().tolist())
+        
+    flolpips_results = np.array(flolpips_results) # [batch_size, num_frames]
+    flolpips = {}
+    flolpips_std = {}
+
+    for clip_timestamp in range(flolpips_results.shape[1]):
+        flolpips[clip_timestamp] = np.mean(flolpips_results[:,clip_timestamp], axis=-1)
+        flolpips_std[clip_timestamp] = np.std(flolpips_results[:,clip_timestamp], axis=-1)
+
+    result = {
+        "value": flolpips,
+        "value_std": flolpips_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+        "result": flolpips_results,
+        "details": flolpips_results.tolist()
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+
+    import json
+    result = calculate_flolpips(videos1, videos2, "cuda:0")
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/cal_fvd.py b/opensora/models/causalvideovae/eval/cal_fvd.py
new file mode 100755
index 000000000..1f1a9806e
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/cal_fvd.py
@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+
+def trans(x):
+    # if greyscale images add channel
+    if x.shape[-3] == 1:
+        x = x.repeat(1, 1, 3, 1, 1)
+
+    # permute BTCHW -> BCTHW
+    x = x.permute(0, 2, 1, 3, 4) 
+
+    return x
+
+def calculate_fvd(videos1, videos2, device, method='styleganv'):
+
+    if method == 'styleganv':
+        from fvd.styleganv.fvd import get_fvd_feats, frechet_distance, load_i3d_pretrained
+    elif method == 'videogpt':
+        from fvd.videogpt.fvd import load_i3d_pretrained
+        from fvd.videogpt.fvd import get_fvd_logits as get_fvd_feats
+        from fvd.videogpt.fvd import frechet_distance
+
+    print("calculate_fvd...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    i3d = load_i3d_pretrained(device=device)
+    fvd_results = []
+
+    # support grayscale input, if grayscale -> channel*3
+    # BTCHW -> BCTHW
+    # videos -> [batch_size, channel, timestamps, h, w]
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    fvd_results = {}
+
+    # for calculate FVD, each clip_timestamp must >= 10
+    for clip_timestamp in tqdm(range(10, videos1.shape[-3]+1)):
+       
+        # get a video clip
+        # videos_clip [batch_size, channel, timestamps[:clip], h, w]
+        videos_clip1 = videos1[:, :, : clip_timestamp]
+        videos_clip2 = videos2[:, :, : clip_timestamp]
+
+        # get FVD features
+        feats1 = get_fvd_feats(videos_clip1, i3d=i3d, device=device)
+        feats2 = get_fvd_feats(videos_clip2, i3d=i3d, device=device)
+      
+        # calculate FVD when timestamps[:clip]
+        fvd_results[clip_timestamp] = frechet_distance(feats1, feats2)
+
+    result = {
+        "value": fvd_results,
+        "video_setting": videos1.shape,
+        "video_setting_name": "batch_size, channel, time, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+    # device = torch.device("cpu")
+
+    import json
+    result = calculate_fvd(videos1, videos2, device, method='videogpt')
+    print(json.dumps(result, indent=4))
+
+    result = calculate_fvd(videos1, videos2, device, method='styleganv')
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
diff --git a/opensora/models/causalvideovae/eval/cal_lpips.py b/opensora/models/causalvideovae/eval/cal_lpips.py
new file mode 100755
index 000000000..10edc610b
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/cal_lpips.py
@@ -0,0 +1,97 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+
+import torch
+import lpips
+
+spatial = True         # Return a spatial map of perceptual distance.
+
+# Linearly calibrated models (LPIPS)
+loss_fn = lpips.LPIPS(net='alex', spatial=spatial) # Can also set net = 'squeeze' or 'vgg'
+# loss_fn = lpips.LPIPS(net='alex', spatial=spatial, lpips=False) # Can also set net = 'squeeze' or 'vgg'
+
+def trans(x):
+    # if greyscale images add channel
+    if x.shape[-3] == 1:
+        x = x.repeat(1, 1, 3, 1, 1)
+
+    # value range [0, 1] -> [-1, 1]
+    x = x * 2 - 1
+
+    return x
+
+def calculate_lpips(videos1, videos2, device):
+    # image should be RGB, IMPORTANT: normalized to [-1,1]
+    print("calculate_lpips...")
+
+    assert videos1.shape == videos2.shape
+
+    # videos [batch_size, timestamps, channel, h, w]
+
+    # support grayscale input, if grayscale -> channel*3
+    # value range [0, 1] -> [-1, 1]
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    lpips_results = []
+
+    for video_num in tqdm(range(videos1.shape[0])):
+        # get a video
+        # video [timestamps, channel, h, w]
+        video1 = videos1[video_num]
+        video2 = videos2[video_num]
+
+        lpips_results_of_a_video = []
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] tensor
+
+            img1 = video1[clip_timestamp].unsqueeze(0).to(device)
+            img2 = video2[clip_timestamp].unsqueeze(0).to(device)
+            
+            loss_fn.to(device)
+
+            # calculate lpips of a video
+            lpips_results_of_a_video.append(loss_fn.forward(img1, img2).mean().detach().cpu().tolist())
+        lpips_results.append(lpips_results_of_a_video)
+    
+    lpips_results = np.array(lpips_results)
+    
+    lpips = {}
+    lpips_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        lpips[clip_timestamp] = np.mean(lpips_results[:,clip_timestamp])
+        lpips_std[clip_timestamp] = np.std(lpips_results[:,clip_timestamp])
+
+
+    result = {
+        "value": lpips,
+        "value_std": lpips_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+    # device = torch.device("cpu")
+
+    import json
+    result = calculate_lpips(videos1, videos2, device)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/cal_psnr.py b/opensora/models/causalvideovae/eval/cal_psnr.py
new file mode 100755
index 000000000..08d6b94db
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/cal_psnr.py
@@ -0,0 +1,97 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import math
+
+def img_psnr_cuda(img1, img2):
+    # [0,1]
+    # compute mse
+    # mse = np.mean((img1-img2)**2)
+    mse = torch.mean((img1 / 1.0 - img2 / 1.0) ** 2)
+    # compute psnr
+    if mse < 1e-10:
+        return 100
+    psnr = 20 * torch.log10(1 / torch.sqrt(mse))
+    return psnr
+
+
+def img_psnr(img1, img2):
+    # [0,1]
+    # compute mse
+    # mse = np.mean((img1-img2)**2)
+    mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2)
+    # compute psnr
+    if mse < 1e-10:
+        return 100
+    psnr = 20 * math.log10(1 / math.sqrt(mse))
+    return psnr
+
+
+def trans(x):
+    return x
+
+def calculate_psnr(videos1, videos2):
+    print("calculate_psnr...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    psnr_results = []
+    
+    for video_num in tqdm(range(videos1.shape[0])):
+        # get a video
+        # video [timestamps, channel, h, w]
+        video1 = videos1[video_num]
+        video2 = videos2[video_num]
+
+        psnr_results_of_a_video = []
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] numpy
+
+            img1 = video1[clip_timestamp].numpy()
+            img2 = video2[clip_timestamp].numpy()
+            
+            # calculate psnr of a video
+            psnr_results_of_a_video.append(img_psnr(img1, img2))
+
+        psnr_results.append(psnr_results_of_a_video)
+    
+    psnr_results = np.array(psnr_results) # [batch_size, num_frames]
+    psnr = {}
+    psnr_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        psnr[clip_timestamp] = np.mean(psnr_results[:,clip_timestamp])
+        psnr_std[clip_timestamp] = np.std(psnr_results[:,clip_timestamp])
+
+    result = {
+        "value": psnr,
+        "value_std": psnr_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+
+    import json
+    result = calculate_psnr(videos1, videos2)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/cal_ssim.py b/opensora/models/causalvideovae/eval/cal_ssim.py
new file mode 100755
index 000000000..b2de811b6
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/cal_ssim.py
@@ -0,0 +1,113 @@
+import numpy as np
+import torch
+from tqdm import tqdm
+import cv2
+ 
+def ssim(img1, img2):
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+    img1 = img1.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1 ** 2
+    mu2_sq = mu2 ** 2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+                                                            (sigma1_sq + sigma2_sq + C2))
+    return ssim_map.mean()
+ 
+ 
+def calculate_ssim_function(img1, img2):
+    # [0,1]
+    # ssim is the only metric extremely sensitive to gray being compared to b/w 
+    if not img1.shape == img2.shape:
+        raise ValueError('Input images must have the same dimensions.')
+    if img1.ndim == 2:
+        return ssim(img1, img2)
+    elif img1.ndim == 3:
+        if img1.shape[0] == 3:
+            ssims = []
+            for i in range(3):
+                ssims.append(ssim(img1[i], img2[i]))
+            return np.array(ssims).mean()                   
+        elif img1.shape[0] == 1:
+            return ssim(np.squeeze(img1), np.squeeze(img2))
+    else:
+        raise ValueError('Wrong input image dimensions.')
+
+def trans(x):
+    return x
+
+def calculate_ssim(videos1, videos2):
+    print("calculate_ssim...")
+
+    # videos [batch_size, timestamps, channel, h, w]
+    
+    assert videos1.shape == videos2.shape
+
+    videos1 = trans(videos1)
+    videos2 = trans(videos2)
+
+    ssim_results = []
+    
+    for video_num in tqdm(range(videos1.shape[0])):
+        # get a video
+        # video [timestamps, channel, h, w]
+        video1 = videos1[video_num]
+        video2 = videos2[video_num]
+
+        ssim_results_of_a_video = []
+        for clip_timestamp in range(len(video1)):
+            # get a img
+            # img [timestamps[x], channel, h, w]
+            # img [channel, h, w] numpy
+
+            img1 = video1[clip_timestamp].numpy()
+            img2 = video2[clip_timestamp].numpy()
+            
+            # calculate ssim of a video
+            ssim_results_of_a_video.append(calculate_ssim_function(img1, img2))
+
+        ssim_results.append(ssim_results_of_a_video)
+
+    ssim_results = np.array(ssim_results)
+
+    ssim = {}
+    ssim_std = {}
+
+    for clip_timestamp in range(len(video1)):
+        ssim[clip_timestamp] = np.mean(ssim_results[:,clip_timestamp])
+        ssim_std[clip_timestamp] = np.std(ssim_results[:,clip_timestamp])
+
+    result = {
+        "value": ssim,
+        "value_std": ssim_std,
+        "video_setting": video1.shape,
+        "video_setting_name": "time, channel, heigth, width",
+    }
+
+    return result
+
+# test code / using example
+
+def main():
+    NUMBER_OF_VIDEOS = 8
+    VIDEO_LENGTH = 50
+    CHANNEL = 3
+    SIZE = 64
+    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
+    device = torch.device("cuda")
+
+    import json
+    result = calculate_ssim(videos1, videos2)
+    print(json.dumps(result, indent=4))
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/eval_clip_score.py b/opensora/models/causalvideovae/eval/eval_clip_score.py
new file mode 100755
index 000000000..ed2f51082
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/eval_clip_score.py
@@ -0,0 +1,225 @@
+"""Calculates the CLIP Scores
+
+The CLIP model is a contrasitively learned language-image model. There is
+an image encoder and a text encoder. It is believed that the CLIP model could 
+measure the similarity of cross modalities. Please find more information from 
+https://github.com/openai/CLIP.
+
+The CLIP Score measures the Cosine Similarity between two embedded features.
+This repository utilizes the pretrained CLIP Model to calculate 
+the mean average of cosine similarities. 
+
+See --help to see further details.
+
+Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.
+
+Copyright 2023 The Hong Kong Polytechnic University
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+import os.path as osp
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+
+import clip
+import torch
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+
+IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
+                    'tif', 'tiff', 'webp'}
+
+TEXT_EXTENSIONS = {'txt'}
+
+
+class DummyDataset(Dataset):
+    
+    FLAGS = ['img', 'txt']
+    def __init__(self, real_path, generated_path,
+                 real_flag: str = 'img',
+                 generated_flag: str = 'img',
+                 transform = None,
+                 tokenizer = None) -> None:
+        super().__init__()
+        assert real_flag in self.FLAGS and generated_flag in self.FLAGS, \
+            'CLIP Score only support modality of {}. However, get {} and {}'.format(
+                self.FLAGS, real_flag, generated_flag
+            )
+        self.real_folder = self._combine_without_prefix(real_path)
+        self.real_flag = real_flag
+        self.fake_foler = self._combine_without_prefix(generated_path)
+        self.generated_flag = generated_flag
+        self.transform = transform
+        self.tokenizer = tokenizer
+        # assert self._check()
+
+    def __len__(self):
+        return len(self.real_folder)
+
+    def __getitem__(self, index):
+        if index >= len(self):
+            raise IndexError
+        real_path = self.real_folder[index]
+        generated_path = self.fake_foler[index]
+        real_data = self._load_modality(real_path, self.real_flag)
+        fake_data = self._load_modality(generated_path, self.generated_flag)
+
+        sample = dict(real=real_data, fake=fake_data)
+        return sample
+    
+    def _load_modality(self, path, modality):
+        if modality == 'img':
+            data = self._load_img(path)
+        elif modality == 'txt':
+            data = self._load_txt(path)
+        else:
+            raise TypeError("Got unexpected modality: {}".format(modality))
+        return data
+
+    def _load_img(self, path):
+        img = Image.open(path)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img
+    
+    def _load_txt(self, path):
+        with open(path, 'r') as fp:
+            data = fp.read()
+            fp.close()
+        if self.tokenizer is not None:
+            data = self.tokenizer(data).squeeze()
+        return data
+
+    def _check(self):
+        for idx in range(len(self)):
+            real_name = self.real_folder[idx].split('.')
+            fake_name = self.fake_folder[idx].split('.')
+            if fake_name != real_name:
+                return False
+        return True
+    
+    def _combine_without_prefix(self, folder_path, prefix='.'):
+        folder = []
+        for name in os.listdir(folder_path):
+            if name[0] == prefix:
+                continue
+            folder.append(osp.join(folder_path, name))
+        folder.sort()
+        return folder
+
+
+@torch.no_grad()
+def calculate_clip_score(dataloader, model, real_flag, generated_flag):
+    score_acc = 0.
+    sample_num = 0.
+    logit_scale = model.logit_scale.exp()
+    for batch_data in tqdm(dataloader):
+        real = batch_data['real']
+        real_features = forward_modality(model, real, real_flag)
+        fake = batch_data['fake']
+        fake_features = forward_modality(model, fake, generated_flag)
+        
+        # normalize features
+        real_features = real_features / real_features.norm(dim=1, keepdim=True).to(torch.float32)
+        fake_features = fake_features / fake_features.norm(dim=1, keepdim=True).to(torch.float32)
+        
+        # calculate scores
+        # score = logit_scale * real_features @ fake_features.t()
+        # score_acc += torch.diag(score).sum()
+        score = logit_scale * (fake_features * real_features).sum()
+        score_acc += score
+        sample_num += real.shape[0]
+    
+    return score_acc / sample_num
+
+        
+def forward_modality(model, data, flag):
+    device = next(model.parameters()).device
+    if flag == 'img':
+        features = model.encode_image(data.to(device))
+    elif flag == 'txt':
+        features = model.encode_text(data.to(device))
+    else:
+        raise TypeError
+    return features
+
+
+def main():
+    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--batch-size', type=int, default=50,
+                    help='Batch size to use')
+    parser.add_argument('--clip-model', type=str, default='ViT-B/32',
+                    help='CLIP model to use')
+    parser.add_argument('--num-workers', type=int, default=8,
+                    help=('Number of processes to use for data loading. '
+                          'Defaults to `min(8, num_cpus)`'))
+    parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+    parser.add_argument('--real_flag', type=str, default='img',
+                    help=('The modality of real path. '
+                          'Default to img'))
+    parser.add_argument('--generated_flag', type=str, default='txt',
+                    help=('The modality of generated path. '
+                          'Default to txt'))
+    parser.add_argument('--real_path', type=str, 
+                    help=('Paths to the real images or '
+                          'to .npz statistic files'))
+    parser.add_argument('--generated_path', type=str,
+                    help=('Paths to the generated images or '
+                          'to .npz statistic files'))
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    if args.num_workers is None:
+        try:
+            num_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            # os.sched_getaffinity is not available under Windows, use
+            # os.cpu_count instead (which may not return the *available* number
+            # of CPUs).
+            num_cpus = os.cpu_count()
+
+        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
+    else:
+        num_workers = args.num_workers
+
+    print('Loading CLIP model: {}'.format(args.clip_model))
+    model, preprocess = clip.load(args.clip_model, device=device)
+    
+    dataset = DummyDataset(args.real_path, args.generated_path,
+                           args.real_flag, args.generated_flag,
+                           transform=preprocess, tokenizer=clip.tokenize)
+    dataloader = DataLoader(dataset, args.batch_size, 
+                            num_workers=num_workers, pin_memory=True)
+    
+    print('Calculating CLIP Score:')
+    clip_score = calculate_clip_score(dataloader, model,
+                                      args.real_flag, args.generated_flag)
+    clip_score = clip_score.cpu().item()
+    print('CLIP Score: ', clip_score)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/opensora/models/causalvideovae/eval/eval_common_metric.py b/opensora/models/causalvideovae/eval/eval_common_metric.py
new file mode 100755
index 000000000..331083cec
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/eval_common_metric.py
@@ -0,0 +1,226 @@
+"""Calculates the CLIP Scores
+
+The CLIP model is a contrasitively learned language-image model. There is
+an image encoder and a text encoder. It is believed that the CLIP model could 
+measure the similarity of cross modalities. Please find more information from 
+https://github.com/openai/CLIP.
+
+The CLIP Score measures the Cosine Similarity between two embedded features.
+This repository utilizes the pretrained CLIP Model to calculate 
+the mean average of cosine similarities. 
+
+See --help to see further details.
+
+Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.
+
+Copyright 2023 The Hong Kong Polytechnic University
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+import os.path as osp
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader, Subset
+from decord import VideoReader, cpu
+import random
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.io import read_video
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms._transforms_video import CenterCropVideo
+from cal_lpips import calculate_lpips
+from cal_fvd import calculate_fvd
+from cal_psnr import calculate_psnr
+from cal_flolpips import calculate_flolpips
+from cal_ssim import calculate_ssim
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+class VideoDataset(Dataset):
+    def __init__(self, 
+                 real_video_dir,
+                 generated_video_dir,
+                 num_frames,
+                 sample_rate = 1,
+                 crop_size=None,
+                 resolution=128,
+                 ) -> None:
+        super().__init__()
+        self.real_video_files = []
+        self.generated_video_files = self._combine_without_prefix(generated_video_dir)
+        for video_file in self.generated_video_files:
+            filename = os.path.basename(video_file)
+            self.real_video_files.append(os.path.join(real_video_dir, filename))
+        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]
+        generated_video_file = self.generated_video_files[index]
+        real_video_tensor  = self._load_video(real_video_file, self.sample_rate)
+        generated_video_tensor  = self._load_video(generated_video_file, 1)
+        return {'real': real_video_tensor, 'generated':generated_video_tensor }
+
+
+    def _load_video(self, video_path, sample_rate):
+        num_frames = self.num_frames
+        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(0, 3, 1, 2) # (T, H, W, C) -> (C, T, H, W)
+        return _preprocess(video_data, short_size=self.short_size, crop_size = self.crop_size)
+
+
+    def _combine_without_prefix(self, folder_path, prefix='.'):
+        folder = []
+        os.makedirs(folder_path, exist_ok=True)
+        for name in os.listdir(folder_path):
+            if name[0] == prefix:
+                continue
+            if osp.isfile(osp.join(folder_path, name)):
+                folder.append(osp.join(folder_path, name))
+        folder.sort()
+        return folder
+
+def _preprocess(video_data, short_size=128, crop_size=None):
+    transform = Compose(
+        [
+            Lambda(lambda x: x / 255.0),
+            ShortSideScale(size=short_size),
+            CenterCropVideo(crop_size=crop_size),
+        ]
+    )
+    video_outputs = transform(video_data)
+    # video_outputs = torch.unsqueeze(video_outputs, 0) # (bz,c,t,h,w)
+    return video_outputs
+
+
+def calculate_common_metric(args, dataloader, device):
+
+    score_list = []
+    for batch_data in tqdm(dataloader): # {'real': real_video_tensor, 'generated':generated_video_tensor }
+        real_videos = batch_data['real'] 
+        generated_videos = batch_data['generated']
+        
+        assert real_videos.shape[2] == generated_videos.shape[2]
+        if args.metric == 'fvd':
+            tmp_list = list(calculate_fvd(real_videos, generated_videos, args.device, method=args.fvd_method)['value'].values())
+        elif args.metric == 'ssim':
+            tmp_list = list(calculate_ssim(real_videos, generated_videos)['value'].values())
+        elif args.metric == 'psnr':
+            tmp_list = list(calculate_psnr(real_videos, generated_videos)['value'].values())
+        elif args.metric == 'flolpips':
+            result = calculate_flolpips(real_videos, generated_videos, args.device)
+            tmp_list = list(result['value'].values())
+        elif args.metric == 'wraperror':
+            result = calculate_wraperror(real_videos, generated_videos, args.device)
+            tmp_list = list(result['value'].values())
+        else:
+            tmp_list  = list(calculate_lpips(real_videos, generated_videos, args.device)['value'].values())
+        score_list += tmp_list
+    return np.mean(score_list)
+        
+def main():
+    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--batch_size', type=int, default=2,
+                    help='Batch size to use')
+    parser.add_argument('--real_video_dir', type=str,
+                    help=('the path of real videos`'))
+    parser.add_argument('--generated_video_dir', type=str,
+                    help=('the path of generated videos`'))
+    parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+    parser.add_argument('--num_workers', type=int, default=8,
+                    help=('Number of processes to use for data loading. '
+                          'Defaults to `min(8, num_cpus)`'))
+    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('--subset_size', type=int, default=None)
+    parser.add_argument("--metric", type=str, default="fvd",choices=['fvd','psnr','ssim','lpips', 'flolpips'])
+    parser.add_argument("--fvd_method", type=str, default='styleganv',choices=['styleganv','videogpt'])
+
+
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda:1' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    if args.num_workers is None:
+        try:
+            num_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            # os.sched_getaffinity is not available under Windows, use
+            # os.cpu_count instead (which may not return the *available* number
+            # of CPUs).
+            num_cpus = os.cpu_count()
+
+        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
+    else:
+        num_workers = args.num_workers
+
+    
+    dataset = VideoDataset(args.real_video_dir,
+                           args.generated_video_dir,
+                            num_frames = args.num_frames,
+                            sample_rate = args.sample_rate,
+                            crop_size=args.crop_size,
+                            resolution=args.resolution) 
+
+    if args.subset_size:
+        indices = range(args.subset_size)
+        dataset = Subset(dataset, indices=indices)
+    
+    dataloader = DataLoader(dataset, args.batch_size, 
+                            num_workers=num_workers, pin_memory=True)
+    
+
+    metric_score = calculate_common_metric(args, dataloader,device)
+    print('metric: ', args.metric, " ",metric_score)
+
+if __name__ == '__main__':
+    main()
diff --git a/opensora/models/causalvideovae/eval/flolpips/correlation/correlation.py b/opensora/models/causalvideovae/eval/flolpips/correlation/correlation.py
new file mode 100755
index 000000000..5f88f0107
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/flolpips/correlation/correlation.py
@@ -0,0 +1,397 @@
+#!/usr/bin/env python
+
+import torch
+
+import cupy#add
+import re
+
+kernel_Correlation_rearrange = '''
+	extern "C" __global__ void kernel_Correlation_rearrange(
+		const int n,
+		const float* input,
+		float* output
+	) {
+	  int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x;
+
+	  if (intIndex >= n) {
+	    return;
+	  }
+
+	  int intSample = blockIdx.z;
+	  int intChannel = blockIdx.y;
+
+	  float fltValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex];
+
+	  __syncthreads();
+
+	  int intPaddedY = (intIndex / SIZE_3(input)) + 4;
+	  int intPaddedX = (intIndex % SIZE_3(input)) + 4;
+	  int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX;
+
+	  output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = fltValue;
+	}
+'''
+
+kernel_Correlation_updateOutput = '''
+	extern "C" __global__ void kernel_Correlation_updateOutput(
+	  const int n,
+	  const float* rbot0,
+	  const float* rbot1,
+	  float* top
+	) {
+	  extern __shared__ char patch_data_char[];
+	  
+	  float *patch_data = (float *)patch_data_char;
+	  
+	  // First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1
+	  int x1 = blockIdx.x + 4;
+	  int y1 = blockIdx.y + 4;
+	  int item = blockIdx.z;
+	  int ch_off = threadIdx.x;
+	  
+	  // Load 3D patch into shared shared memory
+	  for (int j = 0; j < 1; j++) { // HEIGHT
+	    for (int i = 0; i < 1; i++) { // WIDTH
+	      int ji_off = (j + i) * SIZE_3(rbot0);
+	      for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
+	        int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch;
+	        int idxPatchData = ji_off + ch;
+	        patch_data[idxPatchData] = rbot0[idx1];
+	      }
+	    }
+	  }
+	  
+	  __syncthreads();
+	  
+	  __shared__ float sum[32];
+	  
+	  // Compute correlation
+	  for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) {
+	    sum[ch_off] = 0;
+	  
+	    int s2o = top_channel % 9 - 4;
+	    int s2p = top_channel / 9 - 4;
+	    
+	    for (int j = 0; j < 1; j++) { // HEIGHT
+	      for (int i = 0; i < 1; i++) { // WIDTH
+	        int ji_off = (j + i) * SIZE_3(rbot0);
+	        for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
+	          int x2 = x1 + s2o;
+	          int y2 = y1 + s2p;
+	          
+	          int idxPatchData = ji_off + ch;
+	          int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch;
+	          
+	          sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2];
+	        }
+	      }
+	    }
+	    
+	    __syncthreads();
+	    
+	    if (ch_off == 0) {
+	      float total_sum = 0;
+	      for (int idx = 0; idx < 32; idx++) {
+	        total_sum += sum[idx];
+	      }
+	      const int sumelems = SIZE_3(rbot0);
+	      const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x;
+	      top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems;
+	    }
+	  }
+	}
+'''
+
+kernel_Correlation_updateGradFirst = '''
+	#define ROUND_OFF 50000
+
+	extern "C" __global__ void kernel_Correlation_updateGradFirst(
+	  const int n,
+	  const int intSample,
+	  const float* rbot0,
+	  const float* rbot1,
+	  const float* gradOutput,
+	  float* gradFirst,
+	  float* gradSecond
+	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
+	  int n = intIndex % SIZE_1(gradFirst); // channels
+	  int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos
+	  int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos
+	  
+	  // round_off is a trick to enable integer division with ceil, even for negative numbers
+	  // We use a large offset, for the inner part not to become negative.
+	  const int round_off = ROUND_OFF;
+	  const int round_off_s1 = round_off;
+	  
+	  // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
+	  int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
+	  int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
+	  
+	  // Same here:
+	  int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4)
+	  int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4)
+	  
+	  float sum = 0;
+	  if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
+	    xmin = max(0,xmin);
+	    xmax = min(SIZE_3(gradOutput)-1,xmax);
+	    
+	    ymin = max(0,ymin);
+	    ymax = min(SIZE_2(gradOutput)-1,ymax);
+	    
+	    for (int p = -4; p <= 4; p++) {
+	      for (int o = -4; o <= 4; o++) {
+	        // Get rbot1 data:
+	        int s2o = o;
+	        int s2p = p;
+	        int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n;
+	        float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n]
+	        
+	        // Index offset for gradOutput in following loops:
+	        int op = (p+4) * 9 + (o+4); // index[o,p]
+	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
+	        
+	        for (int y = ymin; y <= ymax; y++) {
+	          for (int x = xmin; x <= xmax; x++) {
+	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
+	            sum += gradOutput[idxgradOutput] * bot1tmp;
+	          }
+	        }
+	      }
+	    }
+	  }
+	  const int sumelems = SIZE_1(gradFirst);
+	  const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4);
+	  gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems;
+	} }
+'''
+
+kernel_Correlation_updateGradSecond = '''
+	#define ROUND_OFF 50000
+
+	extern "C" __global__ void kernel_Correlation_updateGradSecond(
+	  const int n,
+	  const int intSample,
+	  const float* rbot0,
+	  const float* rbot1,
+	  const float* gradOutput,
+	  float* gradFirst,
+	  float* gradSecond
+	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
+	  int n = intIndex % SIZE_1(gradSecond); // channels
+	  int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos
+	  int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos
+	  
+	  // round_off is a trick to enable integer division with ceil, even for negative numbers
+	  // We use a large offset, for the inner part not to become negative.
+	  const int round_off = ROUND_OFF;
+	  const int round_off_s1 = round_off;
+	  
+	  float sum = 0;
+	  for (int p = -4; p <= 4; p++) {
+	    for (int o = -4; o <= 4; o++) {
+	      int s2o = o;
+	      int s2p = p;
+	      
+	      //Get X,Y ranges and clamp
+	      // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
+	      int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
+	      int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
+	      
+	      // Same here:
+	      int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o)
+	      int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p)
+          
+	      if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
+	        xmin = max(0,xmin);
+	        xmax = min(SIZE_3(gradOutput)-1,xmax);
+	        
+	        ymin = max(0,ymin);
+	        ymax = min(SIZE_2(gradOutput)-1,ymax);
+	        
+	        // Get rbot0 data:
+	        int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n;
+	        float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n]
+	        
+	        // Index offset for gradOutput in following loops:
+	        int op = (p+4) * 9 + (o+4); // index[o,p]
+	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
+	        
+	        for (int y = ymin; y <= ymax; y++) {
+	          for (int x = xmin; x <= xmax; x++) {
+	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
+	            sum += gradOutput[idxgradOutput] * bot0tmp;
+	          }
+	        }
+	      }
+	    }
+	  }
+	  const int sumelems = SIZE_1(gradSecond);
+	  const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4);
+	  gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems;
+	} }
+'''
+
+def cupy_kernel(strFunction, objVariables):
+	strKernel = globals()[strFunction]
+
+	while True:
+		objMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel)
+
+		if objMatch is None:
+			break
+		# end
+
+		intArg = int(objMatch.group(2))
+
+		strTensor = objMatch.group(4)
+		intSizes = objVariables[strTensor].size()
+
+		strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg]))
+	# end
+
+	while True:
+		objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)
+
+		if objMatch is None:
+			break
+		# end
+
+		intArgs = int(objMatch.group(2))
+		strArgs = objMatch.group(4).split(',')
+
+		strTensor = strArgs[0]
+		intStrides = objVariables[strTensor].stride()
+		strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]
+
+		strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')
+	# end
+
+	return strKernel
+# end
+
+@cupy.memoize(for_each_device=True)
+def cupy_launch(strFunction, strKernel):
+    return cupy.RawKernel(strKernel, strFunction)
+# end
+
+class _FunctionCorrelation(torch.autograd.Function):
+	@staticmethod
+	def forward(self, first, second):
+		rbot0 = first.new_zeros([ first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1] ])
+		rbot1 = first.new_zeros([ first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1] ])
+
+		self.save_for_backward(first, second, rbot0, rbot1)
+
+		first = first.contiguous();	assert(first.is_cuda == True)
+		second = second.contiguous(); assert(second.is_cuda == True)
+
+		output = first.new_zeros([ first.shape[0], 81, first.shape[2], first.shape[3] ])
+
+		if first.is_cuda == True:
+			n = first.shape[2] * first.shape[3]
+			cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
+				'input': first,
+				'output': rbot0
+			}))(
+				grid=tuple([ int((n + 16 - 1) / 16), first.shape[1], first.shape[0] ]),
+				block=tuple([ 16, 1, 1 ]),
+				args=[ n, first.data_ptr(), rbot0.data_ptr() ]
+			)
+
+			n = second.shape[2] * second.shape[3]
+			cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
+				'input': second,
+				'output': rbot1
+			}))(
+				grid=tuple([ int((n + 16 - 1) / 16), second.shape[1], second.shape[0] ]),
+				block=tuple([ 16, 1, 1 ]),
+				args=[ n, second.data_ptr(), rbot1.data_ptr() ]
+			)
+
+			n = output.shape[1] * output.shape[2] * output.shape[3]
+			cupy_launch('kernel_Correlation_updateOutput', cupy_kernel('kernel_Correlation_updateOutput', {
+				'rbot0': rbot0,
+				'rbot1': rbot1,
+				'top': output
+			}))(
+				grid=tuple([ output.shape[3], output.shape[2], output.shape[0] ]),
+				block=tuple([ 32, 1, 1 ]),
+				shared_mem=first.shape[1] * 4,
+				args=[ n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr() ]
+			)
+
+		elif first.is_cuda == False:
+			raise NotImplementedError()
+
+		# end
+
+		return output
+	# end
+
+	@staticmethod
+	def backward(self, gradOutput):
+		first, second, rbot0, rbot1 = self.saved_tensors
+
+		gradOutput = gradOutput.contiguous(); assert(gradOutput.is_cuda == True)
+
+		gradFirst = first.new_zeros([ first.shape[0], first.shape[1], first.shape[2], first.shape[3] ]) if self.needs_input_grad[0] == True else None
+		gradSecond = first.new_zeros([ first.shape[0], first.shape[1], first.shape[2], first.shape[3] ]) if self.needs_input_grad[1] == True else None
+
+		if first.is_cuda == True:
+			if gradFirst is not None:
+				for intSample in range(first.shape[0]):
+					n = first.shape[1] * first.shape[2] * first.shape[3]
+					cupy_launch('kernel_Correlation_updateGradFirst', cupy_kernel('kernel_Correlation_updateGradFirst', {
+						'rbot0': rbot0,
+						'rbot1': rbot1,
+						'gradOutput': gradOutput,
+						'gradFirst': gradFirst,
+						'gradSecond': None
+					}))(
+						grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
+						block=tuple([ 512, 1, 1 ]),
+						args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), gradFirst.data_ptr(), None ]
+					)
+				# end
+			# end
+
+			if gradSecond is not None:
+				for intSample in range(first.shape[0]):
+					n = first.shape[1] * first.shape[2] * first.shape[3]
+					cupy_launch('kernel_Correlation_updateGradSecond', cupy_kernel('kernel_Correlation_updateGradSecond', {
+						'rbot0': rbot0,
+						'rbot1': rbot1,
+						'gradOutput': gradOutput,
+						'gradFirst': None,
+						'gradSecond': gradSecond
+					}))(
+						grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
+						block=tuple([ 512, 1, 1 ]),
+						args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), None, gradSecond.data_ptr() ]
+					)
+				# end
+			# end
+
+		elif first.is_cuda == False:
+			raise NotImplementedError()
+
+		# end
+
+		return gradFirst, gradSecond
+	# end
+# end
+
+def FunctionCorrelation(tenFirst, tenSecond):
+	return _FunctionCorrelation.apply(tenFirst, tenSecond)
+# end
+
+class ModuleCorrelation(torch.nn.Module):
+	def __init__(self):
+		super(ModuleCorrelation, self).__init__()
+	# end
+
+	def forward(self, tenFirst, tenSecond):
+		return _FunctionCorrelation.apply(tenFirst, tenSecond)
+	# end
+# end
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/flolpips/flolpips.py b/opensora/models/causalvideovae/eval/flolpips/flolpips.py
new file mode 100755
index 000000000..66f473149
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/flolpips/flolpips.py
@@ -0,0 +1,308 @@
+
+from __future__ import absolute_import
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .pretrained_networks import vgg16, alexnet, squeezenet
+import torch.nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+import cv2
+
+from .pwcnet import Network as PWCNet
+from .utils import *
+
+def spatial_average(in_tens, keepdim=True):
+    return in_tens.mean([2,3],keepdim=keepdim)
+
+def mw_spatial_average(in_tens, flow, keepdim=True):
+    _,_,h,w = in_tens.shape
+    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
+    flow_mag = torch.sqrt(flow[:,0:1]**2 + flow[:,1:2]**2)
+    flow_mag = flow_mag / torch.sum(flow_mag, dim=[1,2,3], keepdim=True)
+    return torch.sum(in_tens*flow_mag, dim=[2,3],keepdim=keepdim)
+
+
+def mtw_spatial_average(in_tens, flow, texture, keepdim=True):
+    _,_,h,w = in_tens.shape
+    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
+    texture = F.interpolate(texture, (h,w), align_corners=False, mode='bilinear')
+    flow_mag = torch.sqrt(flow[:,0:1]**2 + flow[:,1:2]**2)
+    flow_mag = (flow_mag - flow_mag.min()) / (flow_mag.max() - flow_mag.min()) + 1e-6
+    texture = (texture - texture.min()) / (texture.max() - texture.min()) + 1e-6
+    weight = flow_mag / texture
+    weight /= torch.sum(weight)
+    return torch.sum(in_tens*weight, dim=[2,3],keepdim=keepdim)
+
+
+
+def m2w_spatial_average(in_tens, flow, keepdim=True):
+    _,_,h,w = in_tens.shape
+    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
+    flow_mag = flow[:,0:1]**2 + flow[:,1:2]**2 # B,1,H,W
+    flow_mag = flow_mag / torch.sum(flow_mag)
+    return torch.sum(in_tens*flow_mag, dim=[2,3],keepdim=keepdim)
+
+def upsample(in_tens, out_HW=(64,64)): # assumes scale factor is same for H and W
+    in_H, in_W = in_tens.shape[2], in_tens.shape[3]
+    return nn.Upsample(size=out_HW, mode='bilinear', align_corners=False)(in_tens)
+
+# Learned perceptual metric
+class LPIPS(nn.Module):
+    def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, 
+        pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=False):
+        # lpips - [True] means with linear calibration on top of base network
+        # pretrained - [True] means load linear weights
+
+        super(LPIPS, self).__init__()
+        if(verbose):
+            print('Setting up [%s] perceptual loss: trunk [%s], v[%s], spatial [%s]'%
+                ('LPIPS' if lpips else 'baseline', net, version, 'on' if spatial else 'off'))
+
+        self.pnet_type = net
+        self.pnet_tune = pnet_tune
+        self.pnet_rand = pnet_rand
+        self.spatial = spatial
+        self.lpips = lpips # false means baseline of just averaging all layers
+        self.version = version
+        self.scaling_layer = ScalingLayer()
+
+        if(self.pnet_type in ['vgg','vgg16']):
+            net_type = vgg16
+            self.chns = [64,128,256,512,512]
+        elif(self.pnet_type=='alex'):
+            net_type = alexnet
+            self.chns = [64,192,384,256,256]
+        elif(self.pnet_type=='squeeze'):
+            net_type = squeezenet
+            self.chns = [64,128,256,384,384,512,512]
+        self.L = len(self.chns)
+
+        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)
+
+        if(lpips):
+            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+            self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4]
+            if(self.pnet_type=='squeeze'): # 7 layers for squeezenet
+                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
+                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
+                self.lins+=[self.lin5,self.lin6]
+            self.lins = nn.ModuleList(self.lins)
+
+            if(pretrained):
+                if(model_path is None):
+                    import inspect
+                    import os
+                    model_path = os.path.abspath(os.path.join(inspect.getfile(self.__init__), '..', 'weights/v%s/%s.pth'%(version,net)))
+
+                if(verbose):
+                    print('Loading model from: %s'%model_path)
+                self.load_state_dict(torch.load(model_path, map_location='cpu'), strict=False)          
+
+        if(eval_mode):
+            self.eval()
+
+    def forward(self, in0, in1, retPerLayer=False, normalize=False):
+        if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
+            in0 = 2 * in0  - 1
+            in1 = 2 * in1  - 1
+
+        # v0.0 - original release had a bug, where input was not scaled
+        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
+        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+
+        for kk in range(self.L):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk]-feats1[kk])**2
+
+        if(self.lpips):
+            if(self.spatial):
+                res = [upsample(self.lins[kk](diffs[kk]), out_HW=in0.shape[2:]) for kk in range(self.L)]
+            else:
+                res = [spatial_average(self.lins[kk](diffs[kk]), keepdim=True) for kk in range(self.L)]
+        else:
+            if(self.spatial):
+                res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_HW=in0.shape[2:]) for kk in range(self.L)]
+            else:
+                res = [spatial_average(diffs[kk].sum(dim=1,keepdim=True), keepdim=True) for kk in range(self.L)]
+
+        # val = res[0]
+        # for l in range(1,self.L):
+        #     val += res[l]
+        #     print(val)
+
+        # a = spatial_average(self.lins[kk](diffs[kk]), keepdim=True)
+        # b = torch.max(self.lins[kk](feats0[kk]**2))
+        # for kk in range(self.L):
+        #     a += spatial_average(self.lins[kk](diffs[kk]), keepdim=True)
+        #     b = torch.max(b,torch.max(self.lins[kk](feats0[kk]**2)))
+        # a = a/self.L
+        # from IPython import embed
+        # embed()
+        # return 10*torch.log10(b/a)
+        
+        # if(retPerLayer):
+        #     return (val, res)
+        # else:
+        return torch.sum(torch.cat(res, 1), dim=(1,2,3), keepdims=False)
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030,-.088,-.188])[None,:,None,None])
+        self.register_buffer('scale', torch.Tensor([.458,.448,.450])[None,:,None,None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    ''' A single linear layer which does a 1x1 conv '''
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+
+        layers = [nn.Dropout(),] if(use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),]
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+class Dist2LogitLayer(nn.Module):
+    ''' takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) '''
+    def __init__(self, chn_mid=32, use_sigmoid=True):
+        super(Dist2LogitLayer, self).__init__()
+
+        layers = [nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),]
+        layers += [nn.LeakyReLU(0.2,True),]
+        layers += [nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),]
+        layers += [nn.LeakyReLU(0.2,True),]
+        layers += [nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),]
+        if(use_sigmoid):
+            layers += [nn.Sigmoid(),]
+        self.model = nn.Sequential(*layers)
+
+    def forward(self,d0,d1,eps=0.1):
+        return self.model.forward(torch.cat((d0,d1,d0-d1,d0/(d1+eps),d1/(d0+eps)),dim=1))
+
+class BCERankingLoss(nn.Module):
+    def __init__(self, chn_mid=32):
+        super(BCERankingLoss, self).__init__()
+        self.net = Dist2LogitLayer(chn_mid=chn_mid)
+        # self.parameters = list(self.net.parameters())
+        self.loss = torch.nn.BCELoss()
+
+    def forward(self, d0, d1, judge):
+        per = (judge+1.)/2.
+        self.logit = self.net.forward(d0,d1)
+        return self.loss(self.logit, per)
+
+# L2, DSSIM metrics
+class FakeNet(nn.Module):
+    def __init__(self, use_gpu=True, colorspace='Lab'):
+        super(FakeNet, self).__init__()
+        self.use_gpu = use_gpu
+        self.colorspace = colorspace
+
+class L2(FakeNet):
+    def forward(self, in0, in1, retPerLayer=None):
+        assert(in0.size()[0]==1) # currently only supports batchSize 1
+
+        if(self.colorspace=='RGB'):
+            (N,C,X,Y) = in0.size()
+            value = torch.mean(torch.mean(torch.mean((in0-in1)**2,dim=1).view(N,1,X,Y),dim=2).view(N,1,1,Y),dim=3).view(N)
+            return value
+        elif(self.colorspace=='Lab'):
+            value = l2(tensor2np(tensor2tensorlab(in0.data,to_norm=False)), 
+                tensor2np(tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
+            ret_var = Variable( torch.Tensor((value,) ) )
+            if(self.use_gpu):
+                ret_var = ret_var.cuda()
+            return ret_var
+
+class DSSIM(FakeNet):
+
+    def forward(self, in0, in1, retPerLayer=None):
+        assert(in0.size()[0]==1) # currently only supports batchSize 1
+
+        if(self.colorspace=='RGB'):
+            value = dssim(1.*tensor2im(in0.data), 1.*tensor2im(in1.data), range=255.).astype('float')
+        elif(self.colorspace=='Lab'):
+            value = dssim(tensor2np(tensor2tensorlab(in0.data,to_norm=False)), 
+                tensor2np(tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
+        ret_var = Variable( torch.Tensor((value,) ) )
+        if(self.use_gpu):
+            ret_var = ret_var.cuda()
+        return ret_var
+
+def print_network(net):
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print('Network',net)
+    print('Total number of parameters: %d' % num_params)
+
+
+class FloLPIPS(LPIPS):
+    def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=False):
+        super(FloLPIPS, self).__init__(pretrained, net, version, lpips, spatial, pnet_rand, pnet_tune, use_dropout, model_path, eval_mode, verbose)
+
+    def forward(self, in0, in1, flow, retPerLayer=False, normalize=False):
+        if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
+            in0 = 2 * in0  - 1
+            in1 = 2 * in1  - 1
+
+        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
+        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+
+        for kk in range(self.L):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk]-feats1[kk])**2
+
+        res = [mw_spatial_average(self.lins[kk](diffs[kk]), flow, keepdim=True) for kk in range(self.L)]
+
+        return torch.sum(torch.cat(res, 1), dim=(1,2,3), keepdims=False)
+
+
+
+
+
+class Flolpips(nn.Module):
+    def __init__(self):
+        super(Flolpips, self).__init__()
+        self.loss_fn = FloLPIPS(net='alex',version='0.1')
+        self.flownet = PWCNet()
+    
+    @torch.no_grad()
+    def forward(self, I0, I1, frame_dis, frame_ref):
+        """
+        args:
+            I0: first frame of the triplet, shape: [B, C, H, W]
+            I1: third frame of the triplet, shape: [B, C, H, W]
+            frame_dis: prediction of the intermediate frame, shape: [B, C, H, W]
+            frame_ref: ground-truth of the intermediate frame, shape: [B, C, H, W]
+        """
+        assert I0.size() == I1.size() == frame_dis.size() == frame_ref.size(), \
+                "the 4 input tensors should have same size"
+
+        flow_ref = self.flownet(frame_ref, I0)
+        flow_dis = self.flownet(frame_dis, I0)
+        flow_diff = flow_ref - flow_dis
+        flolpips_wrt_I0 = self.loss_fn.forward(frame_ref, frame_dis, flow_diff, normalize=True)
+
+        flow_ref = self.flownet(frame_ref, I1)
+        flow_dis = self.flownet(frame_dis, I1)
+        flow_diff = flow_ref - flow_dis
+        flolpips_wrt_I1 = self.loss_fn.forward(frame_ref, frame_dis, flow_diff, normalize=True)
+
+        flolpips = (flolpips_wrt_I0 + flolpips_wrt_I1) / 2
+        return flolpips
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/flolpips/pretrained_networks.py b/opensora/models/causalvideovae/eval/flolpips/pretrained_networks.py
new file mode 100755
index 000000000..82f14e206
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/flolpips/pretrained_networks.py
@@ -0,0 +1,180 @@
+from collections import namedtuple
+import torch
+from torchvision import models as tv
+
+class squeezenet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(squeezenet, self).__init__()
+        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.slice6 = torch.nn.Sequential()
+        self.slice7 = torch.nn.Sequential()
+        self.N_slices = 7
+        for x in range(2):
+            self.slice1.add_module(str(x), pretrained_features[x])
+        for x in range(2,5):
+            self.slice2.add_module(str(x), pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), pretrained_features[x])
+        for x in range(10, 11):
+            self.slice5.add_module(str(x), pretrained_features[x])
+        for x in range(11, 12):
+            self.slice6.add_module(str(x), pretrained_features[x])
+        for x in range(12, 13):
+            self.slice7.add_module(str(x), pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        h = self.slice6(h)
+        h_relu6 = h
+        h = self.slice7(h)
+        h_relu7 = h
+        vgg_outputs = namedtuple("SqueezeOutputs", ['relu1','relu2','relu3','relu4','relu5','relu6','relu7'])
+        out = vgg_outputs(h_relu1,h_relu2,h_relu3,h_relu4,h_relu5,h_relu6,h_relu7)
+
+        return out
+
+
+class alexnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(alexnet, self).__init__()
+        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(2):
+            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(2, 5):
+            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(5, 8):
+            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(8, 10):
+            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
+        for x in range(10, 12):
+            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1 = h
+        h = self.slice2(h)
+        h_relu2 = h
+        h = self.slice3(h)
+        h_relu3 = h
+        h = self.slice4(h)
+        h_relu4 = h
+        h = self.slice5(h)
+        h_relu5 = h
+        alexnet_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5'])
+        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
+
+        return out
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+
+        return out
+
+
+
+class resnet(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True, num=18):
+        super(resnet, self).__init__()
+        if(num==18):
+            self.net = tv.resnet18(pretrained=pretrained)
+        elif(num==34):
+            self.net = tv.resnet34(pretrained=pretrained)
+        elif(num==50):
+            self.net = tv.resnet50(pretrained=pretrained)
+        elif(num==101):
+            self.net = tv.resnet101(pretrained=pretrained)
+        elif(num==152):
+            self.net = tv.resnet152(pretrained=pretrained)
+        self.N_slices = 5
+
+        self.conv1 = self.net.conv1
+        self.bn1 = self.net.bn1
+        self.relu = self.net.relu
+        self.maxpool = self.net.maxpool
+        self.layer1 = self.net.layer1
+        self.layer2 = self.net.layer2
+        self.layer3 = self.net.layer3
+        self.layer4 = self.net.layer4
+
+    def forward(self, X):
+        h = self.conv1(X)
+        h = self.bn1(h)
+        h = self.relu(h)
+        h_relu1 = h
+        h = self.maxpool(h)
+        h = self.layer1(h)
+        h_conv2 = h
+        h = self.layer2(h)
+        h_conv3 = h
+        h = self.layer3(h)
+        h_conv4 = h
+        h = self.layer4(h)
+        h_conv5 = h
+
+        outputs = namedtuple("Outputs", ['relu1','conv2','conv3','conv4','conv5'])
+        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
+
+        return out
diff --git a/opensora/models/causalvideovae/eval/flolpips/pwcnet.py b/opensora/models/causalvideovae/eval/flolpips/pwcnet.py
new file mode 100755
index 000000000..374f18e20
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/flolpips/pwcnet.py
@@ -0,0 +1,344 @@
+#!/usr/bin/env python
+
+import torch
+
+import getopt
+import math
+import numpy
+import os
+import PIL
+import PIL.Image
+import sys
+
+# try:
+from .correlation import correlation # the custom cost volume layer
+# except:
+# 	sys.path.insert(0, './correlation'); import correlation # you should consider upgrading python
+# end
+
+##########################################################
+
+# assert(int(str('').join(torch.__version__.split('.')[0:2])) >= 13) # requires at least pytorch version 1.3.0
+
+# torch.set_grad_enabled(False) # make sure to not compute gradients for computational performance
+
+# torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
+
+# ##########################################################
+
+# arguments_strModel = 'default' # 'default', or 'chairs-things'
+# arguments_strFirst = './images/first.png'
+# arguments_strSecond = './images/second.png'
+# arguments_strOut = './out.flo'
+
+# for strOption, strArgument in getopt.getopt(sys.argv[1:], '', [ strParameter[2:] + '=' for strParameter in sys.argv[1::2] ])[0]:
+# 	if strOption == '--model' and strArgument != '': arguments_strModel = strArgument # which model to use
+# 	if strOption == '--first' and strArgument != '': arguments_strFirst = strArgument # path to the first frame
+# 	if strOption == '--second' and strArgument != '': arguments_strSecond = strArgument # path to the second frame
+# 	if strOption == '--out' and strArgument != '': arguments_strOut = strArgument # path to where the output should be stored
+# end
+
+##########################################################
+
+
+
+def backwarp(tenInput, tenFlow):
+	backwarp_tenGrid = {}
+	backwarp_tenPartial = {}
+	if str(tenFlow.shape) not in backwarp_tenGrid:
+		tenHor = torch.linspace(-1.0 + (1.0 / tenFlow.shape[3]), 1.0 - (1.0 / tenFlow.shape[3]), tenFlow.shape[3]).view(1, 1, 1, -1).expand(-1, -1, tenFlow.shape[2], -1)
+		tenVer = torch.linspace(-1.0 + (1.0 / tenFlow.shape[2]), 1.0 - (1.0 / tenFlow.shape[2]), tenFlow.shape[2]).view(1, 1, -1, 1).expand(-1, -1, -1, tenFlow.shape[3])
+
+		backwarp_tenGrid[str(tenFlow.shape)] = torch.cat([ tenHor, tenVer ], 1).cuda()
+	# end
+
+	if str(tenFlow.shape) not in backwarp_tenPartial:
+		backwarp_tenPartial[str(tenFlow.shape)] = tenFlow.new_ones([ tenFlow.shape[0], 1, tenFlow.shape[2], tenFlow.shape[3] ])
+	# end
+
+	tenFlow = torch.cat([ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0) ], 1)
+	tenInput = torch.cat([ tenInput, backwarp_tenPartial[str(tenFlow.shape)] ], 1)
+
+	tenOutput = torch.nn.functional.grid_sample(input=tenInput, grid=(backwarp_tenGrid[str(tenFlow.shape)] + tenFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros', align_corners=False)
+
+	tenMask = tenOutput[:, -1:, :, :]; tenMask[tenMask > 0.999] = 1.0; tenMask[tenMask < 1.0] = 0.0
+
+	return tenOutput[:, :-1, :, :] * tenMask
+# end
+
+##########################################################
+
+class Network(torch.nn.Module):
+	def __init__(self):
+		super(Network, self).__init__()
+
+		class Extractor(torch.nn.Module):
+			def __init__(self):
+				super(Extractor, self).__init__()
+
+				self.netOne = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netTwo = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netThr = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFou = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFiv = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netSix = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=128, out_channels=196, kernel_size=3, stride=2, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+			# end
+
+			def forward(self, tenInput):
+				tenOne = self.netOne(tenInput)
+				tenTwo = self.netTwo(tenOne)
+				tenThr = self.netThr(tenTwo)
+				tenFou = self.netFou(tenThr)
+				tenFiv = self.netFiv(tenFou)
+				tenSix = self.netSix(tenFiv)
+
+				return [ tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix ]
+			# end
+		# end
+
+		class Decoder(torch.nn.Module):
+			def __init__(self, intLevel):
+				super(Decoder, self).__init__()
+
+				intPrevious = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 1]
+				intCurrent = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 0]
+
+				if intLevel < 6: self.netUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1)
+				if intLevel < 6: self.netUpfeat = torch.nn.ConvTranspose2d(in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=4, stride=2, padding=1)
+				if intLevel < 6: self.fltBackwarp = [ None, None, None, 5.0, 2.5, 1.25, 0.625, None ][intLevel + 1]
+
+				self.netOne = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netTwo = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128, out_channels=128, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netThr = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128, out_channels=96, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFou = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96, out_channels=64, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netFiv = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64, out_channels=32, kernel_size=3, stride=1, padding=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
+				)
+
+				self.netSix = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=3, stride=1, padding=1)
+				)
+			# end
+
+			def forward(self, tenFirst, tenSecond, objPrevious):
+				tenFlow = None
+				tenFeat = None
+
+				if objPrevious is None:
+					tenFlow = None
+					tenFeat = None
+
+					tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=tenSecond), negative_slope=0.1, inplace=False)
+
+					tenFeat = torch.cat([ tenVolume ], 1)
+
+				elif objPrevious is not None:
+					tenFlow = self.netUpflow(objPrevious['tenFlow'])
+					tenFeat = self.netUpfeat(objPrevious['tenFeat'])
+
+					tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=backwarp(tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp)), negative_slope=0.1, inplace=False)
+
+					tenFeat = torch.cat([ tenVolume, tenFirst, tenFlow, tenFeat ], 1)
+
+				# end
+
+				tenFeat = torch.cat([ self.netOne(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netTwo(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netThr(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netFou(tenFeat), tenFeat ], 1)
+				tenFeat = torch.cat([ self.netFiv(tenFeat), tenFeat ], 1)
+
+				tenFlow = self.netSix(tenFeat)
+
+				return {
+					'tenFlow': tenFlow,
+					'tenFeat': tenFeat
+				}
+			# end
+		# end
+
+		class Refiner(torch.nn.Module):
+			def __init__(self):
+				super(Refiner, self).__init__()
+
+				self.netMain = torch.nn.Sequential(
+					torch.nn.Conv2d(in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=2, dilation=2),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=4, dilation=4),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=128, out_channels=96, kernel_size=3, stride=1, padding=8, dilation=8),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, stride=1, padding=16, dilation=16),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, dilation=1),
+					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
+					torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1, dilation=1)
+				)
+			# end
+
+			def forward(self, tenInput):
+				return self.netMain(tenInput)
+			# end
+		# end
+
+		self.netExtractor = Extractor()
+
+		self.netTwo = Decoder(2)
+		self.netThr = Decoder(3)
+		self.netFou = Decoder(4)
+		self.netFiv = Decoder(5)
+		self.netSix = Decoder(6)
+
+		self.netRefiner = Refiner()
+
+		self.load_state_dict({ strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/github/pytorch-pwc/network-' + 'default' + '.pytorch').items() })
+	# end
+
+	def forward(self, tenFirst, tenSecond):
+		intWidth = tenFirst.shape[3]
+		intHeight = tenFirst.shape[2]
+
+		intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
+		intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
+
+		tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+		tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+		
+		tenFirst = self.netExtractor(tenPreprocessedFirst)
+		tenSecond = self.netExtractor(tenPreprocessedSecond)
+
+
+		objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None)
+		objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate)
+		objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate)
+		objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate)
+		objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate)
+
+		tenFlow = objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat'])
+		tenFlow = 20.0 * torch.nn.functional.interpolate(input=tenFlow, size=(intHeight, intWidth), mode='bilinear', align_corners=False)
+		tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
+		tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
+
+		return tenFlow
+	# end
+# end
+
+netNetwork = None
+
+##########################################################
+
+def estimate(tenFirst, tenSecond):
+	global netNetwork
+
+	if netNetwork is None:
+		netNetwork = Network().cuda().eval()
+	# end
+
+	assert(tenFirst.shape[1] == tenSecond.shape[1])
+	assert(tenFirst.shape[2] == tenSecond.shape[2])
+
+	intWidth = tenFirst.shape[2]
+	intHeight = tenFirst.shape[1]
+
+	assert(intWidth == 1024) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
+	assert(intHeight == 436) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
+
+	tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth)
+	tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth)
+
+	intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
+	intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
+
+	tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+	tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
+
+	tenFlow = 20.0 * torch.nn.functional.interpolate(input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), size=(intHeight, intWidth), mode='bilinear', align_corners=False)
+
+	tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
+	tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
+
+	return tenFlow[0, :, :, :].cpu()
+# end
+
+##########################################################
+
+# if __name__ == '__main__':
+# 	tenFirst = torch.FloatTensor(numpy.ascontiguousarray(numpy.array(PIL.Image.open(arguments_strFirst))[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) * (1.0 / 255.0)))
+# 	tenSecond = torch.FloatTensor(numpy.ascontiguousarray(numpy.array(PIL.Image.open(arguments_strSecond))[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) * (1.0 / 255.0)))
+
+# 	tenOutput = estimate(tenFirst, tenSecond)
+
+# 	objOutput = open(arguments_strOut, 'wb')
+
+# 	numpy.array([ 80, 73, 69, 72 ], numpy.uint8).tofile(objOutput)
+# 	numpy.array([ tenOutput.shape[2], tenOutput.shape[1] ], numpy.int32).tofile(objOutput)
+# 	numpy.array(tenOutput.numpy().transpose(1, 2, 0), numpy.float32).tofile(objOutput)
+
+# 	objOutput.close()
+# end
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/flolpips/utils.py b/opensora/models/causalvideovae/eval/flolpips/utils.py
new file mode 100755
index 000000000..98d4502d5
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/flolpips/utils.py
@@ -0,0 +1,95 @@
+import numpy as np
+import cv2
+import torch
+
+
+def normalize_tensor(in_feat,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True))
+    return in_feat/(norm_factor+eps)
+
+def l2(p0, p1, range=255.):
+    return .5*np.mean((p0 / range - p1 / range)**2)
+
+def dssim(p0, p1, range=255.):
+    from skimage.measure import compare_ssim
+    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.
+
+def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
+    image_numpy = image_tensor[0].cpu().float().numpy()
+    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
+    return image_numpy.astype(imtype)
+
+def tensor2np(tensor_obj):
+    # change dimension of a tensor object into a numpy array
+    return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))
+
+def np2tensor(np_obj):
+     # change dimenion of np array into tensor array
+    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
+
+def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):
+    # image tensor to lab tensor
+    from skimage import color
+
+    img = tensor2im(image_tensor)
+    img_lab = color.rgb2lab(img)
+    if(mc_only):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+    if(to_norm and not mc_only):
+        img_lab[:,:,0] = img_lab[:,:,0]-50
+        img_lab = img_lab/100.
+
+    return np2tensor(img_lab)
+
+def read_frame_yuv2rgb(stream, width, height, iFrame, bit_depth, pix_fmt='420'):
+    if pix_fmt == '420':
+        multiplier = 1
+        uv_factor = 2
+    elif pix_fmt == '444':
+        multiplier = 2
+        uv_factor = 1
+    else:
+        print('Pixel format {} is not supported'.format(pix_fmt))
+        return
+
+    if bit_depth == 8:
+        datatype = np.uint8
+        stream.seek(iFrame*1.5*width*height*multiplier)
+        Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
+        
+        # read chroma samples and upsample since original is 4:2:0 sampling
+        U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+        V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+
+    else:
+        datatype = np.uint16
+        stream.seek(iFrame*3*width*height*multiplier)
+        Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
+                
+        U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+        V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
+                                reshape((height//uv_factor, width//uv_factor))
+
+    if pix_fmt == '420':
+        yuv = np.empty((height*3//2, width), dtype=datatype)
+        yuv[0:height,:] = Y
+
+        yuv[height:height+height//4,:] = U.reshape(-1, width)
+        yuv[height+height//4:,:] = V.reshape(-1, width)
+
+        if bit_depth != 8:
+            yuv = (yuv/(2**bit_depth-1)*255).astype(np.uint8)
+
+        #convert to rgb
+        rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_I420)
+    
+    else:
+        yvu = np.stack([Y,V,U],axis=2)
+        if bit_depth != 8:
+            yvu = (yvu/(2**bit_depth-1)*255).astype(np.uint8)
+        rgb = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2RGB)
+
+    return rgb
diff --git a/opensora/models/causalvideovae/eval/fvd/styleganv/fvd.py b/opensora/models/causalvideovae/eval/fvd/styleganv/fvd.py
new file mode 100755
index 000000000..3043a2a4a
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/fvd/styleganv/fvd.py
@@ -0,0 +1,90 @@
+import torch
+import os
+import math
+import torch.nn.functional as F
+
+# https://github.com/universome/fvd-comparison
+
+
+def load_i3d_pretrained(device=torch.device('cpu')):
+    i3D_WEIGHTS_URL = "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt"
+    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_torchscript.pt')
+    print(filepath)
+    if not os.path.exists(filepath):
+        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
+        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+    i3d = torch.jit.load(filepath).eval().to(device)
+    i3d = torch.nn.DataParallel(i3d)
+    return i3d
+    
+
+def get_feats(videos, detector, device, bs=10):
+    # videos : torch.tensor BCTHW [0, 1]
+    detector_kwargs = dict(rescale=False, resize=False, return_features=True) # Return raw features before the softmax layer.
+    feats = np.empty((0, 400))
+    with torch.no_grad():
+        for i in range((len(videos)-1)//bs + 1):
+            feats = np.vstack([feats, detector(torch.stack([preprocess_single(video) for video in videos[i*bs:(i+1)*bs]]).to(device), **detector_kwargs).detach().cpu().numpy()])
+    return feats
+
+
+def get_fvd_feats(videos, i3d, device, bs=10):
+    # videos in [0, 1] as torch tensor BCTHW
+    # videos = [preprocess_single(video) for video in videos]
+    embeddings = get_feats(videos, i3d, device, bs)
+    return embeddings
+
+
+def preprocess_single(video, resolution=224, sequence_length=None):
+    # video: CTHW, [0, 1]
+    c, t, h, w = video.shape
+
+    # temporal crop
+    if sequence_length is not None:
+        assert sequence_length <= t
+        video = video[:, :sequence_length]
+
+    # scale shorter side to resolution
+    scale = resolution / min(h, w)
+    if h < w:
+        target_size = (resolution, math.ceil(w * scale))
+    else:
+        target_size = (math.ceil(h * scale), resolution)
+    video = F.interpolate(video, size=target_size, mode='bilinear', align_corners=False)
+
+    # center crop
+    c, t, h, w = video.shape
+    w_start = (w - resolution) // 2
+    h_start = (h - resolution) // 2
+    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
+
+    # [0, 1] -> [-1, 1]
+    video = (video - 0.5) * 2
+
+    return video.contiguous()
+
+
+"""
+Copy-pasted from https://github.com/cvpr2022-stylegan-v/stylegan-v/blob/main/src/metrics/frechet_video_distance.py
+"""
+from typing import Tuple
+from scipy.linalg import sqrtm
+import numpy as np
+
+
+def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    mu = feats.mean(axis=0) # [d]
+    sigma = np.cov(feats, rowvar=False) # [d, d]
+    return mu, sigma
+
+
+def frechet_distance(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
+    mu_gen, sigma_gen = compute_stats(feats_fake)
+    mu_real, sigma_real = compute_stats(feats_real)
+    m = np.square(mu_gen - mu_real).sum()
+    if feats_fake.shape[0]>1:
+        s, _ = sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
+        fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
+    else:
+        fid = np.real(m)
+    return float(fid)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/fvd/videogpt/fvd.py b/opensora/models/causalvideovae/eval/fvd/videogpt/fvd.py
new file mode 100755
index 000000000..e81c82904
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/fvd/videogpt/fvd.py
@@ -0,0 +1,137 @@
+import torch
+import os
+import math
+import torch.nn.functional as F
+import numpy as np
+import einops
+
+def load_i3d_pretrained(device=torch.device('cpu')):
+    i3D_WEIGHTS_URL = "https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI"
+    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_pretrained_400.pt')
+    print(filepath)
+    if not os.path.exists(filepath):
+        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
+        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
+    from .pytorch_i3d import InceptionI3d
+    i3d = InceptionI3d(400, in_channels=3).eval().to(device)
+    i3d.load_state_dict(torch.load(filepath, map_location=device))
+    i3d = torch.nn.DataParallel(i3d)
+    return i3d
+
+def preprocess_single(video, resolution, sequence_length=None):
+    # video: THWC, {0, ..., 255}
+    video = video.permute(0, 3, 1, 2).float() / 255. # TCHW
+    t, c, h, w = video.shape
+
+    # temporal crop
+    if sequence_length is not None:
+        assert sequence_length <= t
+        video = video[:sequence_length]
+
+    # scale shorter side to resolution
+    scale = resolution / min(h, w)
+    if h < w:
+        target_size = (resolution, math.ceil(w * scale))
+    else:
+        target_size = (math.ceil(h * scale), resolution)
+    video = F.interpolate(video, size=target_size, mode='bilinear',
+                          align_corners=False)
+
+    # center crop
+    t, c, h, w = video.shape
+    w_start = (w - resolution) // 2
+    h_start = (h - resolution) // 2
+    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
+    video = video.permute(1, 0, 2, 3).contiguous() # CTHW
+
+    video -= 0.5
+
+    return video
+
+def preprocess(videos, target_resolution=224):
+    # we should tras videos in [0-1] [b c t h w] as th.float 
+    # -> videos in {0, ..., 255} [b t h w c] as np.uint8 array
+    videos = einops.rearrange(videos, 'b c t h w -> b t h w c')
+    videos = (videos*255).numpy().astype(np.uint8)
+
+    b, t, h, w, c = videos.shape
+    videos = torch.from_numpy(videos)
+    videos = torch.stack([preprocess_single(video, target_resolution) for video in videos])
+    return videos * 2 # [-0.5, 0.5] -> [-1, 1]
+
+def get_fvd_logits(videos, i3d, device, bs=10):
+    videos = preprocess(videos)
+    embeddings = get_logits(i3d, videos, device, bs=10)
+    return embeddings
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L161
+def _symmetric_matrix_square_root(mat, eps=1e-10):
+    u, s, v = torch.svd(mat)
+    si = torch.where(s < eps, s, torch.sqrt(s))
+    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
+
+# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L400
+def trace_sqrt_product(sigma, sigma_v):
+    sqrt_sigma = _symmetric_matrix_square_root(sigma)
+    sqrt_a_sigmav_a = torch.matmul(sqrt_sigma, torch.matmul(sigma_v, sqrt_sigma))
+    return torch.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a))
+
+# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
+def cov(m, rowvar=False):
+    '''Estimate a covariance matrix given data.
+
+    Covariance indicates the level to which two variables vary together.
+    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
+    then the covariance matrix element `C_{ij}` is the covariance of
+    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.
+
+    Args:
+        m: A 1-D or 2-D array containing multiple variables and observations.
+            Each row of `m` represents a variable, and each column a single
+            observation of all those variables.
+        rowvar: If `rowvar` is True, then each row represents a
+            variable, with observations in the columns. Otherwise, the
+            relationship is transposed: each column represents a variable,
+            while the rows contain observations.
+
+    Returns:
+        The covariance matrix of the variables.
+    '''
+    if m.dim() > 2:
+        raise ValueError('m has more than 2 dimensions')
+    if m.dim() < 2:
+        m = m.view(1, -1)
+    if not rowvar and m.size(0) != 1:
+        m = m.t()
+
+    fact = 1.0 / (m.size(1) - 1) # unbiased estimate
+    m -= torch.mean(m, dim=1, keepdim=True)
+    mt = m.t()  # if complex: mt = m.t().conj()
+    return fact * m.matmul(mt).squeeze()
+
+
+def frechet_distance(x1, x2):
+    x1 = x1.flatten(start_dim=1)
+    x2 = x2.flatten(start_dim=1)
+    m, m_w = x1.mean(dim=0), x2.mean(dim=0)
+    sigma, sigma_w = cov(x1, rowvar=False), cov(x2, rowvar=False)
+    mean = torch.sum((m - m_w) ** 2)
+    if x1.shape[0]>1:
+        sqrt_trace_component = trace_sqrt_product(sigma, sigma_w)
+        trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component
+        fd = trace + mean
+    else:
+        fd = np.real(mean)
+    return float(fd)
+
+
+def get_logits(i3d, videos, device, bs=10):
+    # assert videos.shape[0] % 16 == 0
+    with torch.no_grad():
+        logits = []
+        for i in range(0, videos.shape[0], bs):
+            batch = videos[i:i + bs].to(device)
+            # logits.append(i3d.module.extract_features(batch)) # wrong
+            logits.append(i3d(batch)) # right
+        logits = torch.cat(logits, dim=0)
+        return logits
diff --git a/opensora/models/causalvideovae/eval/fvd/videogpt/pytorch_i3d.py b/opensora/models/causalvideovae/eval/fvd/videogpt/pytorch_i3d.py
new file mode 100755
index 000000000..58a16cdd7
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/fvd/videogpt/pytorch_i3d.py
@@ -0,0 +1,322 @@
+# Original code from https://github.com/piergiaj/pytorch-i3d
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class MaxPool3dSamePadding(nn.MaxPool3d):
+
+    def compute_pad(self, dim, s):
+        if s % self.stride[dim] == 0:
+            return max(self.kernel_size[dim] - self.stride[dim], 0)
+        else:
+            return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)
+
+    def forward(self, x):
+        # compute 'same' padding
+        (batch, channel, t, h, w) = x.size()
+        out_t = np.ceil(float(t) / float(self.stride[0]))
+        out_h = np.ceil(float(h) / float(self.stride[1]))
+        out_w = np.ceil(float(w) / float(self.stride[2]))
+        pad_t = self.compute_pad(0, t)
+        pad_h = self.compute_pad(1, h)
+        pad_w = self.compute_pad(2, w)
+
+        pad_t_f = pad_t // 2
+        pad_t_b = pad_t - pad_t_f
+        pad_h_f = pad_h // 2
+        pad_h_b = pad_h - pad_h_f
+        pad_w_f = pad_w // 2
+        pad_w_b = pad_w - pad_w_f
+
+        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
+        x = F.pad(x, pad)
+        return super(MaxPool3dSamePadding, self).forward(x)
+
+
+class Unit3D(nn.Module):
+
+    def __init__(self, in_channels,
+                 output_channels,
+                 kernel_shape=(1, 1, 1),
+                 stride=(1, 1, 1),
+                 padding=0,
+                 activation_fn=F.relu,
+                 use_batch_norm=True,
+                 use_bias=False,
+                 name='unit_3d'):
+
+        """Initializes Unit3D module."""
+        super(Unit3D, self).__init__()
+
+        self._output_channels = output_channels
+        self._kernel_shape = kernel_shape
+        self._stride = stride
+        self._use_batch_norm = use_batch_norm
+        self._activation_fn = activation_fn
+        self._use_bias = use_bias
+        self.name = name
+        self.padding = padding
+
+        self.conv3d = nn.Conv3d(in_channels=in_channels,
+                                out_channels=self._output_channels,
+                                kernel_size=self._kernel_shape,
+                                stride=self._stride,
+                                padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function
+                                bias=self._use_bias)
+
+        if self._use_batch_norm:
+            self.bn = nn.BatchNorm3d(self._output_channels, eps=1e-5, momentum=0.001)
+
+    def compute_pad(self, dim, s):
+        if s % self._stride[dim] == 0:
+            return max(self._kernel_shape[dim] - self._stride[dim], 0)
+        else:
+            return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)
+
+
+    def forward(self, x):
+        # compute 'same' padding
+        (batch, channel, t, h, w) = x.size()
+        out_t = np.ceil(float(t) / float(self._stride[0]))
+        out_h = np.ceil(float(h) / float(self._stride[1]))
+        out_w = np.ceil(float(w) / float(self._stride[2]))
+        pad_t = self.compute_pad(0, t)
+        pad_h = self.compute_pad(1, h)
+        pad_w = self.compute_pad(2, w)
+
+        pad_t_f = pad_t // 2
+        pad_t_b = pad_t - pad_t_f
+        pad_h_f = pad_h // 2
+        pad_h_b = pad_h - pad_h_f
+        pad_w_f = pad_w // 2
+        pad_w_b = pad_w - pad_w_f
+
+        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
+        x = F.pad(x, pad)
+
+        x = self.conv3d(x)
+        if self._use_batch_norm:
+            x = self.bn(x)
+        if self._activation_fn is not None:
+            x = self._activation_fn(x)
+        return x
+
+
+
+class InceptionModule(nn.Module):
+    def __init__(self, in_channels, out_channels, name):
+        super(InceptionModule, self).__init__()
+
+        self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0,
+                         name=name+'/Branch_0/Conv3d_0a_1x1')
+        self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_1/Conv3d_0a_1x1')
+        self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3],
+                          name=name+'/Branch_1/Conv3d_0b_3x3')
+        self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_2/Conv3d_0a_1x1')
+        self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3],
+                          name=name+'/Branch_2/Conv3d_0b_3x3')
+        self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],
+                                stride=(1, 1, 1), padding=0)
+        self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0,
+                          name=name+'/Branch_3/Conv3d_0b_1x1')
+        self.name = name
+
+    def forward(self, x):
+        b0 = self.b0(x)
+        b1 = self.b1b(self.b1a(x))
+        b2 = self.b2b(self.b2a(x))
+        b3 = self.b3b(self.b3a(x))
+        return torch.cat([b0,b1,b2,b3], dim=1)
+
+
+class InceptionI3d(nn.Module):
+    """Inception-v1 I3D architecture.
+    The model is introduced in:
+        Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset
+        Joao Carreira, Andrew Zisserman
+        https://arxiv.org/pdf/1705.07750v1.pdf.
+    See also the Inception architecture, introduced in:
+        Going deeper with convolutions
+        Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,
+        Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.
+        http://arxiv.org/pdf/1409.4842v1.pdf.
+    """
+
+    # Endpoints of the model in order. During construction, all the endpoints up
+    # to a designated `final_endpoint` are returned in a dictionary as the
+    # second return value.
+    VALID_ENDPOINTS = (
+        'Conv3d_1a_7x7',
+        'MaxPool3d_2a_3x3',
+        'Conv3d_2b_1x1',
+        'Conv3d_2c_3x3',
+        'MaxPool3d_3a_3x3',
+        'Mixed_3b',
+        'Mixed_3c',
+        'MaxPool3d_4a_3x3',
+        'Mixed_4b',
+        'Mixed_4c',
+        'Mixed_4d',
+        'Mixed_4e',
+        'Mixed_4f',
+        'MaxPool3d_5a_2x2',
+        'Mixed_5b',
+        'Mixed_5c',
+        'Logits',
+        'Predictions',
+    )
+
+    def __init__(self, num_classes=400, spatial_squeeze=True,
+                 final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5):
+        """Initializes I3D model instance.
+        Args:
+          num_classes: The number of outputs in the logit layer (default 400, which
+              matches the Kinetics dataset).
+          spatial_squeeze: Whether to squeeze the spatial dimensions for the logits
+              before returning (default True).
+          final_endpoint: The model contains many possible endpoints.
+              `final_endpoint` specifies the last endpoint for the model to be built
+              up to. In addition to the output at `final_endpoint`, all the outputs
+              at endpoints up to `final_endpoint` will also be returned, in a
+              dictionary. `final_endpoint` must be one of
+              InceptionI3d.VALID_ENDPOINTS (default 'Logits').
+          name: A string (optional). The name of this module.
+        Raises:
+          ValueError: if `final_endpoint` is not recognized.
+        """
+
+        if final_endpoint not in self.VALID_ENDPOINTS:
+            raise ValueError('Unknown final endpoint %s' % final_endpoint)
+
+        super(InceptionI3d, self).__init__()
+        self._num_classes = num_classes
+        self._spatial_squeeze = spatial_squeeze
+        self._final_endpoint = final_endpoint
+        self.logits = None
+
+        if self._final_endpoint not in self.VALID_ENDPOINTS:
+            raise ValueError('Unknown final endpoint %s' % self._final_endpoint)
+
+        self.end_points = {}
+        end_point = 'Conv3d_1a_7x7'
+        self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7],
+                                            stride=(2, 2, 2), padding=(3,3,3),  name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_2a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Conv3d_2b_1x1'
+        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0,
+                                       name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Conv3d_2c_3x3'
+        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1,
+                                       name=name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_3a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_3b'
+        self.end_points[end_point] = InceptionModule(192, [64,96,128,16,32,32], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_3c'
+        self.end_points[end_point] = InceptionModule(256, [128,128,192,32,96,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_4a_3x3'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4b'
+        self.end_points[end_point] = InceptionModule(128+192+96+64, [192,96,208,16,48,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4c'
+        self.end_points[end_point] = InceptionModule(192+208+48+64, [160,112,224,24,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4d'
+        self.end_points[end_point] = InceptionModule(160+224+64+64, [128,128,256,24,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4e'
+        self.end_points[end_point] = InceptionModule(128+256+64+64, [112,144,288,32,64,64], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_4f'
+        self.end_points[end_point] = InceptionModule(112+288+64+64, [256,160,320,32,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'MaxPool3d_5a_2x2'
+        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2),
+                                                             padding=0)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_5b'
+        self.end_points[end_point] = InceptionModule(256+320+128+128, [256,160,320,32,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Mixed_5c'
+        self.end_points[end_point] = InceptionModule(256+320+128+128, [384,192,384,48,128,128], name+end_point)
+        if self._final_endpoint == end_point: return
+
+        end_point = 'Logits'
+        self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7],
+                                     stride=(1, 1, 1))
+        self.dropout = nn.Dropout(dropout_keep_prob)
+        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
+                             kernel_shape=[1, 1, 1],
+                             padding=0,
+                             activation_fn=None,
+                             use_batch_norm=False,
+                             use_bias=True,
+                             name='logits')
+
+        self.build()
+
+
+    def replace_logits(self, num_classes):
+        self._num_classes = num_classes
+        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
+                             kernel_shape=[1, 1, 1],
+                             padding=0,
+                             activation_fn=None,
+                             use_batch_norm=False,
+                             use_bias=True,
+                             name='logits')
+
+
+    def build(self):
+        for k in self.end_points.keys():
+            self.add_module(k, self.end_points[k])
+
+    def forward(self, x):
+        for end_point in self.VALID_ENDPOINTS:
+            if end_point in self.end_points:
+                x = self._modules[end_point](x) # use _modules to work with dataparallel
+
+        x = self.logits(self.dropout(self.avg_pool(x)))
+        if self._spatial_squeeze:
+            logits = x.squeeze(3).squeeze(3)
+        logits = logits.mean(dim=2)
+        # logits is batch X time X classes, which is what we want to work with
+        return logits
+
+
+    def extract_features(self, x):
+        for end_point in self.VALID_ENDPOINTS:
+            if end_point in self.end_points:
+                x = self._modules[end_point](x)
+        return self.avg_pool(x)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/script/cal_clip_score.sh b/opensora/models/causalvideovae/eval/script/cal_clip_score.sh
new file mode 100755
index 000000000..8bc241c15
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/script/cal_clip_score.sh
@@ -0,0 +1,23 @@
+# clip_score cross modality
+python eval_clip_score.py \
+    --real_path path/to/image \
+    --generated_path path/to/text \
+    --batch-size 50 \
+    --device "cuda"
+
+# clip_score within the same modality
+python eval_clip_score.py \
+    --real_path path/to/textA \
+    --generated_path path/to/textB \
+    --real_flag txt \
+    --generated_flag txt \
+    --batch-size 50 \
+    --device "cuda"
+
+python eval_clip_score.py \
+    --real_path path/to/imageA \
+    --generated_path path/to/imageB \
+    --real_flag img \
+    --generated_flag img \
+    --batch-size 50 \
+    --device "cuda"
diff --git a/opensora/models/causalvideovae/eval/script/cal_fvd.sh b/opensora/models/causalvideovae/eval/script/cal_fvd.sh
new file mode 100755
index 000000000..e344242a0
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/script/cal_fvd.sh
@@ -0,0 +1,9 @@
+python eval_common_metric.py \
+    --real_video_dir path/to/imageA\
+    --generated_video_dir path/to/imageB \
+    --batch_size 10 \
+    --crop_size 64 \
+    --num_frames 20 \
+    --device 'cuda' \
+    --metric 'fvd' \
+    --fvd_method 'styleganv'
diff --git a/opensora/models/causalvideovae/eval/script/cal_lpips.sh b/opensora/models/causalvideovae/eval/script/cal_lpips.sh
new file mode 100755
index 000000000..9ebaf26e8
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/script/cal_lpips.sh
@@ -0,0 +1,8 @@
+python eval_common_metric.py \
+    --real_video_dir path/to/imageA\
+    --generated_video_dir path/to/imageB \
+    --batch_size 10 \
+    --num_frames 20 \
+    --crop_size 64 \
+    --device 'cuda' \
+    --metric 'lpips'
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/script/cal_psnr.sh b/opensora/models/causalvideovae/eval/script/cal_psnr.sh
new file mode 100755
index 000000000..a60a5d1c1
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/script/cal_psnr.sh
@@ -0,0 +1,9 @@
+
+python eval_common_metric.py \
+    --real_video_dir /data/xiaogeng_liu/data/video1 \
+    --generated_video_dir /data/xiaogeng_liu/data/video2 \
+    --batch_size 10 \
+    --num_frames 20 \
+    --crop_size 64 \
+    --device 'cuda' \
+    --metric 'psnr'
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/eval/script/cal_ssim.sh b/opensora/models/causalvideovae/eval/script/cal_ssim.sh
new file mode 100755
index 000000000..404d8da9a
--- /dev/null
+++ b/opensora/models/causalvideovae/eval/script/cal_ssim.sh
@@ -0,0 +1,8 @@
+python eval_common_metric.py \
+    --real_video_dir /data/xiaogeng_liu/data/video1 \
+    --generated_video_dir /data/xiaogeng_liu/data/video2 \
+    --batch_size 10 \
+    --num_frames 20 \
+    --crop_size 64 \
+    --device 'cuda' \
+    --metric 'ssim'
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/__init__.py b/opensora/models/causalvideovae/model/__init__.py
new file mode 100755
index 000000000..7ca9d1546
--- /dev/null
+++ b/opensora/models/causalvideovae/model/__init__.py
@@ -0,0 +1,4 @@
+from .causal_vae import (
+    CausalVAEModel, CausalVAEModelWrapper
+)
+from .ema_model import EMA
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/causal_vae/__init__.py b/opensora/models/causalvideovae/model/causal_vae/__init__.py
new file mode 100755
index 000000000..c1aeb800d
--- /dev/null
+++ b/opensora/models/causalvideovae/model/causal_vae/__init__.py
@@ -0,0 +1,29 @@
+from .modeling_causalvae import CausalVAEModel
+
+from einops import rearrange
+from torch import nn
+
+class CausalVAEModelWrapper(nn.Module):
+    def __init__(self, model_path, subfolder=None, cache_dir=None, use_ema=False, **kwargs):
+        super(CausalVAEModelWrapper, self).__init__()
+        # if os.path.exists(ckpt):
+        # self.vae = CausalVAEModel.load_from_checkpoint(ckpt)
+        self.vae = CausalVAEModel.from_pretrained(model_path, subfolder=subfolder, cache_dir=cache_dir, **kwargs)
+        if use_ema:
+            self.vae.init_from_ema(model_path)
+            self.vae = self.vae.ema
+    def encode(self, x):  # b c t h w
+        # x = self.vae.encode(x).sample()
+        x = self.vae.encode(x).sample().mul_(0.18215)
+        return x
+    def decode(self, x):
+        # x = self.vae.decode(x)
+        x = self.vae.decode(x / 0.18215)
+        x = rearrange(x, 'b c t h w -> b t c h w').contiguous()
+        return x
+
+    def dtype(self):
+        return self.vae.dtype
+    #
+    # def device(self):
+    #     return self.vae.device
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py b/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py
new file mode 100755
index 000000000..a6f20cc2b
--- /dev/null
+++ b/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py
@@ -0,0 +1,628 @@
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+from ..modeling_videobase import VideoBaseAE
+from ..modules import Normalize
+from ..modules.ops import nonlinearity
+from typing import List, Tuple
+import torch.nn as nn
+
+from ..utils.module_utils import resolve_str_to_obj, Module
+from ..utils.distrib_utils import DiagonalGaussianDistribution
+from ..utils.scheduler_utils import cosine_scheduler
+from ...utils.utils import custom_to_video
+
+import torch
+from diffusers.configuration_utils import register_to_config
+from copy import deepcopy
+import os
+import glob
+
+import numpy as np
+from ...eval.cal_psnr import calculate_psnr
+from decord import VideoReader, cpu
+from pytorchvideo.transforms import ShortSideScale
+from torchvision.io import read_video
+from torchvision.transforms import Lambda, Compose
+from torchvision.transforms._transforms_video import CenterCropVideo
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        z_channels: int,
+        hidden_size: int,
+        hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
+        attn_resolutions: Tuple[int] = (16,),
+        conv_in: Module = "Conv2d",
+        conv_out: Module = "CasualConv3d",
+        attention: Module = "AttnBlock",
+        resnet_blocks: Tuple[Module] = (
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock3D",
+        ),
+        spatial_downsample: Tuple[Module] = (
+            "Downsample",
+            "Downsample",
+            "Downsample",
+            "",
+        ),
+        temporal_downsample: Tuple[Module] = ("", "", "TimeDownsampleRes2x", ""),
+        mid_resnet: Module = "ResnetBlock3D",
+        dropout: float = 0.0,
+        resolution: int = 256,
+        num_res_blocks: int = 2,
+        double_z: bool = True,
+    ) -> None:
+        super().__init__()
+        assert len(resnet_blocks) == len(hidden_size_mult), print(
+            hidden_size_mult, resnet_blocks
+        )
+        # ---- Config ----
+        self.num_resolutions = len(hidden_size_mult)
+        self.resolution = resolution
+        self.num_res_blocks = num_res_blocks
+        
+        # ---- In ----
+        self.conv_in = resolve_str_to_obj(conv_in)(
+            3, hidden_size, kernel_size=3, stride=1, padding=1
+        )
+
+        # ---- Downsample ----
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(hidden_size_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = hidden_size * in_ch_mult[i_level]
+            block_out = hidden_size * hidden_size_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    resolve_str_to_obj(resnet_blocks[i_level])(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(resolve_str_to_obj(attention)(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if spatial_downsample[i_level]:
+                down.downsample = resolve_str_to_obj(spatial_downsample[i_level])(
+                    block_in, block_in
+                )
+                curr_res = curr_res // 2
+            if temporal_downsample[i_level]:
+                down.time_downsample = resolve_str_to_obj(temporal_downsample[i_level])(
+                    block_in, block_in
+                )
+            self.down.append(down)
+
+        # ---- Mid ----
+        self.mid = nn.Module()
+        self.mid.block_1 = resolve_str_to_obj(mid_resnet)(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = resolve_str_to_obj(attention)(block_in)
+        self.mid.block_2 = resolve_str_to_obj(mid_resnet)(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=dropout,
+        )
+        # ---- Out ----
+        self.norm_out = Normalize(block_in)
+        self.conv_out = resolve_str_to_obj(conv_out)(
+            block_in,
+            2 * z_channels if double_z else z_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, x):
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if hasattr(self.down[i_level], "downsample"):
+                hs.append(self.down[i_level].downsample(hs[-1]))
+            if hasattr(self.down[i_level], "time_downsample"):
+                hs_down = self.down[i_level].time_downsample(hs[-1])
+                hs.append(hs_down)
+
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        if npu_config is None:
+            h = self.norm_out(h)
+        else:
+            h = npu_config.run_group_norm(self.norm_out, h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        z_channels: int,
+        hidden_size: int,
+        hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
+        attn_resolutions: Tuple[int] = (16,),
+        conv_in: Module = "Conv2d",
+        conv_out: Module = "CasualConv3d",
+        attention: Module = "AttnBlock",
+        resnet_blocks: Tuple[Module] = (
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+        ),
+        spatial_upsample: Tuple[Module] = (
+            "",
+            "SpatialUpsample2x",
+            "SpatialUpsample2x",
+            "SpatialUpsample2x",
+        ),
+        temporal_upsample: Tuple[Module] = ("", "", "", "TimeUpsampleRes2x"),
+        mid_resnet: Module = "ResnetBlock3D",
+        dropout: float = 0.0,
+        resolution: int = 256,
+        num_res_blocks: int = 2,
+    ):
+        super().__init__()
+        # ---- Config ----
+        self.num_resolutions = len(hidden_size_mult)
+        self.resolution = resolution
+        self.num_res_blocks = num_res_blocks
+
+        # ---- In ----
+        block_in = hidden_size * hidden_size_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.conv_in = resolve_str_to_obj(conv_in)(
+            z_channels, block_in, kernel_size=3, padding=1
+        )
+
+        # ---- Mid ----
+        self.mid = nn.Module()
+        self.mid.block_1 = resolve_str_to_obj(mid_resnet)(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=dropout,
+        )
+        self.mid.attn_1 = resolve_str_to_obj(attention)(block_in)
+        self.mid.block_2 = resolve_str_to_obj(mid_resnet)(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=dropout,
+        )
+
+        # ---- Upsample ----
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = hidden_size * hidden_size_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    resolve_str_to_obj(resnet_blocks[i_level])(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        dropout=dropout,
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(resolve_str_to_obj(attention)(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if spatial_upsample[i_level]:
+                up.upsample = resolve_str_to_obj(spatial_upsample[i_level])(
+                    block_in, block_in
+                )
+                curr_res = curr_res * 2
+            if temporal_upsample[i_level]:
+                up.time_upsample = resolve_str_to_obj(temporal_upsample[i_level])(
+                    block_in, block_in
+                )
+            self.up.insert(0, up)
+
+        # ---- Out ----
+        self.norm_out = Normalize(block_in)
+        self.conv_out = resolve_str_to_obj(conv_out)(
+            block_in, 3, kernel_size=3, padding=1
+        )
+
+    def forward(self, z):
+        h = self.conv_in(z)
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if hasattr(self.up[i_level], "upsample"):
+                h = self.up[i_level].upsample(h)
+            if hasattr(self.up[i_level], "time_upsample"):
+                h = self.up[i_level].time_upsample(h)
+        if npu_config is None:
+            h = self.norm_out(h)
+        else:
+            h = npu_config.run_group_norm(self.norm_out, h)
+        h = nonlinearity(h)
+        if npu_config is None:
+            h = self.conv_out(h)
+        else:
+            h_dtype = h.dtype
+            h = npu_config.run_conv3d(self.conv_out, h, h_dtype)
+        return h
+
+
+class CausalVAEModel(VideoBaseAE):
+
+    @register_to_config
+    def __init__(
+        self,
+        hidden_size: int = 128,
+        z_channels: int = 4,
+        hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
+        attn_resolutions: Tuple[int] = [],
+        dropout: float = 0.0,
+        resolution: int = 256,
+        double_z: bool = True,
+        embed_dim: int = 4,
+        num_res_blocks: int = 2,
+        q_conv: str = "CausalConv3d",
+        encoder_conv_in: Module = "CausalConv3d",
+        encoder_conv_out: Module = "CausalConv3d",
+        encoder_attention: Module = "AttnBlock3D",
+        encoder_resnet_blocks: Tuple[Module] = (
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+        ),
+        encoder_spatial_downsample: Tuple[Module] = (
+            "SpatialDownsample2x",
+            "SpatialDownsample2x",
+            "SpatialDownsample2x",
+            "",
+        ),
+        encoder_temporal_downsample: Tuple[Module] = (
+            "",
+            "TimeDownsample2x",
+            "TimeDownsample2x",
+            "",
+        ),
+        encoder_mid_resnet: Module = "ResnetBlock3D",
+        decoder_conv_in: Module = "CausalConv3d",
+        decoder_conv_out: Module = "CausalConv3d",
+        decoder_attention: Module = "AttnBlock3D",
+        decoder_resnet_blocks: Tuple[Module] = (
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+        ),
+        decoder_spatial_upsample: Tuple[Module] = (
+            "",
+            "SpatialUpsample2x",
+            "SpatialUpsample2x",
+            "SpatialUpsample2x",
+        ),
+        decoder_temporal_upsample: Tuple[Module] = ("", "", "TimeUpsample2x", "TimeUpsample2x"),
+        decoder_mid_resnet: Module = "ResnetBlock3D",
+        use_quant_layer: bool = True
+    ) -> None:
+        super().__init__()
+        
+        
+        self.tile_sample_min_size = 256
+        self.tile_sample_min_size_t = 33
+        self.tile_latent_min_size = int(self.tile_sample_min_size / (2 ** (len(hidden_size_mult) - 1)))
+        
+        # t_down_ratio = [i for i in encoder_temporal_downsample if len(i) > 0]
+        # self.tile_latent_min_size_t = int((self.tile_sample_min_size_t-1) / (2 ** len(t_down_ratio))) + 1
+        self.tile_latent_min_size_t = 16
+        self.tile_overlap_factor = 0.125
+        self.use_tiling = False
+        
+        self.use_quant_layer = use_quant_layer
+
+        self.encoder = Encoder(
+            z_channels=z_channels,
+            hidden_size=hidden_size,
+            hidden_size_mult=hidden_size_mult,
+            attn_resolutions=attn_resolutions,
+            conv_in=encoder_conv_in,
+            conv_out=encoder_conv_out,
+            attention=encoder_attention,
+            resnet_blocks=encoder_resnet_blocks,
+            spatial_downsample=encoder_spatial_downsample,
+            temporal_downsample=encoder_temporal_downsample,
+            mid_resnet=encoder_mid_resnet,
+            dropout=dropout,
+            resolution=resolution,
+            num_res_blocks=num_res_blocks,
+            double_z=double_z,
+        )
+
+        self.decoder = Decoder(
+            z_channels=z_channels,
+            hidden_size=hidden_size,
+            hidden_size_mult=hidden_size_mult,
+            attn_resolutions=attn_resolutions,
+            conv_in=decoder_conv_in,
+            conv_out=decoder_conv_out,
+            attention=decoder_attention,
+            resnet_blocks=decoder_resnet_blocks,
+            spatial_upsample=decoder_spatial_upsample,
+            temporal_upsample=decoder_temporal_upsample,
+            mid_resnet=decoder_mid_resnet,
+            dropout=dropout,
+            resolution=resolution,
+            num_res_blocks=num_res_blocks,
+        )
+        if self.use_quant_layer:
+            quant_conv_cls = resolve_str_to_obj(q_conv)
+            self.quant_conv = quant_conv_cls(2 * z_channels, 2 * embed_dim, 1)
+            self.post_quant_conv = quant_conv_cls(embed_dim, z_channels, 1)
+
+    def get_encoder(self):
+        if self.use_quant_layer:
+            return [self.quant_conv, self.encoder]
+        return [self.encoder]
+            
+    def get_decoder(self):
+        if self.use_quant_layer:
+            return [self.post_quant_conv, self.decoder]
+        return [self.decoder]
+    
+    def encode(self, x):
+        if self.use_tiling and (
+            x.shape[-1] > self.tile_sample_min_size
+            or x.shape[-2] > self.tile_sample_min_size
+            or x.shape[-3] > self.tile_sample_min_size_t
+        ):
+            return self.tiled_encode(x)
+        h = self.encoder(x)
+        if self.use_quant_layer:
+            h = self.quant_conv(h)
+        posterior = DiagonalGaussianDistribution(h)
+        return posterior
+
+    def decode(self, z):
+        if self.use_tiling and (
+            z.shape[-1] > self.tile_latent_min_size
+            or z.shape[-2] > self.tile_latent_min_size
+            or z.shape[-3] > self.tile_latent_min_size_t
+        ):
+            return self.tiled_decode(z)
+        if self.use_quant_layer:
+            z = self.post_quant_conv(z)
+        dec = self.decoder(z)
+        return dec
+
+    def forward(self, input, sample_posterior=True):
+        posterior = self.encode(input)
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        return dec, posterior
+    
+    def on_train_start(self):
+        self.ema = deepcopy(self) if self.save_ema==True else None
+    
+    def get_last_layer(self):
+        if hasattr(self.decoder.conv_out, "conv"):
+            return self.decoder.conv_out.conv.weight
+        else:
+            return self.decoder.conv_out.weight
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
+                1 - y / blend_extent
+            ) + b[:, :, :, y, :] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[4], b.shape[4], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
+                1 - x / blend_extent
+            ) + b[:, :, :, :, x] * (x / blend_extent)
+        return b
+
+    def tiled_encode(self, x):
+        t = x.shape[2]
+        t_chunk_idx = [i for i in range(0, t, self.tile_sample_min_size_t-1)]
+        if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
+            t_chunk_start_end = [[0, t]]
+        else:
+            t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i+1]+1] for i in range(len(t_chunk_idx)-1)]
+            if t_chunk_start_end[-1][-1] > t:
+                t_chunk_start_end[-1][-1] = t
+            elif t_chunk_start_end[-1][-1] < t:
+                last_start_end = [t_chunk_idx[-1], t]
+                t_chunk_start_end.append(last_start_end)
+        moments = []
+        for idx, (start, end) in enumerate(t_chunk_start_end):
+            chunk_x = x[:, :, start: end]
+            if idx != 0:
+                moment = self.tiled_encode2d(chunk_x, return_moments=True)[:, :, 1:]
+            else:
+                moment = self.tiled_encode2d(chunk_x, return_moments=True)
+            moments.append(moment)
+        moments = torch.cat(moments, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+        return posterior
+    
+    def tiled_decode(self, x):
+        t = x.shape[2]
+        t_chunk_idx = [i for i in range(0, t, self.tile_latent_min_size_t-1)]
+        if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
+            t_chunk_start_end = [[0, t]]
+        else:
+            t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i+1]+1] for i in range(len(t_chunk_idx)-1)]
+            if t_chunk_start_end[-1][-1] > t:
+                t_chunk_start_end[-1][-1] = t
+            elif t_chunk_start_end[-1][-1] < t:
+                last_start_end = [t_chunk_idx[-1], t]
+                t_chunk_start_end.append(last_start_end)
+        dec_ = []
+        for idx, (start, end) in enumerate(t_chunk_start_end):
+            chunk_x = x[:, :, start: end]
+            if idx != 0:
+                dec = self.tiled_decode2d(chunk_x)[:, :, 1:]
+            else:
+                dec = self.tiled_decode2d(chunk_x)
+            dec_.append(dec)
+        dec_ = torch.cat(dec_, dim=2)
+        return dec_
+
+    def tiled_encode2d(self, x, return_moments=False):
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[3], overlap_size):
+            row = []
+            for j in range(0, x.shape[4], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.use_quant_layer:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=4))
+
+        moments = torch.cat(result_rows, dim=3)
+        posterior = DiagonalGaussianDistribution(moments)
+        if return_moments:
+            return moments
+        return posterior
+
+    def tiled_decode2d(self, z):
+
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[3], overlap_size):
+            row = []
+            for j in range(0, z.shape[4], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                if self.use_quant_layer:
+                    tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=4))
+
+        dec = torch.cat(result_rows, dim=3)
+        return dec
+
+    def enable_tiling(self, use_tiling: bool = True):
+        self.use_tiling = use_tiling
+
+    def disable_tiling(self):
+        self.enable_tiling(False)
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")
+        print("init from " + path)
+        
+        if "ema_state_dict" in sd and len(sd['ema_state_dict']) > 0 and os.environ.get("NOT_USE_EMA_MODEL", 0) == 0:
+            print("Load from ema model!")
+            sd = sd["ema_state_dict"]
+            sd = {key.replace("module.", ""): value for key, value in sd.items()}
+        elif "state_dict" in sd:
+            print("Load from normal model!")
+            if "gen_model" in sd["state_dict"]:
+                sd = sd["state_dict"]["gen_model"]
+            else:
+                sd = sd["state_dict"]
+                
+        keys = list(sd.keys())
+        
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        
+        miss, unexpected = self.load_state_dict(sd, strict=False)
+        assert len(miss) == 0, f"miss key: {miss}"
+        if len(unexpected) > 0:
+            for i in unexpected:
+                assert 'loss' in i, "unexpected key: {i}"
+
+
+
+
diff --git a/opensora/models/causalvideovae/model/configuration_videobase.py b/opensora/models/causalvideovae/model/configuration_videobase.py
new file mode 100755
index 000000000..f25a25c5b
--- /dev/null
+++ b/opensora/models/causalvideovae/model/configuration_videobase.py
@@ -0,0 +1,44 @@
+import json
+import yaml
+from typing import TypeVar, Dict, Any
+from diffusers import ConfigMixin
+
+T = TypeVar('T', bound='VideoBaseConfiguration')
+class VideoBaseConfiguration(ConfigMixin):
+    config_name = "VideoBaseConfiguration"
+    _nested_config_fields: Dict[str, Any] = {}
+    
+    def __init__(self, **kwargs):
+        pass
+    
+    def to_dict(self) -> Dict[str, Any]:
+        d = {}
+        for key, value in vars(self).items():
+            if isinstance(value, VideoBaseConfiguration):
+                d[key] = value.to_dict()  # Serialize nested VideoBaseConfiguration instances
+            elif isinstance(value, tuple):
+                d[key] = list(value)
+            else:
+                d[key] = value
+        return d
+    
+    def to_yaml_file(self, yaml_path: str):
+        with open(yaml_path, 'w') as yaml_file:
+            yaml.dump(self.to_dict(), yaml_file, default_flow_style=False)
+    
+    @classmethod
+    def load_from_yaml(cls: T, yaml_path: str) -> T:
+        with open(yaml_path, 'r') as yaml_file:
+            config_dict = yaml.safe_load(yaml_file)
+        for field, field_type in cls._nested_config_fields.items():
+            if field in config_dict:
+                config_dict[field] = field_type.load_from_dict(config_dict[field])
+        return cls(**config_dict)
+
+    @classmethod
+    def load_from_dict(cls: T, config_dict: Dict[str, Any]) -> T:
+        # Process nested configuration objects
+        for field, field_type in cls._nested_config_fields.items():
+            if field in config_dict:
+                config_dict[field] = field_type.load_from_dict(config_dict[field])
+        return cls(**config_dict)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/dataset_videobase.py b/opensora/models/causalvideovae/model/dataset_videobase.py
new file mode 100755
index 000000000..564bc1f72
--- /dev/null
+++ b/opensora/models/causalvideovae/model/dataset_videobase.py
@@ -0,0 +1,107 @@
+import os.path as osp
+import random
+from glob import glob
+
+from torchvision import transforms
+import numpy as np
+import torch
+import torch.utils.data as data
+import torch.nn.functional as F
+from torchvision.transforms import Lambda
+
+from ..dataset.transform import ToTensorVideo, CenterCropVideo
+from ..utils.dataset_utils import DecordInit
+
+def TemporalRandomCrop(total_frames, size):
+    """
+    Performs a random temporal crop on a video sequence.
+
+    This function randomly selects a continuous frame sequence of length `size` from a video sequence.
+    `total_frames` indicates the total number of frames in the video sequence, and `size` represents the length of the frame sequence to be cropped.
+
+    Parameters:
+    - total_frames (int): The total number of frames in the video sequence.
+    - size (int): The length of the frame sequence to be cropped.
+
+    Returns:
+    - (int, int): A tuple containing two integers. The first integer is the starting frame index of the cropped sequence,
+                  and the second integer is the ending frame index (inclusive) of the cropped sequence.
+    """
+    rand_end = max(0, total_frames - size - 1)
+    begin_index = random.randint(0, rand_end)
+    end_index = min(begin_index + size, total_frames)
+    return begin_index, end_index
+
+def resize(x, resolution):
+    height, width = x.shape[-2:]
+    resolution = min(2 * resolution, height, width)
+    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
+
+class VideoDataset(data.Dataset):
+    """ Generic dataset for videos files stored in folders
+    Returns BCTHW videos in the range [-0.5, 0.5] """
+    video_exts = ['avi', 'mp4', 'webm']
+    def __init__(self, video_folder, sequence_length, image_folder=None, train=True, resolution=64, sample_rate=1, dynamic_sample=True):
+
+        self.train = train
+        self.sequence_length = sequence_length
+        self.sample_rate = sample_rate
+        self.resolution = resolution
+        self.v_decoder = DecordInit()
+        self.video_folder = video_folder
+        self.dynamic_sample = dynamic_sample
+
+        self.transform = transforms.Compose([
+            ToTensorVideo(),
+            # Lambda(lambda x: resize(x, self.resolution)),
+            CenterCropVideo(self.resolution),
+            Lambda(lambda x: 2.0 * x - 1.0)
+        ])
+        print('Building datasets...')
+        self.samples = self._make_dataset()
+
+    def _make_dataset(self):
+        samples = []
+        samples += sum([glob(osp.join(self.video_folder, '**', f'*.{ext}'), recursive=True)
+                            for ext in self.video_exts], [])
+        return samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        video_path = self.samples[idx]
+        try:
+            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
+            return dict(video=video, label="")
+        except Exception as e:
+            print(f'Error with {e}, {video_path}')
+            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
+        if self.dynamic_sample:
+            sample_rate = random.randint(1, self.sample_rate)
+        else:
+            sample_rate = self.sample_rate
+        size = self.sequence_length * sample_rate
+        start_frame_ind, end_frame_ind = TemporalRandomCrop(total_frames, size)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.sequence_length, 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
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/ema_model.py b/opensora/models/causalvideovae/model/ema_model.py
new file mode 100755
index 000000000..c4cf50172
--- /dev/null
+++ b/opensora/models/causalvideovae/model/ema_model.py
@@ -0,0 +1,31 @@
+class EMA:
+    def __init__(self, model, decay):
+        self.model = model
+        self.decay = decay
+        self.shadow = {}
+        self.backup = {}
+        
+    def register(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    def update(self):
+        for name, param in self.model.named_parameters():
+            if name in self.shadow:
+                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
+                self.shadow[name] = new_average.clone()
+
+    def apply_shadow(self):
+        for name, param in self.model.named_parameters():
+            if name in self.shadow:
+                self.backup[name] = param.data
+                param.data = self.shadow[name]
+
+    def restore(self):
+        for name, param in self.model.named_parameters():
+            if name in self.shadow:
+                param.data = self.backup[name]
+        self.backup = {}
+        
+        
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/losses/__init__.py b/opensora/models/causalvideovae/model/losses/__init__.py
new file mode 100755
index 000000000..07e681263
--- /dev/null
+++ b/opensora/models/causalvideovae/model/losses/__init__.py
@@ -0,0 +1 @@
+from .perceptual_loss import SimpleLPIPS, LPIPSWithDiscriminator, LPIPSWithDiscriminator3D
diff --git a/opensora/models/causalvideovae/model/losses/discriminator.py b/opensora/models/causalvideovae/model/losses/discriminator.py
new file mode 100755
index 000000000..247a61070
--- /dev/null
+++ b/opensora/models/causalvideovae/model/losses/discriminator.py
@@ -0,0 +1,181 @@
+import functools
+import torch.nn as nn
+from ..modules.normalize import ActNorm
+from ..modules.conv import CausalConv3d
+from einops import rearrange
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+def weights_init_conv(m):
+    if hasattr(m, 'conv'):
+        m = m.conv
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix
+        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm2d
+        else:
+            norm_layer = ActNorm
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [
+            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)
+    
+class NLayerDiscriminator3D(nn.Module):
+    """Defines a 3D PatchGAN discriminator as in Pix2Pix but for 3D inputs."""
+    def __init__(self, input_nc=1, ndf=64, n_layers=3, use_actnorm=False):
+        """
+        Construct a 3D PatchGAN discriminator
+
+        Parameters:
+            input_nc (int)  -- the number of channels in input volumes
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            use_actnorm (bool) -- flag to use actnorm instead of batchnorm
+        """
+        super(NLayerDiscriminator3D, self).__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm3d
+        else:
+            raise NotImplementedError("Not implemented.")
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func != nn.BatchNorm3d
+        else:
+            use_bias = norm_layer != nn.BatchNorm3d
+
+        kw = 3
+        padw = 1
+        sequence = [nn.Conv3d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv3d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=(kw, kw, kw), stride=(2 if n==1 else 1,2,2), padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv3d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=(kw, kw, kw), stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [nn.Conv3d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)
+
+
+
+
+
+# class NLayerDiscriminator3D(nn.Module):
+#     """Defines a 3D PatchGAN discriminator as in Pix2Pix but for 3D inputs."""
+#     def __init__(self, input_nc=1, ndf=64, n_layers=3, use_actnorm=False):
+#         """
+#         Construct a 3D PatchGAN discriminator
+
+#         Parameters:
+#             input_nc (int)  -- the number of channels in input volumes
+#             ndf (int)       -- the number of filters in the last conv layer
+#             n_layers (int)  -- the number of conv layers in the discriminator
+#             use_actnorm (bool) -- flag to use actnorm instead of batchnorm
+#         """
+#         super(NLayerDiscriminator3D, self).__init__()
+#         if not use_actnorm:
+#             norm_layer = nn.BatchNorm3d
+#         else:
+#             raise NotImplementedError("Not implemented.")
+#         if type(norm_layer) == functools.partial:
+#             use_bias = norm_layer.func != nn.BatchNorm3d
+#         else:
+#             use_bias = norm_layer != nn.BatchNorm3d
+
+#         kw = 4
+#         padw = 1
+#         sequence = [CausalConv3d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+#         nf_mult = 1
+#         nf_mult_prev = 1
+#         for n in range(1, n_layers):  # gradually increase the number of filters
+#             nf_mult_prev = nf_mult
+#             nf_mult = min(2 ** n, 8)
+#             sequence += [
+#                 CausalConv3d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=(kw, kw, kw), stride=(2 if n==1 else 1,2,2), padding=padw, bias=use_bias),
+#                 norm_layer(ndf * nf_mult),
+#                 nn.LeakyReLU(0.2, True)
+#             ]
+
+#         nf_mult_prev = nf_mult
+#         nf_mult = min(2 ** n_layers, 8)
+#         sequence += [
+#             CausalConv3d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=(kw, kw, kw), stride=1, padding=padw, bias=use_bias),
+#             norm_layer(ndf * nf_mult),
+#             nn.LeakyReLU(0.2, True)
+#         ]
+
+#         sequence += [CausalConv3d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+#         self.main = nn.Sequential(*sequence)
+
+#     def forward(self, input):
+#         """Standard forward."""
+#         return self.main(input)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/losses/lpips.py b/opensora/models/causalvideovae/model/losses/lpips.py
new file mode 100755
index 000000000..b51a73881
--- /dev/null
+++ b/opensora/models/causalvideovae/model/losses/lpips.py
@@ -0,0 +1,120 @@
+"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from collections import namedtuple
+from ...utils.taming_download import get_ckpt_path
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name, ".cache/lpips")
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    @classmethod
+    def from_pretrained(cls, name="vgg_lpips"):
+        if name != "vgg_lpips":
+            raise NotImplementedError
+        model = cls()
+        ckpt = get_ckpt_path(name)
+        model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        return model
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x,eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True))
+    return x/(norm_factor+eps)
+
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2,3],keepdim=keepdim)
diff --git a/opensora/models/causalvideovae/model/losses/perceptual_loss.py b/opensora/models/causalvideovae/model/losses/perceptual_loss.py
new file mode 100755
index 000000000..eb99daae2
--- /dev/null
+++ b/opensora/models/causalvideovae/model/losses/perceptual_loss.py
@@ -0,0 +1,389 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .lpips import LPIPS
+from einops import rearrange
+from .discriminator import NLayerDiscriminator, weights_init, NLayerDiscriminator3D
+
+
+def hinge_d_loss(logits_real, logits_fake):
+    loss_real = torch.mean(F.relu(1.0 - logits_real))
+    loss_fake = torch.mean(F.relu(1.0 + logits_fake))
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def vanilla_d_loss(logits_real, logits_fake):
+    d_loss = 0.5 * (
+        torch.mean(torch.nn.functional.softplus(-logits_real))
+        + torch.mean(torch.nn.functional.softplus(logits_fake))
+    )
+    return d_loss
+
+
+def hinge_d_loss_with_exemplar_weights(logits_real, logits_fake, weights):
+    assert weights.shape[0] == logits_real.shape[0] == logits_fake.shape[0]
+    loss_real = torch.mean(F.relu(1.0 - logits_real), dim=[1, 2, 3])
+    loss_fake = torch.mean(F.relu(1.0 + logits_fake), dim=[1, 2, 3])
+    loss_real = (weights * loss_real).sum() / weights.sum()
+    loss_fake = (weights * loss_fake).sum() / weights.sum()
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def adopt_weight(weight, global_step, threshold=0, value=0.0):
+    if global_step < threshold:
+        weight = value
+    return weight
+
+
+def measure_perplexity(predicted_indices, n_embed):
+    # src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py
+    # eval cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally
+    encodings = F.one_hot(predicted_indices, n_embed).float().reshape(-1, n_embed)
+    avg_probs = encodings.mean(0)
+    perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp()
+    cluster_use = torch.sum(avg_probs > 0)
+    return perplexity, cluster_use
+
+
+def l1(x, y):
+    return torch.abs(x - y)
+
+
+def l2(x, y):
+    return torch.pow((x - y), 2)
+
+
+class LPIPSWithDiscriminator(nn.Module):
+    def __init__(
+        self,
+        disc_start,
+        logvar_init=0.0,
+        kl_weight=1.0,
+        pixelloss_weight=1.0,
+        perceptual_weight=1.0,
+        # --- Discriminator Loss ---
+        disc_num_layers=3,
+        disc_in_channels=3,
+        disc_factor=1.0,
+        disc_weight=1.0,
+        use_actnorm=False,
+        disc_conditional=False,
+        disc_loss="hinge",
+        loss_type: str = "l1"
+    ):
+
+        super().__init__()
+        assert disc_loss in ["hinge", "vanilla"]
+        self.kl_weight = kl_weight
+        self.pixel_weight = pixelloss_weight
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+        self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init)
+
+        self.discriminator = NLayerDiscriminator(
+            input_nc=disc_in_channels, n_layers=disc_num_layers, use_actnorm=use_actnorm
+        ).apply(weights_init)
+        self.discriminator_iter_start = disc_start
+        self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
+        self.disc_factor = disc_factor
+        self.discriminator_weight = disc_weight
+        self.disc_conditional = disc_conditional
+        self.loss_func = l1 if loss_type == "l1" else l2
+        
+    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+        layer = last_layer if last_layer is not None else self.last_layer[0]
+
+        nll_grads = torch.autograd.grad(nll_loss, layer, retain_graph=True)[0]
+        g_grads = torch.autograd.grad(g_loss, layer, retain_graph=True)[0]
+        
+        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+        d_weight = d_weight * self.discriminator_weight
+        return d_weight
+
+    def forward(
+        self,
+        inputs,
+        reconstructions,
+        posteriors,
+        optimizer_idx,
+        global_step,
+        split="train",
+        weights=None,
+        last_layer=None,
+        cond=None,
+    ):
+        # GAN Part
+        inputs = rearrange(inputs, "b c t h w -> (b t) c h w").contiguous()
+        reconstructions = rearrange(
+            reconstructions, "b c t h w -> (b t) c h w"
+        ).contiguous()
+        if optimizer_idx == 0:
+            rec_loss = self.loss_func(inputs, reconstructions)
+            if self.perceptual_weight > 0:
+                p_loss = self.perceptual_loss(inputs, reconstructions)
+                rec_loss = rec_loss + self.perceptual_weight * p_loss
+            nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+            weighted_nll_loss = nll_loss
+            if weights is not None:
+                weighted_nll_loss = weights * nll_loss
+            weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+            nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+            kl_loss = posteriors.kl()
+            kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+            
+            logits_fake = self.discriminator(reconstructions)
+            g_loss = -torch.mean(logits_fake)
+            if global_step >= self.discriminator_iter_start:
+                if self.disc_factor > 0.0:
+                    d_weight = self.calculate_adaptive_weight(
+                        nll_loss, g_loss, last_layer=last_layer
+                    )
+                else:
+                    d_weight = torch.tensor(1.0)
+            else:
+                d_weight = torch.tensor(0.0)
+                g_loss = torch.tensor(0.0, requires_grad=True)
+                
+            disc_factor = adopt_weight(
+                self.disc_factor, global_step, threshold=self.discriminator_iter_start
+            )
+            loss = (
+                weighted_nll_loss
+                + self.kl_weight * kl_loss
+                + d_weight * disc_factor * g_loss
+            )
+            log = {
+                "{}/total_loss".format(split): loss.clone().detach().mean(),
+                "{}/logvar".format(split): self.logvar.detach(),
+                "{}/kl_loss".format(split): kl_loss.detach().mean(),
+                "{}/nll_loss".format(split): nll_loss.detach().mean(),
+                "{}/rec_loss".format(split): rec_loss.detach().mean(),
+                "{}/d_weight".format(split): d_weight.detach(),
+                "{}/disc_factor".format(split): torch.tensor(disc_factor),
+                "{}/g_loss".format(split): g_loss.detach().mean(),
+            }
+            return loss, log
+
+        if optimizer_idx == 1:
+            if cond is None:
+                logits_real = self.discriminator(inputs.contiguous().detach())
+                logits_fake = self.discriminator(reconstructions.contiguous().detach())
+            else:
+                logits_real = self.discriminator(
+                    torch.cat((inputs.contiguous().detach(), cond), dim=1)
+                )
+                logits_fake = self.discriminator(
+                    torch.cat((reconstructions.contiguous().detach(), cond), dim=1)
+                )
+
+            disc_factor = adopt_weight(
+                self.disc_factor, global_step, threshold=self.discriminator_iter_start
+            )
+            d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
+
+            log = {
+                "{}/disc_loss".format(split): d_loss.clone().detach().mean(),
+                "{}/logits_real".format(split): logits_real.detach().mean(),
+                "{}/logits_fake".format(split): logits_fake.detach().mean(),
+            }
+            return d_loss, log
+
+
+class LPIPSWithDiscriminator3D(nn.Module):
+    def __init__(
+        self,
+        disc_start,
+        logvar_init=0.0,
+        kl_weight=1.0,
+        pixelloss_weight=1.0,
+        perceptual_weight=1.0,
+        # --- Discriminator Loss ---
+        disc_num_layers=3,
+        disc_in_channels=3,
+        disc_factor=1.0,
+        disc_weight=1.0,
+        use_actnorm=False,
+        disc_conditional=False,
+        disc_loss="hinge",
+        learn_logvar: bool = False,
+        loss_type: str = "l1"
+    ):
+
+        super().__init__()
+        assert disc_loss in ["hinge", "vanilla"]
+        self.kl_weight = kl_weight
+        self.pixel_weight = pixelloss_weight
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+        self.logvar = nn.Parameter(
+            torch.full((), logvar_init), requires_grad=learn_logvar
+        )
+        self.discriminator = NLayerDiscriminator3D(
+            input_nc=disc_in_channels, n_layers=disc_num_layers, use_actnorm=use_actnorm
+        ).apply(weights_init)
+        self.discriminator_iter_start = disc_start
+        self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
+        self.disc_factor = disc_factor
+        self.discriminator_weight = disc_weight
+        self.disc_conditional = disc_conditional
+        self.loss_func = l1 if loss_type == "l1" else l2
+        
+    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+        layer = last_layer if last_layer is not None else self.last_layer[0]
+
+        nll_grads = torch.autograd.grad(nll_loss, layer, retain_graph=True)[0]
+        g_grads = torch.autograd.grad(g_loss, layer, retain_graph=True)[0]
+        
+        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+        d_weight = d_weight * self.discriminator_weight
+        return d_weight
+
+    def forward(
+        self,
+        inputs,
+        reconstructions,
+        posteriors,
+        optimizer_idx,
+        global_step,
+        split="train",
+        weights=None,
+        last_layer=None,
+        cond=None
+    ):
+        
+        t = inputs.shape[2]
+        # GAN Part
+        if optimizer_idx == 0:
+            inputs = rearrange(inputs, "b c t h w -> (b t) c h w").contiguous()
+            reconstructions = rearrange(
+                reconstructions, "b c t h w -> (b t) c h w"
+            ).contiguous()
+            rec_loss = self.loss_func(inputs, reconstructions)
+            if self.perceptual_weight > 0:
+                p_loss = self.perceptual_loss(inputs, reconstructions)
+                rec_loss = rec_loss + self.perceptual_weight * p_loss
+            nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+            weighted_nll_loss = nll_loss
+            if weights is not None:
+                weighted_nll_loss = weights * nll_loss
+            weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+            nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+            kl_loss = posteriors.kl()
+            kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+
+            inputs = rearrange(inputs, "(b t) c h w -> b c t h w", t=t).contiguous()
+            reconstructions = rearrange(
+                reconstructions, "(b t) c h w -> b c t h w", t=t
+            ).contiguous()
+            
+            logits_fake = self.discriminator(reconstructions)
+            g_loss = -torch.mean(logits_fake)
+            if global_step >= self.discriminator_iter_start:
+                if self.disc_factor > 0.0:
+                    d_weight = self.calculate_adaptive_weight(
+                        nll_loss, g_loss, last_layer=last_layer
+                    )
+                else:
+                    d_weight = torch.tensor(1.0)
+            else:
+                d_weight = torch.tensor(0.0)
+                g_loss = torch.tensor(0.0, requires_grad=True)
+                
+            disc_factor = adopt_weight(
+                self.disc_factor, global_step, threshold=self.discriminator_iter_start
+            )
+            loss = (
+                weighted_nll_loss
+                + self.kl_weight * kl_loss
+                + d_weight * disc_factor * g_loss
+            )
+            log = {
+                "{}/total_loss".format(split): loss.clone().detach().mean(),
+                "{}/logvar".format(split): self.logvar.detach(),
+                "{}/kl_loss".format(split): kl_loss.detach().mean(),
+                "{}/nll_loss".format(split): nll_loss.detach().mean(),
+                "{}/rec_loss".format(split): rec_loss.detach().mean(),
+                "{}/d_weight".format(split): d_weight.detach(),
+                "{}/disc_factor".format(split): torch.tensor(disc_factor),
+                "{}/g_loss".format(split): g_loss.detach().mean(),
+            }
+            return loss, log
+        elif optimizer_idx == 1:
+            logits_real = self.discriminator(inputs.contiguous().detach())
+            logits_fake = self.discriminator(reconstructions.contiguous().detach())
+
+            disc_factor = adopt_weight(
+                self.disc_factor, global_step, threshold=self.discriminator_iter_start
+            )
+            
+            d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
+
+            log = {
+                "{}/disc_loss".format(split): d_loss.clone().detach().mean(),
+                "{}/logits_real".format(split): logits_real.detach().mean(),
+                "{}/logits_fake".format(split): logits_fake.detach().mean(),
+            }
+            return d_loss, log
+
+
+class SimpleLPIPS(nn.Module):
+    def __init__(
+        self,
+        logvar_init=0.0,
+        kl_weight=1.0,
+        pixelloss_weight=1.0,
+        perceptual_weight=1.0,
+        disc_loss="hinge",
+        learn_logvar: bool = False,
+        **kwargs
+    ):
+
+        super().__init__()
+        assert disc_loss in ["hinge", "vanilla"]
+        self.kl_weight = kl_weight
+        self.pixel_weight = pixelloss_weight
+        self.perceptual_loss = LPIPS().eval()
+        self.perceptual_weight = perceptual_weight
+        self.logvar = nn.Parameter(
+            torch.full((), logvar_init), requires_grad=learn_logvar
+        )
+
+    def forward(
+        self,
+        inputs,
+        reconstructions,
+        posteriors,
+        split="train",
+        weights=None,
+    ):
+        inputs = rearrange(inputs, "b c t h w -> (b t) c h w").contiguous()
+        reconstructions = rearrange(
+            reconstructions, "b c t h w -> (b t) c h w"
+        ).contiguous()
+        rec_loss = torch.abs(inputs - reconstructions)
+        if self.perceptual_weight > 0:
+            p_loss = self.perceptual_loss(inputs, reconstructions)
+            rec_loss = rec_loss + self.perceptual_weight * p_loss
+        nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+        weighted_nll_loss = nll_loss
+        if weights is not None:
+            weighted_nll_loss = weights * nll_loss
+        weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+        nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+        kl_loss = posteriors.kl()
+        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+        loss = weighted_nll_loss + self.kl_weight * kl_loss
+        log = {
+            "{}/total_loss".format(split): loss.clone().detach().mean(),
+            "{}/logvar".format(split): self.logvar.detach(),
+            "{}/kl_loss".format(split): kl_loss.detach().mean(),
+            "{}/nll_loss".format(split): nll_loss.detach().mean(),
+            "{}/rec_loss".format(split): rec_loss.detach().mean(),
+        }
+        if self.perceptual_weight > 0:
+            log.update({"{}/p_loss".format(split): p_loss.detach().mean()})
+        return loss, log
diff --git a/opensora/models/causalvideovae/model/modeling_videobase.py b/opensora/models/causalvideovae/model/modeling_videobase.py
new file mode 100755
index 000000000..260575675
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modeling_videobase.py
@@ -0,0 +1,53 @@
+import torch
+from diffusers import ModelMixin, ConfigMixin
+from torch import nn
+import os
+import json
+import pytorch_lightning as pl
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+from typing import Optional, Union
+import glob
+
+
+class VideoBaseAE(ModelMixin, ConfigMixin):
+    config_name = "config.json"
+    
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+    
+    def encode(self, x: torch.Tensor, *args, **kwargs):
+        pass
+
+    def decode(self, encoding: torch.Tensor, *args, **kwargs):
+        pass
+    
+    @property
+    def num_training_steps(self) -> int:
+        """Total training steps inferred from datamodule and devices."""
+        if self.trainer.max_steps:
+            return self.trainer.max_steps
+    
+        limit_batches = self.trainer.limit_train_batches
+        batches = len(self.train_dataloader())
+        batches = min(batches, limit_batches) if isinstance(limit_batches, int) else int(limit_batches * batches)     
+    
+        num_devices = max(1, self.trainer.num_gpus, self.trainer.num_processes)
+        if self.trainer.tpu_cores:
+            num_devices = max(num_devices, self.trainer.tpu_cores)
+    
+        effective_accum = self.trainer.accumulate_grad_batches * num_devices
+        return (batches // effective_accum) * self.trainer.max_epochs
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
+        ckpt_files = glob.glob(os.path.join(pretrained_model_name_or_path, '*.ckpt'))
+        if ckpt_files:
+            # Adapt to checkpoint
+            last_ckpt_file = ckpt_files[-1]
+            config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
+            model = cls.from_config(config_file)
+            model.init_from_ckpt(last_ckpt_file)
+            return model
+        else:
+            return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/__init__.py b/opensora/models/causalvideovae/model/modules/__init__.py
new file mode 100755
index 000000000..7ef41edd2
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/__init__.py
@@ -0,0 +1,6 @@
+from .block import Block
+from .attention import *
+from .conv import *
+from .normalize import *
+from .resnet_block import *
+from .updownsample import *
diff --git a/opensora/models/causalvideovae/model/modules/attention.py b/opensora/models/causalvideovae/model/modules/attention.py
new file mode 100755
index 000000000..73c4c3793
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/attention.py
@@ -0,0 +1,239 @@
+import torch.nn as nn
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+from .normalize import Normalize
+from .conv import CausalConv3d
+import torch
+import numpy as np
+from einops import rearrange
+from .block import Block
+from .ops import video_to_image
+
+class LinearAttention(Block):
+    def __init__(self, dim, heads=4, dim_head=32):
+        super().__init__()
+        self.heads = heads
+        hidden_dim = dim_head * heads
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        qkv = self.to_qkv(x)
+        q, k, v = rearrange(
+            qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
+        )
+        k = k.softmax(dim=-1)
+        context = torch.einsum("bhdn,bhen->bhde", k, v)
+        out = torch.einsum("bhde,bhdn->bhen", context, q)
+        out = rearrange(
+            out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
+        )
+        return self.to_out(out)
+
+
+class LinAttnBlock(LinearAttention):
+    """to match AttnBlock usage"""
+
+    def __init__(self, in_channels):
+        super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
+
+
+class AttnBlock3D(Block):
+    """Compatible with old versions, there are issues, use with caution."""
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+        self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+        self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+        self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, t, h, w = q.shape
+        q = q.reshape(b * t, c, h * w)
+        q = q.permute(0, 2, 1)  # b,hw,c
+        k = k.reshape(b * t, c, h * w)  # b,c,hw
+        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b * t, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, t, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+class AttnBlock3DFix(nn.Module):
+    """
+    Thanks to https://github.com/PKU-YuanGroup/Open-Sora-Plan/pull/172.
+    """
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+        self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+        self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+        self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
+
+    def forward(self, x):
+        h_ = x
+        if npu_config is None:
+            h_ = self.norm(h_)
+        else:
+            h_ = npu_config.run_group_norm(self.norm, h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        # q: (b c t h w) -> (b t c h w) -> (b*t c h*w) -> (b*t h*w c)
+        b, c, t, h, w = q.shape
+        q = q.permute(0, 2, 1, 3, 4)
+        q = q.reshape(b * t, c, h * w)
+        q = q.permute(0, 2, 1)
+
+        # k: (b c t h w) -> (b t c h w) -> (b*t c h*w)
+        k = k.permute(0, 2, 1, 3, 4)
+        k = k.reshape(b * t, c, h * w)
+
+        # w: (b*t hw hw)
+        w_ = torch.bmm(q, k)
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        # v: (b c t h w) -> (b t c h w) -> (bt c hw)
+        # w_: (bt hw hw) -> (bt hw hw)
+        v = v.permute(0, 2, 1, 3, 4)
+        v = v.reshape(b * t, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+
+        # h_: (b*t c hw) -> (b t c h w) -> (b c t h w)
+        h_ = h_.reshape(b, t, c, h, w)
+        h_ = h_.permute(0, 2, 1, 3 ,4)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class AttnBlock(Block):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = torch.nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        
+    @video_to_image
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)  # b,hw,c
+        k = k.reshape(b, c, h * w)  # b,c,hw
+        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+class TemporalAttnBlock(Block):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = torch.nn.Conv3d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.v = torch.nn.Conv3d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.proj_out = torch.nn.Conv3d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, t, h, w = q.shape
+        q = rearrange(q, "b c t h w -> (b h w) t c")
+        k = rearrange(k, "b c t h w -> (b h w) c t")
+        v = rearrange(v, "b c t h w -> (b h w) c t")
+        w_ = torch.bmm(q, k)
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        w_ = w_.permute(0, 2, 1)
+        h_ = torch.bmm(v, w_)
+        h_ = rearrange(h_, "(b h w) c t -> b c t h w", h=h, w=w)
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+
+def make_attn(in_channels, attn_type="vanilla"):
+    assert attn_type in ["vanilla", "linear", "none", "vanilla3D"], f"attn_type {attn_type} unknown"
+    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+    print(attn_type)
+    if attn_type == "vanilla":
+        return AttnBlock(in_channels)
+    elif attn_type == "vanilla3D":
+        return AttnBlock3D(in_channels)
+    elif attn_type == "none":
+        return nn.Identity(in_channels)
+    else:
+        return LinAttnBlock(in_channels)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/block.py b/opensora/models/causalvideovae/model/modules/block.py
new file mode 100755
index 000000000..e93672d20
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/block.py
@@ -0,0 +1,5 @@
+import torch.nn as nn
+
+class Block(nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/conv.py b/opensora/models/causalvideovae/model/modules/conv.py
new file mode 100755
index 000000000..71d6a4705
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/conv.py
@@ -0,0 +1,127 @@
+import torch.nn as nn
+from typing import Union, Tuple
+import torch.nn.functional as F
+import torch
+from .block import Block
+from .ops import cast_tuple
+from einops import rearrange
+from .ops import video_to_image
+from torch.utils.checkpoint import checkpoint
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+
+class Conv2d(nn.Conv2d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int]] = 3,
+        stride: Union[int, Tuple[int]] = 1,
+        padding: Union[str, int, Tuple[int]] = 0,
+        dilation: Union[int, Tuple[int]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+            padding_mode,
+            device,
+            dtype,
+        )
+        
+    @video_to_image
+    def forward(self, x):
+        return super().forward(x)
+        
+
+
+class CausalConv3d(nn.Module):
+    def __init__(
+        self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], init_method="random", **kwargs
+    ):
+        super().__init__()
+        self.kernel_size = cast_tuple(kernel_size, 3)
+        self.time_kernel_size = self.kernel_size[0]
+        self.chan_in = chan_in
+        self.chan_out = chan_out
+        stride = kwargs.pop("stride", 1)
+        padding = kwargs.pop("padding", 0)
+        padding = list(cast_tuple(padding, 3))
+        padding[0] = 0
+        stride = cast_tuple(stride, 3)
+        self.conv = nn.Conv3d(chan_in, chan_out, self.kernel_size, stride=stride, padding=padding)
+        self.pad = nn.ReplicationPad2d((0, 0, self.time_kernel_size - 1, 0))
+        self._init_weights(init_method)
+        
+    def _init_weights(self, init_method):
+        ks = torch.tensor(self.kernel_size)
+        if init_method == "avg":
+            assert (
+                self.kernel_size[1] == 1 and self.kernel_size[2] == 1
+            ), "only support temporal up/down sample"
+            assert self.chan_in == self.chan_out, "chan_in must be equal to chan_out"
+            weight = torch.zeros((self.chan_out, self.chan_in, *self.kernel_size))
+
+            eyes = torch.concat(
+                [
+                    torch.eye(self.chan_in).unsqueeze(-1) * 1/3,
+                    torch.eye(self.chan_in).unsqueeze(-1) * 1/3,
+                    torch.eye(self.chan_in).unsqueeze(-1) * 1/3,
+                ],
+                dim=-1,
+            )
+            weight[:, :, :, 0, 0] = eyes
+
+            self.conv.weight = nn.Parameter(
+                weight,
+                requires_grad=True,
+            )
+        elif init_method == "zero":
+            self.conv.weight = nn.Parameter(
+                torch.zeros((self.chan_out, self.chan_in, *self.kernel_size)),
+                requires_grad=True,
+            )
+        if self.conv.bias is not None:
+            nn.init.constant_(self.conv.bias, 0)
+            
+    def forward(self, x):
+        if npu_config is not None and npu_config.on_npu:
+            x_dtype = x.dtype
+            first_frame_pad = x[:, :, :1, :, :].repeat(
+                (1, 1, self.time_kernel_size - 1, 1, 1)
+            )  # b c t h w
+            x = torch.concatenate((first_frame_pad, x), dim=2)  # 3 + 16
+            return npu_config.run_conv3d(self.conv, x, x_dtype)
+        else:
+            # 1 + 16   16 as video, 1 as image
+            first_frame_pad = x[:, :, :1, :, :].repeat(
+                (1, 1, self.time_kernel_size - 1, 1, 1)
+            )  # b c t h w
+            x = torch.concatenate((first_frame_pad, x), dim=2)  # 3 + 16
+            return self.conv(x)
+    
+    
+class CausalConv3d_GC(CausalConv3d):
+    def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int]], init_method="random", **kwargs):
+        super().__init__(chan_in, chan_out, kernel_size, init_method, **kwargs)
+    def forward(self, x):
+        # 1 + 16   16 as video, 1 as image
+        first_frame_pad = x[:, :, :1, :, :].repeat(
+            (1, 1, self.time_kernel_size - 1, 1, 1)
+        )   # b c t h w
+        x = torch.concatenate((first_frame_pad, x), dim=2)  # 3 + 16
+        return checkpoint(self.conv, x)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/normalize.py b/opensora/models/causalvideovae/model/modules/normalize.py
new file mode 100755
index 000000000..7c8c05f0f
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/normalize.py
@@ -0,0 +1,101 @@
+import torch
+import torch.nn as nn
+from .block import Block
+
+class GroupNorm(Block):
+    def __init__(self, num_channels, num_groups=32, eps=1e-6, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.norm = torch.nn.GroupNorm(
+            num_groups=num_groups, num_channels=num_channels, eps=1e-6, affine=True
+        )
+    def forward(self, x):
+        return self.norm(x)
+
+def Normalize(in_channels, num_groups=32):
+    return torch.nn.GroupNorm(
+        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+    )
+
+class ActNorm(nn.Module):
+    def __init__(self, num_features, logdet=False, affine=True,
+                 allow_reverse_init=False):
+        assert affine
+        super().__init__()
+        self.logdet = logdet
+        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+        self.allow_reverse_init = allow_reverse_init
+
+        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
+
+    def initialize(self, input):
+        with torch.no_grad():
+            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+            mean = (
+                flatten.mean(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+            std = (
+                flatten.std(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+
+            self.loc.data.copy_(-mean)
+            self.scale.data.copy_(1 / (std + 1e-6))
+
+    def forward(self, input, reverse=False):
+        if reverse:
+            return self.reverse(input)
+        if len(input.shape) == 2:
+            input = input[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        _, _, height, width = input.shape
+
+        if self.training and self.initialized.item() == 0:
+            self.initialize(input)
+            self.initialized.fill_(1)
+
+        h = self.scale * (input + self.loc)
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+
+        if self.logdet:
+            log_abs = torch.log(torch.abs(self.scale))
+            logdet = height*width*torch.sum(log_abs)
+            logdet = logdet * torch.ones(input.shape[0]).to(input)
+            return h, logdet
+
+        return h
+
+    def reverse(self, output):
+        if self.training and self.initialized.item() == 0:
+            if not self.allow_reverse_init:
+                raise RuntimeError(
+                    "Initializing ActNorm in reverse direction is "
+                    "disabled by default. Use allow_reverse_init=True to enable."
+                )
+            else:
+                self.initialize(output)
+                self.initialized.fill_(1)
+
+        if len(output.shape) == 2:
+            output = output[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        h = output / self.scale - self.loc
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+        return h
diff --git a/opensora/models/causalvideovae/model/modules/ops.py b/opensora/models/causalvideovae/model/modules/ops.py
new file mode 100755
index 000000000..e5987e740
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/ops.py
@@ -0,0 +1,40 @@
+import torch
+from einops import rearrange
+
+def video_to_image(func):
+    def wrapper(self, x, *args, **kwargs):
+        if x.dim() == 5:
+            t = x.shape[2]
+            x = rearrange(x, "b c t h w -> (b t) c h w")
+            x = func(self, x, *args, **kwargs)
+            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+        return x
+    return wrapper
+
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+def cast_tuple(t, length=1):
+    return t if isinstance(t, tuple) or isinstance(t, list) else ((t,) * length)
+
+def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True):
+    n_dims = len(x.shape)
+    if src_dim < 0:
+        src_dim = n_dims + src_dim
+    if dest_dim < 0:
+        dest_dim = n_dims + dest_dim
+    assert 0 <= src_dim < n_dims and 0 <= dest_dim < n_dims
+    dims = list(range(n_dims))
+    del dims[src_dim]
+    permutation = []
+    ctr = 0
+    for i in range(n_dims):
+        if i == dest_dim:
+            permutation.append(src_dim)
+        else:
+            permutation.append(dims[ctr])
+            ctr += 1
+    x = x.permute(permutation)
+    if make_contiguous:
+        x = x.contiguous()
+    return x
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/quant.py b/opensora/models/causalvideovae/model/modules/quant.py
new file mode 100755
index 000000000..bb702cee5
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/quant.py
@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+import numpy as np
+import torch.nn.functional as F
+from .ops import shift_dim
+
+class Codebook(nn.Module):
+    def __init__(self, n_codes, embedding_dim):
+        super().__init__()
+        self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim))
+        self.register_buffer("N", torch.zeros(n_codes))
+        self.register_buffer("z_avg", self.embeddings.data.clone())
+
+        self.n_codes = n_codes
+        self.embedding_dim = embedding_dim
+        self._need_init = True
+
+    def _tile(self, x):
+        d, ew = x.shape
+        if d < self.n_codes:
+            n_repeats = (self.n_codes + d - 1) // d
+            std = 0.01 / np.sqrt(ew)
+            x = x.repeat(n_repeats, 1)
+            x = x + torch.randn_like(x) * std
+        return x
+
+    def _init_embeddings(self, z):
+        # z: [b, c, t, h, w]
+        self._need_init = False
+        flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)
+        y = self._tile(flat_inputs)
+
+        d = y.shape[0]
+        _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
+        if dist.is_initialized():
+            dist.broadcast(_k_rand, 0)
+        self.embeddings.data.copy_(_k_rand)
+        self.z_avg.data.copy_(_k_rand)
+        self.N.data.copy_(torch.ones(self.n_codes))
+
+    def forward(self, z):
+        # z: [b, c, t, h, w]
+        if self._need_init and self.training:
+            self._init_embeddings(z)
+        flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)
+        distances = (
+            (flat_inputs**2).sum(dim=1, keepdim=True)
+            - 2 * flat_inputs @ self.embeddings.t()
+            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
+        )
+
+        encoding_indices = torch.argmin(distances, dim=1)
+        encode_onehot = F.one_hot(encoding_indices, self.n_codes).type_as(flat_inputs)
+        encoding_indices = encoding_indices.view(z.shape[0], *z.shape[2:])
+
+        embeddings = F.embedding(encoding_indices, self.embeddings)
+        embeddings = shift_dim(embeddings, -1, 1)
+
+        commitment_loss = 0.25 * F.mse_loss(z, embeddings.detach())
+
+        # EMA codebook update
+        if self.training:
+            n_total = encode_onehot.sum(dim=0)
+            encode_sum = flat_inputs.t() @ encode_onehot
+            if dist.is_initialized():
+                dist.all_reduce(n_total)
+                dist.all_reduce(encode_sum)
+
+            self.N.data.mul_(0.99).add_(n_total, alpha=0.01)
+            self.z_avg.data.mul_(0.99).add_(encode_sum.t(), alpha=0.01)
+
+            n = self.N.sum()
+            weights = (self.N + 1e-7) / (n + self.n_codes * 1e-7) * n
+            encode_normalized = self.z_avg / weights.unsqueeze(1)
+            self.embeddings.data.copy_(encode_normalized)
+
+            y = self._tile(flat_inputs)
+            _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
+            if dist.is_initialized():
+                dist.broadcast(_k_rand, 0)
+
+            usage = (self.N.view(self.n_codes, 1) >= 1).float()
+            self.embeddings.data.mul_(usage).add_(_k_rand * (1 - usage))
+
+        embeddings_st = (embeddings - z).detach() + z
+
+        avg_probs = torch.mean(encode_onehot, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        return dict(
+            embeddings=embeddings_st,
+            encodings=encoding_indices,
+            commitment_loss=commitment_loss,
+            perplexity=perplexity,
+        )
+
+    def dictionary_lookup(self, encodings):
+        embeddings = F.embedding(encodings, self.embeddings)
+        return embeddings
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/resnet_block.py b/opensora/models/causalvideovae/model/modules/resnet_block.py
new file mode 100755
index 000000000..ce70334f1
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/resnet_block.py
@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+from einops import rearrange, pack, unpack
+from .normalize import Normalize
+from .ops import nonlinearity, video_to_image
+from .conv import CausalConv3d, CausalConv3d_GC
+from .block import Block
+from torch.utils.checkpoint import checkpoint
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+    
+class ResnetBlock2D(Block):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+                 dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+                
+    @video_to_image
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        x = x + h
+        return x
+
+class ResnetBlock3D(Block):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
+            else:
+                self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
+
+    def forward(self, x):
+        h = x
+        if npu_config is None:
+            h = self.norm1(h)
+        else:
+            h = npu_config.run_group_norm(self.norm1, h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        if npu_config is None:
+            h = self.norm2(h)
+        else:
+            h = npu_config.run_group_norm(self.norm2, h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x + h
+    
+    
+class ResnetBlock3D_GC(Block):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
+            else:
+                self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
+                
+    def forward(self, x):
+        return checkpoint(self._forward, x, use_reentrant=True)
+    
+    def _forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x + h
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/modules/updownsample.py b/opensora/models/causalvideovae/model/modules/updownsample.py
new file mode 100755
index 000000000..86a48cc38
--- /dev/null
+++ b/opensora/models/causalvideovae/model/modules/updownsample.py
@@ -0,0 +1,385 @@
+from typing import Union, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .resnet_block import ResnetBlock3D
+from .attention import TemporalAttnBlock
+from .normalize import Normalize
+from .ops import cast_tuple, video_to_image
+from .conv import CausalConv3d, CausalConv3d_GC
+from einops import rearrange
+from .block import Block
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+
+
+class Upsample(Block):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.with_conv = True
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+            
+    @video_to_image
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+class Downsample(Block):
+    def __init__(self, in_channels, out_channels, undown=False):
+        super().__init__()
+        self.with_conv = True
+        self.undown = undown
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            if self.undown:
+                self.conv = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+            else:
+                self.conv = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+    @video_to_image
+    def forward(self, x):
+        if self.with_conv:
+            if self.undown:
+                if npu_config is not None and npu_config.on_npu:
+                    x_dtype = x.dtype
+                    x = x.to(npu_config.replaced_type)
+                    x = npu_config.run_conv3d(self.conv, x, x_dtype)
+                else:
+                    x = self.conv(x)
+            else:
+                pad = (0, 1, 0, 1)
+                if npu_config is not None and npu_config.on_npu:
+                    x_dtype = x.dtype
+                    x = x.to(npu_config.replaced_type)
+                    x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+                    x = npu_config.run_conv3d(self.conv, x, x_dtype)
+                else:
+                    x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+                    x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+class SpatialDownsample2x(Block):
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int]] = (3, 3),
+        stride: Union[int, Tuple[int]] = (2, 2),
+        **kwargs
+    ):
+        super().__init__()
+        kernel_size = cast_tuple(kernel_size, 2)
+        stride = cast_tuple(stride, 2)
+        self.chan_in = chan_in
+        self.chan_out = chan_out
+        self.kernel_size = kernel_size
+        self.conv = CausalConv3d(
+            self.chan_in,
+            self.chan_out,
+            (1,) + self.kernel_size,
+            stride=(1, ) + stride,
+            padding=0
+        )
+
+    def forward(self, x):
+        pad = (0,1,0,1,0,0)
+        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+class SpatialUpsample2x_GC(Block):
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int]] = (3, 3),
+        stride: Union[int, Tuple[int]] = (1, 1),
+        unup=False,
+    ):
+        super().__init__()
+        self.chan_in = chan_in
+        self.chan_out = chan_out
+        self.kernel_size = kernel_size
+        self.unup = unup
+        self.conv = CausalConv3d_GC(
+            self.chan_in,
+            self.chan_out,
+            (1,) + self.kernel_size,
+            stride=(1, ) + stride,
+            padding=1
+        )
+
+    def forward(self, x):
+        if not self.unup:
+            t = x.shape[2]
+            x = rearrange(x, "b c t h w -> b (c t) h w")
+            x = F.interpolate(x, scale_factor=(2,2), mode="nearest")
+            x = rearrange(x, "b (c t) h w -> b c t h w", t=t)
+        x = self.conv(x)
+        return x
+    
+
+class SpatialUpsample2x(Block):
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: Union[int, Tuple[int]] = (3, 3),
+        stride: Union[int, Tuple[int]] = (1, 1),
+        unup=False,
+    ):
+        super().__init__()
+        self.chan_in = chan_in
+        self.chan_out = chan_out
+        self.kernel_size = kernel_size
+        self.unup = unup
+        self.conv = CausalConv3d(
+            self.chan_in,
+            self.chan_out,
+            (1,) + self.kernel_size,
+            stride=(1, ) + stride,
+            padding=1
+        )
+
+    def forward(self, x):
+        if not self.unup:
+            t = x.shape[2]
+            x = rearrange(x, "b c t h w -> b (c t) h w")
+            x = F.interpolate(x, scale_factor=(2,2), mode="nearest")
+            x = rearrange(x, "b (c t) h w -> b c t h w", t=t)
+        x = self.conv(x)
+        return x
+    
+class TimeDownsample2x(Block):
+    def __init__(
+        self,
+        chan_in,
+        chan_out,
+        kernel_size: int = 3
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        if npu_config is not None and npu_config.on_npu:
+            self.avg_pool = nn.AvgPool2d((kernel_size, 1), stride=(2, 1))
+            self.pad = nn.ReplicationPad3d((0, 0, 0, 0, self.kernel_size - 1, 0))
+        else:
+            self.conv = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1))
+        
+    def forward(self, x):
+        if npu_config is not None and npu_config.on_npu:
+            n, c, d, h, w = x.shape
+            x = self.pad(x)
+            x = x.view(n * c, -1, h * w)
+            pooled = self.avg_pool(x)
+            output = pooled.view(n, c, -1, h, w)
+            return output
+        else:
+            first_frame_pad = x[:, :, :1, :, :].repeat(
+                (1, 1, self.kernel_size - 1, 1, 1)
+            )
+            x = torch.concatenate((first_frame_pad, x), dim=2)
+            return self.conv(x)
+
+class TimeUpsample2x(Block):
+    def __init__(
+        self,
+        chan_in,
+        chan_out
+    ):
+        super().__init__()
+    def forward(self, x):
+        if x.size(2) > 1:
+            x,x_= x[:,:,:1],x[:,:,1:]
+            x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
+            x = torch.concat([x, x_], dim=2)
+        return x
+    
+class TimeDownsampleRes2x(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: int = 3,
+        mix_factor: float = 2.0,
+    ):
+        super().__init__()
+        self.kernel_size = cast_tuple(kernel_size, 3)
+        if npu_config is not None and npu_config.on_npu:
+            self.avg_pool = nn.AvgPool2d((kernel_size, 1), stride=(2, 1))
+            self.pad = nn.ReplicationPad3d((0, 0, 0, 0, kernel_size - 1, 0))
+        else:
+            self.avg_pool = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1))
+        self.conv = nn.Conv3d(
+            in_channels, out_channels, self.kernel_size, stride=(2,1,1), padding=(0,1,1)
+        )
+        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
+    
+    def forward(self, x):
+        alpha = torch.sigmoid(self.mix_factor)
+        if npu_config is not None and npu_config.on_npu:
+            n, c, d, h, w = x.shape
+            x_dtype = x.dtype
+            x = x.to(npu_config.replaced_type)
+            x = self.pad(x)
+            pad_x = x.view(n, c, -1, h, w)
+            avg_x = self.avg_pool(x.view(n * c, -1, h * w)).view(n, c, -1, h, w).to(x_dtype)
+            conv_x = npu_config.run_conv3d(self.conv, pad_x, x_dtype)
+            return alpha * avg_x + (1 - alpha) * conv_x
+        else:
+            first_frame_pad = x[:, :, :1, :, :].repeat(
+                (1, 1, self.kernel_size[0] - 1, 1, 1)
+            )
+            x = torch.concatenate((first_frame_pad, x), dim=2)
+            return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(x)
+
+class TimeUpsampleRes2x(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: int = 3,
+        mix_factor: float = 2.0,
+    ):
+        super().__init__()
+        self.conv = CausalConv3d(
+            in_channels, out_channels, kernel_size, padding=1
+        )
+        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
+        
+    def forward(self, x):
+        alpha = torch.sigmoid(self.mix_factor)
+        if x.size(2) > 1:
+            x,x_= x[:,:,:1],x[:,:,1:]
+            if npu_config is not None and npu_config.on_npu:
+                x_dtype = x_.dtype
+                x_ = x_.to(npu_config.replaced_type)
+                x_ = F.interpolate(x_, scale_factor=(2, 1, 1), mode='trilinear')
+                x_ = x_.to(x_dtype)
+            else:
+                x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
+            x = torch.concat([x, x_], dim=2)
+        return alpha * x + (1-alpha) * self.conv(x)
+
+class TimeDownsampleResAdv2x(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: int = 3,
+        mix_factor: float = 1.5,
+    ):
+        super().__init__()
+        self.kernel_size = cast_tuple(kernel_size, 3)
+        self.avg_pool = nn.AvgPool3d((kernel_size,1,1), stride=(2,1,1))
+        self.attn = TemporalAttnBlock(in_channels)
+        self.res = ResnetBlock3D(in_channels=in_channels, out_channels=in_channels, dropout=0.0)
+        self.conv = nn.Conv3d(
+            in_channels, out_channels, self.kernel_size, stride=(2,1,1), padding=(0,1,1)
+        )
+        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
+    
+    def forward(self, x):
+        first_frame_pad = x[:, :, :1, :, :].repeat(
+            (1, 1, self.kernel_size[0] - 1, 1, 1)
+        )
+        x = torch.concatenate((first_frame_pad, x), dim=2)
+        alpha = torch.sigmoid(self.mix_factor)
+        return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(self.attn((self.res(x))))
+
+class TimeUpsampleResAdv2x(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size: int = 3,
+        mix_factor: float = 1.5,
+    ):
+        super().__init__()
+        self.res = ResnetBlock3D(in_channels=in_channels, out_channels=in_channels, dropout=0.0)
+        self.attn = TemporalAttnBlock(in_channels)
+        self.norm = Normalize(in_channels=in_channels)
+        self.conv = CausalConv3d(
+            in_channels, out_channels, kernel_size, padding=1
+        )
+        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
+        
+    def forward(self, x):
+        if x.size(2) > 1:
+            x,x_= x[:,:,:1],x[:,:,1:]
+            if npu_config is not None and npu_config.on_npu:
+                x_dtype = x_.dtype
+                x_ = x_.to(npu_config.replaced_type)
+                x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
+                x_ = x_.to(x_dtype)
+            else:
+                x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
+            x = torch.concat([x, x_], dim=2)
+        alpha = torch.sigmoid(self.mix_factor)
+        return alpha * x + (1 - alpha) * self.conv(self.attn(self.res(x)))
+
+
+class Spatial2xTime2x3DDownsample(Block):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=2)
+
+    def forward(self, x):
+        pad = (0,1,0,1,0,0)
+        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+class Spatial2x3DDownsample(Block):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=(1,2,2))
+
+    def forward(self, x):
+        pad = (0,1,0,1,0,0)
+        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+    
+
+class Spatial2x3DUpsample(Block):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=(1,2,2), mode='trilinear')
+        return self.conv(x)
+    
+class Spatial2xTime2x3DUpsample(Block):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
+    def forward(self, x):
+        if x.size(2) > 1:
+            x,x_= x[:,:,:1],x[:,:,1:]
+            x_= F.interpolate(x_, scale_factor=(2,2,2), mode='trilinear')
+            x = F.interpolate(x, scale_factor=(1,2,2), mode='trilinear')
+            x = torch.concat([x, x_], dim=2)
+        else:
+             x = F.interpolate(x, scale_factor=(1,2,2), mode='trilinear')
+        return self.conv(x)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/trainer_videobase.py b/opensora/models/causalvideovae/model/trainer_videobase.py
new file mode 100755
index 000000000..b6c4b0c7d
--- /dev/null
+++ b/opensora/models/causalvideovae/model/trainer_videobase.py
@@ -0,0 +1,26 @@
+from transformers import Trainer
+import torch.nn.functional as F
+from typing import Optional
+import os
+import torch
+from transformers.utils import WEIGHTS_NAME
+import json
+
+class VideoBaseTrainer(Trainer):
+
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        output_dir = output_dir if output_dir is not None else self.args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+        if state_dict is None:
+            state_dict = self.model.state_dict()
+        
+        # get model config
+        model_config = self.model.config.to_dict()
+        
+        # add more information
+        model_config['model'] = self.model.__class__.__name__
+        
+        with open(os.path.join(output_dir, "config.json"), "w") as file:
+            json.dump(self.model.config.to_dict(), file)
+        torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
+        torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
diff --git a/opensora/models/causalvideovae/model/utils/__init__.py b/opensora/models/causalvideovae/model/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/models/causalvideovae/model/utils/distrib_utils.py b/opensora/models/causalvideovae/model/utils/distrib_utils.py
new file mode 100755
index 000000000..760c0673f
--- /dev/null
+++ b/opensora/models/causalvideovae/model/utils/distrib_utils.py
@@ -0,0 +1,42 @@
+import torch
+import numpy as np
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(torch.pow(self.mean, 2)
+                                       + self.var - 1.0 - self.logvar,
+                                       dim=[1, 2, 3])
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=[1, 2, 3])
+
+    def nll(self, sample, dims=[1,2,3]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims)
+
+    def mode(self):
+        return self.mean
diff --git a/opensora/models/causalvideovae/model/utils/module_utils.py b/opensora/models/causalvideovae/model/utils/module_utils.py
new file mode 100755
index 000000000..3f843e4f3
--- /dev/null
+++ b/opensora/models/causalvideovae/model/utils/module_utils.py
@@ -0,0 +1,17 @@
+import importlib
+
+Module = str
+MODULES_BASE = "opensora.models.causalvideovae.model.modules."
+
+def resolve_str_to_obj(str_val, append=True):
+    if append:
+        str_val = MODULES_BASE + str_val
+    module_name, class_name = str_val.rsplit('.', 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, class_name)
+
+def create_instance(module_class_str: str, **kwargs):
+    module_name, class_name = module_class_str.rsplit('.', 1)
+    module = importlib.import_module(module_name)
+    class_ = getattr(module, class_name)
+    return class_(**kwargs)
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/utils/scheduler_utils.py b/opensora/models/causalvideovae/model/utils/scheduler_utils.py
new file mode 100755
index 000000000..0650a3f79
--- /dev/null
+++ b/opensora/models/causalvideovae/model/utils/scheduler_utils.py
@@ -0,0 +1,7 @@
+import torch
+
+def cosine_scheduler(step, max_steps, value_base=1, value_end=0):
+    step = torch.tensor(step)
+    cosine_value = 0.5 * (1 + torch.cos(torch.pi * step / max_steps))
+    value = value_end + (value_base - value_end) * cosine_value
+    return value
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/model/utils/video_utils.py b/opensora/models/causalvideovae/model/utils/video_utils.py
new file mode 100755
index 000000000..4fbd3f04c
--- /dev/null
+++ b/opensora/models/causalvideovae/model/utils/video_utils.py
@@ -0,0 +1,10 @@
+import torch
+import numpy as np
+
+def tensor_to_video(x):
+    x = (x * 2 - 1).detach().cpu()
+    x = torch.clamp(x, -1, 1)
+    x = (x + 1) / 2
+    x = x.permute(1, 0, 2, 3).float().numpy() # c t h w -> t c h w
+    x = (255 * x).astype(np.uint8)
+    return x
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/sample/rec_video_vae.py b/opensora/models/causalvideovae/sample/rec_video_vae.py
new file mode 100644
index 000000000..e2c8a549b
--- /dev/null
+++ b/opensora/models/causalvideovae/sample/rec_video_vae.py
@@ -0,0 +1,277 @@
+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
+import glob
+sys.path.append(".")
+from causalvideovae.model 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):
+    video_exts = ['avi', 'mp4', 'webm']
+    
+    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):
+        samples = []
+        samples += sum([glob.glob(os.path.join(folder_path, '**', f'*.{ext}'), recursive=True)
+                            for ext in self.video_exts], [])
+        samples.sort()
+        return samples
+
+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)
+    print(vqvae)
+    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)
+    
\ No newline at end of file
diff --git a/opensora/models/causalvideovae/utils/__init__.py b/opensora/models/causalvideovae/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/models/causalvideovae/utils/dataset_utils.py b/opensora/models/causalvideovae/utils/dataset_utils.py
new file mode 100755
index 000000000..1309989b1
--- /dev/null
+++ b/opensora/models/causalvideovae/utils/dataset_utils.py
@@ -0,0 +1,160 @@
+import math
+from einops import rearrange
+import decord
+from torch.nn import functional as F
+import torch
+
+
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+class DecordInit(object):
+    """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader."""
+
+    def __init__(self, num_threads=1):
+        self.num_threads = num_threads
+        self.ctx = decord.cpu(0)
+
+    def __call__(self, filename):
+        """Perform the Decord initialization.
+        Args:
+            results (dict): The resulting dict to be modified and passed
+                to the next transform in pipeline.
+        """
+        reader = decord.VideoReader(filename,
+                                    ctx=self.ctx,
+                                    num_threads=self.num_threads)
+        return reader
+
+    def __repr__(self):
+        repr_str = (f'{self.__class__.__name__}('
+                    f'sr={self.sr},'
+                    f'num_threads={self.num_threads})')
+        return repr_str
+
+def pad_to_multiple(number, ds_stride):
+    remainder = number % ds_stride
+    if remainder == 0:
+        return number
+    else:
+        padding = ds_stride - remainder
+        return number + padding
+
+class Collate:
+    def __init__(self, args):
+        self.max_image_size = args.max_image_size
+        self.ae_stride = args.ae_stride
+        self.ae_stride_t = args.ae_stride_t
+        self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride)
+        self.ae_stride_1hw = (1, self.ae_stride, self.ae_stride)
+
+        self.patch_size = args.patch_size
+        self.patch_size_t = args.patch_size_t
+        self.patch_size_thw = (self.patch_size_t, self.patch_size, self.patch_size)
+        self.patch_size_1hw = (1, self.patch_size, self.patch_size)
+
+        self.num_frames = args.num_frames
+        self.use_image_num = args.use_image_num
+        self.max_thw = (self.num_frames, self.max_image_size, self.max_image_size)
+        self.max_1hw = (1, self.max_image_size, self.max_image_size)
+
+    def package(self, batch):
+        # import ipdb;ipdb.set_trace()
+        batch_tubes_vid = [i['video_data']['video'] for i in batch]  # b [c t h w]
+        input_ids_vid = torch.stack([i['video_data']['input_ids'] for i in batch])  # b 1 l
+        cond_mask_vid = torch.stack([i['video_data']['cond_mask'] for i in batch])  # b 1 l
+        batch_tubes_img, input_ids_img, cond_mask_img = None, None, None
+        if self.use_image_num != 0: 
+            batch_tubes_img = [j for i in batch for j in i['image_data']['image']]  # b*num_img [c 1 h w]
+            input_ids_img = torch.stack([i['image_data']['input_ids'] for i in batch])  # b image_num l
+            cond_mask_img = torch.stack([i['image_data']['cond_mask'] for i in batch])  # b image_num l
+        return batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img
+
+    def __call__(self, batch):
+        batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img = self.package(batch)
+
+        # import ipdb;ipdb.set_trace()
+        ds_stride = self.ae_stride * self.patch_size
+        t_ds_stride = self.ae_stride_t * self.patch_size_t
+        if self.use_image_num == 0:
+            pad_batch_tubes, attention_mask = self.process(batch_tubes_vid, t_ds_stride, ds_stride, 
+                                                      self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
+            # attention_mask: b t h w
+            input_ids, cond_mask = input_ids_vid.squeeze(1), cond_mask_vid.squeeze(1)  # b 1 l -> b l
+        else:
+            pad_batch_tubes_vid, attention_mask_vid = self.process(batch_tubes_vid, t_ds_stride, ds_stride, 
+                                                                   self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
+            # attention_mask_vid: b t h w
+            pad_batch_tubes_img, attention_mask_img = self.process(batch_tubes_img, 1, ds_stride, 
+                                                                   self.max_1hw, self.ae_stride_1hw, self.patch_size_1hw, extra_1=False)
+            pad_batch_tubes_img = rearrange(pad_batch_tubes_img, '(b i) c 1 h w -> b c i h w', i=self.use_image_num)
+            attention_mask_img = rearrange(attention_mask_img, '(b i) 1 h w -> b i h w', i=self.use_image_num)
+            pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=2)  # concat at temporal, video first
+            # attention_mask_img: b num_img h w
+            attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1)  # b t+num_img h w
+            input_ids = torch.cat([input_ids_vid, input_ids_img], dim=1)  # b 1+num_img hw
+            cond_mask = torch.cat([cond_mask_vid, cond_mask_img], dim=1)  # b 1+num_img hw
+        return pad_batch_tubes, attention_mask, input_ids, cond_mask
+
+    def process(self, batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw, extra_1):
+        
+        # pad to max multiple of ds_stride
+        batch_input_size = [i.shape for i in batch_tubes]  # [(c t h w), (c t h w)]
+        max_t, max_h, max_w = max_thw
+        pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1 if extra_1 else max_t, t_ds_stride), \
+                                          pad_to_multiple(max_h, ds_stride), \
+                                          pad_to_multiple(max_w, ds_stride)
+        pad_max_t = pad_max_t + 1 if extra_1 else pad_max_t
+        each_pad_t_h_w = [[pad_max_t - i.shape[1],
+                           pad_max_h - i.shape[2],
+                           pad_max_w - i.shape[3]] for i in batch_tubes]
+        pad_batch_tubes = [F.pad(im,
+                                 (0, pad_w,
+                                  0, pad_h,
+                                  0, pad_t), value=0) for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)]
+        pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0)
+
+        # make attention_mask
+        # first_channel_first_frame, first_channel_other_frame = pad_batch_tubes[:, :1, :1], pad_batch_tubes[:, :1, 1:]  # first channel to make attention_mask
+        # attention_mask_first_frame = F.max_pool3d(first_channel_first_frame, kernel_size=(1, *ae_stride_thw[1:]), stride=(1, *ae_stride_thw[1:]))
+        # if first_channel_other_frame.numel() != 0:
+        #     attention_mask_other_frame = F.max_pool3d(first_channel_other_frame, kernel_size=ae_stride_thw, stride=ae_stride_thw)
+        #     attention_mask = torch.cat([attention_mask_first_frame, attention_mask_other_frame], dim=2)
+        # else:
+        #     attention_mask = attention_mask_first_frame
+        # attention_mask_ = attention_mask[:, 0].bool().float()  # b t h w, do not channel
+
+        # import ipdb;ipdb.set_trace()
+        max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
+        max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
+                           max_tube_size[1] // ae_stride_thw[1],
+                           max_tube_size[2] // ae_stride_thw[2]]
+        valid_latent_size = [[int(math.ceil((i[1]-1) / ae_stride_thw[0])) + 1 if extra_1 else int(math.ceil(i[1] / ae_stride_thw[0])),
+                            int(math.ceil(i[2] / ae_stride_thw[1])),
+                            int(math.ceil(i[3] / ae_stride_thw[2]))] for i in batch_input_size]
+        attention_mask = [F.pad(torch.ones(i),
+                                (0, max_latent_size[2] - i[2],
+                                 0, max_latent_size[1] - i[1],
+                                 0, max_latent_size[0] - i[0]), value=0) for i in valid_latent_size]
+        attention_mask = torch.stack(attention_mask)  # b t h w
+
+
+        # max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
+        # max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
+        #                    max_tube_size[1] // ae_stride_thw[1],
+        #                    max_tube_size[2] // ae_stride_thw[2]]
+        # max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]),
+        #                             max_latent_size[1] // patch_size_thw[1],
+        #                             max_latent_size[2] // patch_size_thw[2]]
+        # valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)),
+        #                                int(math.ceil(i[2] / ds_stride)),
+        #                                int(math.ceil(i[3] / ds_stride))] for i in batch_input_size]
+        # attention_mask = [F.pad(torch.ones(i),
+        #                         (0, max_patchify_latent_size[2] - i[2],
+        #                          0, max_patchify_latent_size[1] - i[1],
+        #                          0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size]
+        # attention_mask = torch.stack(attention_mask)  # b t h w
+
+        return pad_batch_tubes, attention_mask
diff --git a/opensora/models/causalvideovae/utils/downloader.py b/opensora/models/causalvideovae/utils/downloader.py
new file mode 100755
index 000000000..f2ac4017b
--- /dev/null
+++ b/opensora/models/causalvideovae/utils/downloader.py
@@ -0,0 +1,18 @@
+import gdown
+import os
+
+opensora_cache_home = os.path.expanduser(
+    os.getenv("OPENSORA_HOME", os.path.join("~/.cache", "opensora"))
+)
+
+
+def gdown_download(id, fname, cache_dir=None):
+    cache_dir = opensora_cache_home if not cache_dir else cache_dir
+
+    os.makedirs(cache_dir, exist_ok=True)
+    destination = os.path.join(cache_dir, fname)
+    if os.path.exists(destination):
+        return destination
+
+    gdown.download(id=id, output=destination, quiet=False)
+    return destination
diff --git a/opensora/models/causalvideovae/utils/utils.py b/opensora/models/causalvideovae/utils/utils.py
new file mode 100755
index 000000000..f9aa60754
--- /dev/null
+++ b/opensora/models/causalvideovae/utils/utils.py
@@ -0,0 +1,627 @@
+import os
+
+import torch
+
+import os
+import math
+import torch
+import logging
+import random
+import subprocess
+import numpy as np
+import torch.distributed as dist
+
+# from torch._six import inf
+from torch import inf
+from PIL import Image
+from typing import Union, Iterable
+from collections import OrderedDict
+from torch.utils.tensorboard import SummaryWriter
+
+from diffusers.utils import is_bs4_available, is_ftfy_available
+
+import html
+import re
+import urllib.parse as ul
+
+from torch.utils.data import Dataset
+from torchvision.transforms import Lambda, Compose
+from decord import VideoReader, cpu
+
+from torchvision.transforms._transforms_video import CenterCropVideo
+from torch.nn import functional as F
+import cv2
+import numpy.typing as npt
+
+if is_bs4_available():
+    from bs4 import BeautifulSoup
+
+if is_ftfy_available():
+    import ftfy
+
+_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]
+
+def find_model(model_name):
+    """
+    Finds a pre-trained Latte model, downloading it if necessary. Alternatively, loads a model from a local path.
+    """
+    assert os.path.isfile(model_name), f'Could not find Latte checkpoint at {model_name}'
+    checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
+
+    # if "ema" in checkpoint:  # supports checkpoints from train.py
+    #     print('Using Ema!')
+    #     checkpoint = checkpoint["ema"]
+    # else:
+    print('Using model!')
+    checkpoint = checkpoint['model']
+    return checkpoint
+
+#################################################################################
+#                             Training Clip Gradients                           #
+#################################################################################
+import deepspeed
+def print_grad_norm(model):
+    # 计算并打印梯度范数
+    # model_engine = accelerator.deepspeed_engine_wrapped.engine
+    # gradients = model_engine.get_gradients()
+    # grad_norm = get_grad_norm(gradients)
+    # 计算并打印梯度范数
+    grad_norm = 0
+    n_grad = 0
+    for name, param in model.named_parameters():
+        grad_data = deepspeed.utils.safe_get_full_grad(param)
+        # self.print_tensor_stats(grad_data, name=name)
+
+        if grad_data is not None:
+            param_norm = grad_data.norm(2)
+            grad_norm += param_norm.item() ** 2
+            n_grad += 1
+    grad_norm = (grad_norm / n_grad) ** (1. / 2)
+
+    # self.print_msg('=' * 50)
+    print(f'Gradient Norm is : {grad_norm}')
+
+def get_grad_norm(
+        parameters: _tensor_or_tensors, norm_type: float = 2.0) -> torch.Tensor:
+    r"""
+    Copy from torch.nn.utils.clip_grad_norm_
+
+    Clips gradient norm of an iterable of parameters.
+
+    The norm is computed over all gradients together, as if they were
+    concatenated into a single vector. Gradients are modified in-place.
+
+    Args:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        max_norm (float or int): max norm of the gradients
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+        error_if_nonfinite (bool): if True, an error is thrown if the total
+            norm of the gradients from :attr:`parameters` is ``nan``,
+            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+
+    Returns:
+        Total norm of the parameter gradients (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    grads = [p.grad for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(grads) == 0:
+        return torch.tensor(0.)
+    device = grads[0].device
+    if norm_type == inf:
+        norms = [g.detach().abs().max().to(device) for g in grads]
+        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
+    return total_norm
+
+
+def clip_grad_norm_(
+        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
+        error_if_nonfinite: bool = False, clip_grad=True) -> torch.Tensor:
+    r"""
+    Copy from torch.nn.utils.clip_grad_norm_
+
+    Clips gradient norm of an iterable of parameters.
+
+    The norm is computed over all gradients together, as if they were
+    concatenated into a single vector. Gradients are modified in-place.
+
+    Args:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        max_norm (float or int): max norm of the gradients
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+        error_if_nonfinite (bool): if True, an error is thrown if the total
+            norm of the gradients from :attr:`parameters` is ``nan``,
+            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+
+    Returns:
+        Total norm of the parameter gradients (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    grads = [p.grad for p in parameters if p.grad is not None]
+    max_norm = float(max_norm)
+    norm_type = float(norm_type)
+    if len(grads) == 0:
+        return torch.tensor(0.)
+    device = grads[0].device
+    if norm_type == inf:
+        norms = [g.detach().abs().max().to(device) for g in grads]
+        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
+
+    if clip_grad:
+        if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
+            raise RuntimeError(
+                f'The total norm of order {norm_type} for gradients from '
+                '`parameters` is non-finite, so it cannot be clipped. To disable '
+                'this error and scale the gradients by the non-finite norm anyway, '
+                'set `error_if_nonfinite=False`')
+        clip_coef = max_norm / (total_norm + 1e-6)
+        # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
+        # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
+        # when the gradients do not reside in CPU memory.
+        clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
+        for g in grads:
+            g.detach().mul_(clip_coef_clamped.to(g.device))
+        # gradient_cliped = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
+        # print(gradient_cliped)
+    return total_norm
+
+
+def get_experiment_dir(root_dir, args):
+    # if args.pretrained is not None and 'Latte-XL-2-256x256.pt' not in args.pretrained:
+    #     root_dir += '-WOPRE'
+    if args.use_compile:
+        root_dir += '-Compile'  # speedup by torch compile
+    if args.attention_mode:
+        root_dir += f'-{args.attention_mode.upper()}'
+    # if args.enable_xformers_memory_efficient_attention:
+    #     root_dir += '-Xfor'
+    if args.gradient_checkpointing:
+        root_dir += '-Gc'
+    if args.mixed_precision:
+        root_dir += f'-{args.mixed_precision.upper()}'
+    root_dir += f'-{args.max_image_size}'
+    return root_dir
+
+def get_precision(args):
+    if args.mixed_precision == "bf16":
+        dtype = torch.bfloat16
+    elif args.mixed_precision == "fp16":
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+    return dtype
+
+#################################################################################
+#                             Training Logger                                   #
+#################################################################################
+
+def create_logger(logging_dir):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    if dist.get_rank() == 0:  # real logger
+        logging.basicConfig(
+            level=logging.INFO,
+            # format='[\033[34m%(asctime)s\033[0m] %(message)s',
+            format='[%(asctime)s] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S',
+            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+        )
+        logger = logging.getLogger(__name__)
+
+    else:  # dummy logger (does nothing)
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+def create_tensorboard(tensorboard_dir):
+    """
+    Create a tensorboard that saves losses.
+    """
+    if dist.get_rank() == 0:  # real tensorboard
+        # tensorboard
+        writer = SummaryWriter(tensorboard_dir)
+
+    return writer
+
+
+def write_tensorboard(writer, *args):
+    '''
+    write the loss information to a tensorboard file.
+    Only for pytorch DDP mode.
+    '''
+    if dist.get_rank() == 0:  # real tensorboard
+        writer.add_scalar(args[0], args[1], args[2])
+
+
+#################################################################################
+#                      EMA Update/ DDP Training Utils                           #
+#################################################################################
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def cleanup():
+    """
+    End DDP training.
+    """
+    dist.destroy_process_group()
+
+
+def setup_distributed(backend="nccl", port=None):
+    """Initialize distributed training environment.
+    support both slurm and torch.distributed.launch
+    see torch.distributed.init_process_group() for more details
+    """
+    num_gpus = torch.cuda.device_count()
+
+    if "SLURM_JOB_ID" in os.environ:
+        rank = int(os.environ["SLURM_PROCID"])
+        world_size = int(os.environ["SLURM_NTASKS"])
+        node_list = os.environ["SLURM_NODELIST"]
+        addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1")
+        # specify master port
+        if port is not None:
+            os.environ["MASTER_PORT"] = str(port)
+        elif "MASTER_PORT" not in os.environ:
+            # os.environ["MASTER_PORT"] = "29566"
+            os.environ["MASTER_PORT"] = str(29567 + num_gpus)
+        if "MASTER_ADDR" not in os.environ:
+            os.environ["MASTER_ADDR"] = addr
+        os.environ["WORLD_SIZE"] = str(world_size)
+        os.environ["LOCAL_RANK"] = str(rank % num_gpus)
+        os.environ["RANK"] = str(rank)
+    else:
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+
+    # torch.cuda.set_device(rank % num_gpus)
+
+    dist.init_process_group(
+        backend=backend,
+        world_size=world_size,
+        rank=rank,
+    )
+
+
+#################################################################################
+#                             Testing  Utils                                    #
+#################################################################################
+
+def save_video_grid(video, nrow=None):
+    b, t, h, w, c = video.shape
+
+    if nrow is None:
+        nrow = math.ceil(math.sqrt(b))
+    ncol = math.ceil(b / nrow)
+    padding = 1
+    video_grid = torch.zeros((t, (padding + h) * nrow + padding,
+                              (padding + w) * ncol + padding, c), dtype=torch.uint8)
+
+    print(video_grid.shape)
+    for i in range(b):
+        r = i // ncol
+        c = i % ncol
+        start_r = (padding + h) * r
+        start_c = (padding + w) * c
+        video_grid[:, start_r:start_r + h, start_c:start_c + w] = video[i]
+
+    return video_grid
+
+
+#################################################################################
+#                             MMCV  Utils                                    #
+#################################################################################
+
+
+def collect_env():
+    # Copyright (c) OpenMMLab. All rights reserved.
+    from mmcv.utils import collect_env as collect_base_env
+    from mmcv.utils import get_git_hash
+    """Collect the information of the running environments."""
+
+    env_info = collect_base_env()
+    env_info['MMClassification'] = get_git_hash()[:7]
+
+    for name, val in env_info.items():
+        print(f'{name}: {val}')
+
+    print(torch.cuda.get_arch_list())
+    print(torch.version.cuda)
+
+
+#################################################################################
+#                          Pixart-alpha  Utils                                  #
+#################################################################################
+
+
+bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
+
+def text_preprocessing(text):
+    # The exact text cleaning as was in the training stage:
+    text = clean_caption(text)
+    text = clean_caption(text)
+    return text
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+def clean_caption(caption):
+    caption = str(caption)
+    caption = ul.unquote_plus(caption)
+    caption = caption.strip().lower()
+    caption = re.sub('<person>', 'person', caption)
+    # urls:
+    caption = re.sub(
+        r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+        '', caption)  # regex for urls
+    caption = re.sub(
+        r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+        '', caption)  # regex for urls
+    # html:
+    caption = BeautifulSoup(caption, features='html.parser').text
+
+    # @<nickname>
+    caption = re.sub(r'@[\w\d]+\b', '', caption)
+
+    # 31C0—31EF CJK Strokes
+    # 31F0—31FF Katakana Phonetic Extensions
+    # 3200—32FF Enclosed CJK Letters and Months
+    # 3300—33FF CJK Compatibility
+    # 3400—4DBF CJK Unified Ideographs Extension A
+    # 4DC0—4DFF Yijing Hexagram Symbols
+    # 4E00—9FFF CJK Unified Ideographs
+    caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
+    caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
+    caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
+    caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
+    caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
+    caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
+    caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
+    #######################################################
+
+    # все виды тире / all types of dash --> "-"
+    caption = re.sub(
+        r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
+        '-', caption)
+
+    # кавычки к одному стандарту
+    caption = re.sub(r'[`´«»“”¨]', '"', caption)
+    caption = re.sub(r'[‘’]', "'", caption)
+
+    # &quot;
+    caption = re.sub(r'&quot;?', '', caption)
+    # &amp
+    caption = re.sub(r'&amp', '', caption)
+
+    # ip adresses:
+    caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
+
+    # article ids:
+    caption = re.sub(r'\d:\d\d\s+$', '', caption)
+
+    # \n
+    caption = re.sub(r'\\n', ' ', caption)
+
+    # "#123"
+    caption = re.sub(r'#\d{1,3}\b', '', caption)
+    # "#12345.."
+    caption = re.sub(r'#\d{5,}\b', '', caption)
+    # "123456.."
+    caption = re.sub(r'\b\d{6,}\b', '', caption)
+    # filenames:
+    caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
+
+    #
+    caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
+    caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
+
+    caption = re.sub(bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+    caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
+
+    # this-is-my-cute-cat / this_is_my_cute_cat
+    regex2 = re.compile(r'(?:\-|\_)')
+    if len(re.findall(regex2, caption)) > 3:
+        caption = re.sub(regex2, ' ', caption)
+
+    caption = basic_clean(caption)
+
+    caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
+    caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
+    caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
+
+    caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
+    caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
+    caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
+    caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
+    caption = re.sub(r'\bpage\s+\d+\b', '', caption)
+
+    caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
+
+    caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
+
+    caption = re.sub(r'\b\s+\:\s+', r': ', caption)
+    caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
+    caption = re.sub(r'\s+', ' ', caption)
+
+    caption.strip()
+
+    caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
+    caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
+    caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
+    caption = re.sub(r'^\.\S+$', '', caption)
+
+    return caption.strip()
+
+
+#################################################################################
+#                          eval PSNR when training                                  #
+#################################################################################
+
+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)),
+            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]
+
+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
+
+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
\ No newline at end of file
diff --git a/opensora/models/diffusion/__init__.py b/opensora/models/diffusion/__init__.py
new file mode 100644
index 000000000..6ef61842a
--- /dev/null
+++ b/opensora/models/diffusion/__init__.py
@@ -0,0 +1,21 @@
+
+from .opensora.modeling_opensora import OpenSora_models
+from .udit.modeling_udit import UDiT_models
+from .opensora.modeling_inpaint import OpenSoraInpaint_models
+
+
+Diffusion_models = {}
+Diffusion_models.update(OpenSora_models)
+Diffusion_models.update(UDiT_models)
+Diffusion_models.update(OpenSoraInpaint_models)
+
+from .opensora.modeling_opensora import OpenSora_models_class
+from .udit.modeling_udit import UDiT_models_class
+from .opensora.modeling_inpaint import OpenSoraInpaint_models_class
+
+
+Diffusion_models_class = {}
+Diffusion_models_class.update(OpenSora_models_class)
+Diffusion_models_class.update(UDiT_models_class)
+Diffusion_models_class.update(OpenSoraInpaint_models_class)
+    
\ No newline at end of file
diff --git a/opensora/models/diffusion/opensora/__init__.py b/opensora/models/diffusion/opensora/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/models/diffusion/opensora/modeling_inpaint.py b/opensora/models/diffusion/opensora/modeling_inpaint.py
new file mode 100644
index 000000000..38f7b9c4e
--- /dev/null
+++ b/opensora/models/diffusion/opensora/modeling_inpaint.py
@@ -0,0 +1,373 @@
+import os
+import numpy as np
+from torch import nn
+import torch
+from einops import rearrange, repeat
+from typing import Any, Dict, Optional, Tuple
+from torch.nn import functional as F
+from diffusers.models.transformer_2d import Transformer2DModelOutput
+from diffusers.utils import is_torch_version, deprecate
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormSingle
+from diffusers.models.embeddings import PixArtAlphaTextProjection
+
+from opensora.models.diffusion.opensora.modules import OverlapPatchEmbed3D, OverlapPatchEmbed2D, PatchEmbed2D, BasicTransformerBlock
+from opensora.utils.utils import to_2tuple
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+    from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+from PIL import Image
+import numpy as np
+from enum import Enum, auto
+import glob
+
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+
+def zero_module(module):
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module
+
+def reconstitute_checkpoint(pretrained_checkpoint, model_state_dict):
+    pretrained_keys = set(list(pretrained_checkpoint.keys()))
+    model_keys = set(list(model_state_dict.keys()))
+    common_keys = list(pretrained_keys & model_keys)
+    checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()}
+    return checkpoint
+
+
+class OpenSoraInpaint(OpenSoraT2V):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        cross_attention_dim: Optional[int] = None,
+        attention_bias: bool = False,
+        sample_size: Optional[int] = None,
+        sample_size_t: Optional[int] = None,
+        num_vector_embeds: Optional[int] = None,
+        patch_size: Optional[int] = None,
+        patch_size_t: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        attention_type: str = "default",
+        caption_channels: int = None,
+        interpolation_scale_h: float = None,
+        interpolation_scale_w: float = None,
+        interpolation_scale_t: float = None,
+        use_additional_conditions: Optional[bool] = None,
+        attention_mode: str = 'xformers', 
+        downsampler: str = None, 
+        use_rope: bool = False,
+        use_stable_fp32: bool = False,
+        # inpaint
+        vae_scale_factor_t: int = 4,
+    ):
+        super().__init__(
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            num_layers=num_layers,
+            dropout=dropout,
+            norm_num_groups=norm_num_groups,
+            cross_attention_dim=cross_attention_dim,
+            attention_bias=attention_bias,
+            sample_size=sample_size,
+            sample_size_t=sample_size_t,
+            num_vector_embeds=num_vector_embeds,
+            patch_size=patch_size,
+            patch_size_t=patch_size_t,
+            activation_fn=activation_fn,
+            num_embeds_ada_norm=num_embeds_ada_norm,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            double_self_attention=double_self_attention,
+            upcast_attention=upcast_attention,
+            norm_type=norm_type,
+            norm_elementwise_affine=norm_elementwise_affine,
+            norm_eps=norm_eps,
+            attention_type=attention_type,
+            caption_channels=caption_channels,
+            interpolation_scale_h=interpolation_scale_h,
+            interpolation_scale_w=interpolation_scale_w,
+            interpolation_scale_t=interpolation_scale_t,
+            use_additional_conditions=use_additional_conditions,
+            attention_mode=attention_mode,
+            downsampler=downsampler,
+            use_rope=use_rope,
+            use_stable_fp32=use_stable_fp32,
+        )
+
+        self.vae_scale_factor_t = vae_scale_factor_t
+        # init masked_video and mask conv_in
+        self._init_patched_inputs_for_inpainting()
+
+    def _init_patched_inputs_for_inpainting(self):
+
+        assert self.config.sample_size_t is not None, "OpenSoraInpaint over patched input must provide sample_size_t"
+        assert self.config.sample_size is not None, "OpenSoraInpaint over patched input must provide sample_size"
+        #assert not (self.config.sample_size_t == 1 and self.config.patch_size_t == 2), "Image do not need patchfy in t-dim"
+
+        self.num_frames = self.config.sample_size_t
+        self.config.sample_size = to_2tuple(self.config.sample_size)
+        self.height = self.config.sample_size[0]
+        self.width = self.config.sample_size[1]
+        self.patch_size_t = self.config.patch_size_t
+        self.patch_size = self.config.patch_size
+        interpolation_scale_t = ((self.config.sample_size_t - 1) // 16 + 1) if self.config.sample_size_t % 2 == 1 else self.config.sample_size_t / 16
+        interpolation_scale_t = (
+            self.config.interpolation_scale_t if self.config.interpolation_scale_t is not None else interpolation_scale_t
+        )
+        interpolation_scale = (
+            self.config.interpolation_scale_h if self.config.interpolation_scale_h is not None else self.config.sample_size[0] / 30, 
+            self.config.interpolation_scale_w if self.config.interpolation_scale_w is not None else self.config.sample_size[1] / 40, 
+        )
+        
+        if self.config.downsampler is not None and len(self.config.downsampler) == 9:
+            self.pos_embed_mask = nn.ModuleList(
+                [
+                    OverlapPatchEmbed3D(
+                        num_frames=self.config.sample_size_t,
+                        height=self.config.sample_size[0],
+                        width=self.config.sample_size[1],
+                        patch_size_t=self.config.patch_size_t,
+                        patch_size=self.config.patch_size,
+                        in_channels=self.in_channels,
+                        embed_dim=self.inner_dim,
+                        interpolation_scale=interpolation_scale, 
+                        interpolation_scale_t=interpolation_scale_t,
+                        use_abs_pos=not self.config.use_rope, 
+                    ),
+                    zero_module(nn.Linear(self.inner_dim, self.inner_dim, bias=False)),
+                ]
+            )
+            self.pos_embed_masked_video = nn.ModuleList(
+                [
+                    OverlapPatchEmbed3D(
+                        num_frames=self.config.sample_size_t,
+                        height=self.config.sample_size[0],
+                        width=self.config.sample_size[1],
+                        patch_size_t=self.config.patch_size_t,
+                        patch_size=self.config.patch_size,
+                        in_channels=self.in_channels,
+                        embed_dim=self.inner_dim,
+                        interpolation_scale=interpolation_scale, 
+                        interpolation_scale_t=interpolation_scale_t,
+                        use_abs_pos=not self.config.use_rope, 
+                    ),
+                    zero_module(nn.Linear(self.inner_dim, self.inner_dim, bias=False)),
+                ]
+            )
+        elif self.config.downsampler is not None and len(self.config.downsampler) == 7:
+            self.pos_embed_mask = nn.ModuleList(
+                [
+                    OverlapPatchEmbed2D(
+                        num_frames=self.config.sample_size_t,
+                        height=self.config.sample_size[0],
+                        width=self.config.sample_size[1],
+                        patch_size_t=self.config.patch_size_t,
+                        patch_size=self.config.patch_size,
+                        in_channels=self.in_channels,
+                        embed_dim=self.inner_dim,
+                        interpolation_scale=interpolation_scale, 
+                        interpolation_scale_t=interpolation_scale_t,
+                        use_abs_pos=not self.config.use_rope, 
+                    ),
+                    zero_module(nn.Linear(self.inner_dim, self.inner_dim, bias=False)),
+                ]
+            )
+            
+            self.pos_embed_masked_video = nn.ModuleList(
+                [
+                    OverlapPatchEmbed2D(
+                        num_frames=self.config.sample_size_t,
+                        height=self.config.sample_size[0],
+                        width=self.config.sample_size[1],
+                        patch_size_t=self.config.patch_size_t,
+                        patch_size=self.config.patch_size,
+                        in_channels=self.in_channels,
+                        embed_dim=self.inner_dim,
+                        interpolation_scale=interpolation_scale, 
+                        interpolation_scale_t=interpolation_scale_t,
+                        use_abs_pos=not self.config.use_rope, 
+                    ),
+                    zero_module(nn.Linear(self.inner_dim, self.inner_dim, bias=False)),
+                ]
+            )
+        
+        else:
+            self.pos_embed_mask = nn.ModuleList(
+                [
+                    PatchEmbed2D(
+                        num_frames=self.config.sample_size_t,
+                        height=self.config.sample_size[0],
+                        width=self.config.sample_size[1],
+                        patch_size_t=self.config.patch_size_t,
+                        patch_size=self.config.patch_size,
+                        in_channels=self.in_channels,
+                        embed_dim=self.inner_dim,
+                        interpolation_scale=interpolation_scale, 
+                        interpolation_scale_t=interpolation_scale_t,
+                        use_abs_pos=not self.config.use_rope, 
+                    ),
+                    zero_module(nn.Linear(self.inner_dim, self.inner_dim, bias=False)),
+                ]
+            )
+            self.pos_embed_masked_video = nn.ModuleList(
+                [
+                    PatchEmbed2D(
+                        num_frames=self.config.sample_size_t,
+                        height=self.config.sample_size[0],
+                        width=self.config.sample_size[1],
+                        patch_size_t=self.config.patch_size_t,
+                        patch_size=self.config.patch_size,
+                        in_channels=self.in_channels,
+                        embed_dim=self.inner_dim,
+                        interpolation_scale=interpolation_scale, 
+                        interpolation_scale_t=interpolation_scale_t,
+                        use_abs_pos=not self.config.use_rope, 
+                    ),
+                    zero_module(nn.Linear(self.inner_dim, self.inner_dim, bias=False)),
+                ]
+            )
+
+    def _operate_on_patched_inputs(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num, vip_tokens=None):
+        # inpaint
+        assert hidden_states.shape[1] == 2 * self.config.in_channels + self.vae_scale_factor_t
+        assert hidden_states.shape[2] > 1, "OpenSoraInpaint only supports video input"
+        in_channels = self.config.in_channels
+        hidden_states, hidden_states_masked_vid, hidden_states_mask = hidden_states[:, :in_channels], hidden_states[:, in_channels: 2 * in_channels], hidden_states[:, 2 * in_channels:]
+
+        hidden_states_vid, hidden_states_img = self.pos_embed(hidden_states.to(self.dtype), frame)
+        hidden_states_masked_vid, _ = self.pos_embed_masked_video[0](hidden_states_masked_vid.to(self.dtype), frame)
+        hidden_states_masked_vid = self.pos_embed_masked_video[1](hidden_states_masked_vid)
+
+        hidden_states_mask, _ = self.pos_embed_mask[0](hidden_states_mask.to(self.dtype), frame)
+        hidden_states_mask = self.pos_embed_mask[1](hidden_states_mask)
+
+        hidden_states_vid = hidden_states_vid + hidden_states_masked_vid + hidden_states_mask
+        
+        timestep_vid, timestep_img = None, None
+        embedded_timestep_vid, embedded_timestep_img = None, None
+        encoder_hidden_states_vid, encoder_hidden_states_img = None, None
+
+        if self.adaln_single is not None:
+            if self.use_additional_conditions and added_cond_kwargs is None:
+                raise ValueError(
+                    "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`."
+                )
+            timestep, embedded_timestep = self.adaln_single(
+                timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
+            )  # b 6d, b d
+            if hidden_states_vid is None:
+                timestep_img = timestep
+                embedded_timestep_img = embedded_timestep
+            else:
+                timestep_vid = timestep
+                embedded_timestep_vid = embedded_timestep
+                if hidden_states_img is not None:
+                    timestep_img = repeat(timestep, 'b d -> (b i) d', i=use_image_num).contiguous()
+                    embedded_timestep_img = repeat(embedded_timestep, 'b d -> (b i) d', i=use_image_num).contiguous()
+
+        if self.caption_projection is not None:
+            encoder_hidden_states = self.caption_projection(encoder_hidden_states)  # b, 1+use_image_num, l, d or b, 1, l, d
+
+            if hidden_states_vid is None:
+                encoder_hidden_states_img = rearrange(encoder_hidden_states, 'b 1 l d -> (b 1) l d')
+            else:
+                encoder_hidden_states_vid = rearrange(encoder_hidden_states[:, :1], 'b 1 l d -> (b 1) l d')
+                if hidden_states_img is not None:
+                    encoder_hidden_states_img = rearrange(encoder_hidden_states[:, 1:], 'b i l d -> (b i) l d')
+
+
+        return hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img
+    
+    def transformer_model_custom_load_state_dict(self, pretrained_model_path, load_mask=False):
+        pretrained_model_path = os.path.join(pretrained_model_path, 'diffusion_pytorch_model.*')
+        pretrained_model_path = glob.glob(pretrained_model_path)
+        assert len(pretrained_model_path) > 0, f"Cannot find pretrained model in {pretrained_model_path}"
+        pretrained_model_path = pretrained_model_path[0]
+
+        print(f'Loading {self.__class__.__name__} pretrained weights...')
+        print(f'Loading pretrained model from {pretrained_model_path}...')
+        model_state_dict = self.state_dict()
+        if 'safetensors' in pretrained_model_path:  # pixart series
+            from safetensors.torch import load_file as safe_load
+            # import ipdb;ipdb.set_trace()
+            pretrained_checkpoint = safe_load(pretrained_model_path, device="cpu")
+        else:  # latest stage training weight
+            pretrained_checkpoint = torch.load(pretrained_model_path, map_location='cpu')
+            if 'model' in pretrained_checkpoint:
+                pretrained_checkpoint = pretrained_checkpoint['model']
+        checkpoint = reconstitute_checkpoint(pretrained_checkpoint, model_state_dict)
+
+        if not 'pos_embed_masked_video.0.weight' in checkpoint:
+            checkpoint['pos_embed_masked_video.0.proj.weight'] = checkpoint['pos_embed.proj.weight']
+            checkpoint['pos_embed_masked_video.0.proj.bias'] = checkpoint['pos_embed.proj.bias']
+        if not 'pos_embed_mask.0.proj.weight' in checkpoint and load_mask:
+            checkpoint['pos_embed_mask.0.proj.weight'] = checkpoint['pos_embed.proj.weight']
+            checkpoint['pos_embed_mask.0.proj.bias'] = checkpoint['pos_embed.proj.bias']
+
+        missing_keys, unexpected_keys = self.load_state_dict(checkpoint, strict=False)
+        print(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}')
+        print(f'Successfully load {len(self.state_dict()) - len(missing_keys)}/{len(model_state_dict)} keys from {pretrained_model_path}!')
+
+    def custom_load_state_dict(self, pretrained_model_path, load_mask=False):
+        assert isinstance(pretrained_model_path, dict), "pretrained_model_path must be a dict"
+
+        pretrained_transformer_model_path = pretrained_model_path.get('transformer_model', None)
+
+        self.transformer_model_custom_load_state_dict(pretrained_transformer_model_path, load_mask)
+
+def OpenSoraInpaint_S_122(**kwargs):
+    return OpenSoraInpaint(num_layers=28, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1536, **kwargs)
+
+def OpenSoraInpaint_B_122(**kwargs):
+    return OpenSoraInpaint(num_layers=32, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1920, **kwargs)
+
+def OpenSoraInpaint_L_122(**kwargs):
+    return OpenSoraInpaint(num_layers=40, attention_head_dim=128, num_attention_heads=16, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2048, **kwargs)
+
+def OpenSoraInpaint_ROPE_L_122(**kwargs):
+    return OpenSoraInpaint(num_layers=32, attention_head_dim=96, num_attention_heads=24, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2304, **kwargs)
+
+OpenSoraInpaint_models = {
+    "OpenSoraInpaint-S/122": OpenSoraInpaint_S_122,  #       1.1B
+    "OpenSoraInpaint-B/122": OpenSoraInpaint_B_122,
+    "OpenSoraInpaint-L/122": OpenSoraInpaint_L_122,
+    "OpenSoraInpaint-ROPE-L/122": OpenSoraInpaint_ROPE_L_122,
+}
+
+OpenSoraInpaint_models_class = {
+    "OpenSoraInpaint-S/122": OpenSoraInpaint,
+    "OpenSoraInpaint-B/122": OpenSoraInpaint,
+    "OpenSoraInpaint-L/122": OpenSoraInpaint,
+    "OpenSoraInpaint-ROPE-L/122": OpenSoraInpaint,
+}
\ No newline at end of file
diff --git a/opensora/models/diffusion/opensora/modeling_opensora.py b/opensora/models/diffusion/opensora/modeling_opensora.py
new file mode 100644
index 000000000..612002a66
--- /dev/null
+++ b/opensora/models/diffusion/opensora/modeling_opensora.py
@@ -0,0 +1,734 @@
+import os
+import numpy as np
+from torch import nn
+import torch
+from einops import rearrange, repeat
+from typing import Any, Dict, Optional, Tuple
+from torch.nn import functional as F
+from diffusers.models.transformer_2d import Transformer2DModelOutput
+from diffusers.utils import is_torch_version, deprecate
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormSingle
+from diffusers.models.embeddings import PixArtAlphaTextProjection
+from opensora.models.diffusion.opensora.modules import OverlapPatchEmbed3D, OverlapPatchEmbed2D, PatchEmbed2D, BasicTransformerBlock
+from opensora.utils.utils import to_2tuple
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+    from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+class OpenSoraT2V(ModelMixin, ConfigMixin):
+    """
+    A 2D Transformer model for image-like data.
+
+    Parameters:
+        num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
+        in_channels (`int`, *optional*):
+            The number of channels in the input and output (specify if the input is **continuous**).
+        num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
+            This is fixed during training since it is used to learn a number of position embeddings.
+        num_vector_embeds (`int`, *optional*):
+            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
+            Includes the class for the masked latent pixel.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
+        num_embeds_ada_norm ( `int`, *optional*):
+            The number of diffusion steps used during training. Pass if at least one of the norm_layers is
+            `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
+            added to the hidden states.
+
+            During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
+        attention_bias (`bool`, *optional*):
+            Configure if the `TransformerBlocks` attention should contain a bias parameter.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        cross_attention_dim: Optional[int] = None,
+        attention_bias: bool = False,
+        sample_size: Optional[int] = None,
+        sample_size_t: Optional[int] = None,
+        num_vector_embeds: Optional[int] = None,
+        patch_size: Optional[int] = None,
+        patch_size_t: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        attention_type: str = "default",
+        caption_channels: int = None,
+        interpolation_scale_h: float = None,
+        interpolation_scale_w: float = None,
+        interpolation_scale_t: float = None,
+        use_additional_conditions: Optional[bool] = None,
+        attention_mode: str = 'xformers', 
+        downsampler: str = None, 
+        use_rope: bool = False,
+        use_stable_fp32: bool = False,
+    ):
+        super().__init__()
+
+        # Validate inputs.
+        if patch_size is not None:
+            if norm_type not in ["ada_norm", "ada_norm_zero", "ada_norm_single"]:
+                raise NotImplementedError(
+                    f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'."
+                )
+            elif norm_type in ["ada_norm", "ada_norm_zero"] and num_embeds_ada_norm is None:
+                raise ValueError(
+                    f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None."
+                )
+
+        # Set some common variables used across the board.
+        self.use_rope = use_rope
+        self.use_linear_projection = use_linear_projection
+        self.interpolation_scale_t = interpolation_scale_t
+        self.interpolation_scale_h = interpolation_scale_h
+        self.interpolation_scale_w = interpolation_scale_w
+        self.downsampler = downsampler
+        self.caption_channels = caption_channels
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.gradient_checkpointing = False
+        self.config.hidden_size = self.inner_dim
+        use_additional_conditions = False
+        # if use_additional_conditions is None:
+            # if norm_type == "ada_norm_single" and sample_size == 128:
+            #     use_additional_conditions = True
+            # else:
+            # use_additional_conditions = False
+        self.use_additional_conditions = use_additional_conditions
+
+        # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
+        # Define whether input is continuous or discrete depending on configuration
+        assert in_channels is not None and patch_size is not None
+
+        if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
+            deprecation_message = (
+                f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
+                " incorrectly set to `'layer_norm'`. Make sure to set `norm_type` to `'ada_norm'` in the config."
+                " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
+                " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
+                " would be very nice if you could open a Pull request for the `transformer/config.json` file"
+            )
+            deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
+            norm_type = "ada_norm"
+
+        # 2. Initialize the right blocks.
+        # Initialize the output blocks and other projection blocks when necessary.
+        self._init_patched_inputs(norm_type=norm_type)
+
+    def _init_patched_inputs(self, norm_type):
+        assert self.config.sample_size_t is not None, "OpenSoraT2V over patched input must provide sample_size_t"
+        assert self.config.sample_size is not None, "OpenSoraT2V over patched input must provide sample_size"
+        #assert not (self.config.sample_size_t == 1 and self.config.patch_size_t == 2), "Image do not need patchfy in t-dim"
+
+        self.num_frames = self.config.sample_size_t
+        self.config.sample_size = to_2tuple(self.config.sample_size)
+        self.height = self.config.sample_size[0]
+        self.width = self.config.sample_size[1]
+        self.patch_size_t = self.config.patch_size_t
+        self.patch_size = self.config.patch_size
+        interpolation_scale_t = ((self.config.sample_size_t - 1) // 16 + 1) if self.config.sample_size_t % 2 == 1 else self.config.sample_size_t / 16
+        interpolation_scale_t = (
+            self.config.interpolation_scale_t if self.config.interpolation_scale_t is not None else interpolation_scale_t
+        )
+        interpolation_scale = (
+            self.config.interpolation_scale_h if self.config.interpolation_scale_h is not None else self.config.sample_size[0] / 30, 
+            self.config.interpolation_scale_w if self.config.interpolation_scale_w is not None else self.config.sample_size[1] / 40, 
+        )
+        # if self.config.sample_size_t > 1:
+        #     self.pos_embed = PatchEmbed3D(
+        #         num_frames=self.config.sample_size_t,
+        #         height=self.config.sample_size[0],
+        #         width=self.config.sample_size[1],
+        #         patch_size_t=self.config.patch_size_t,
+        #         patch_size=self.config.patch_size,
+        #         in_channels=self.in_channels,
+        #         embed_dim=self.inner_dim,
+        #         interpolation_scale=interpolation_scale, 
+        #         interpolation_scale_t=interpolation_scale_t,
+        #     )
+        # else:
+        if self.config.downsampler is not None and len(self.config.downsampler) == 9:
+            self.pos_embed = OverlapPatchEmbed3D(
+                num_frames=self.config.sample_size_t,
+                height=self.config.sample_size[0],
+                width=self.config.sample_size[1],
+                patch_size_t=self.config.patch_size_t,
+                patch_size=self.config.patch_size,
+                in_channels=self.in_channels,
+                embed_dim=self.inner_dim,
+                interpolation_scale=interpolation_scale, 
+                interpolation_scale_t=interpolation_scale_t,
+                use_abs_pos=not self.config.use_rope, 
+            )
+        elif self.config.downsampler is not None and len(self.config.downsampler) == 7:
+            self.pos_embed = OverlapPatchEmbed2D(
+                num_frames=self.config.sample_size_t,
+                height=self.config.sample_size[0],
+                width=self.config.sample_size[1],
+                patch_size_t=self.config.patch_size_t,
+                patch_size=self.config.patch_size,
+                in_channels=self.in_channels,
+                embed_dim=self.inner_dim,
+                interpolation_scale=interpolation_scale, 
+                interpolation_scale_t=interpolation_scale_t,
+                use_abs_pos=not self.config.use_rope, 
+            )
+        
+        else:
+            self.pos_embed = PatchEmbed2D(
+                num_frames=self.config.sample_size_t,
+                height=self.config.sample_size[0],
+                width=self.config.sample_size[1],
+                patch_size_t=self.config.patch_size_t,
+                patch_size=self.config.patch_size,
+                in_channels=self.in_channels,
+                embed_dim=self.inner_dim,
+                interpolation_scale=interpolation_scale, 
+                interpolation_scale_t=interpolation_scale_t,
+                use_abs_pos=not self.config.use_rope, 
+            )
+        interpolation_scale_thw = (interpolation_scale_t, *interpolation_scale)
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim,
+                    self.config.num_attention_heads,
+                    self.config.attention_head_dim,
+                    dropout=self.config.dropout,
+                    cross_attention_dim=self.config.cross_attention_dim,
+                    activation_fn=self.config.activation_fn,
+                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+                    attention_bias=self.config.attention_bias,
+                    only_cross_attention=self.config.only_cross_attention,
+                    double_self_attention=self.config.double_self_attention,
+                    upcast_attention=self.config.upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=self.config.norm_elementwise_affine,
+                    norm_eps=self.config.norm_eps,
+                    attention_type=self.config.attention_type,
+                    attention_mode=self.config.attention_mode, 
+                    downsampler=self.config.downsampler, 
+                    use_rope=self.config.use_rope, 
+                    interpolation_scale_thw=interpolation_scale_thw, 
+                )
+                for _ in range(self.config.num_layers)
+            ]
+        )
+
+        if self.config.norm_type != "ada_norm_single":
+            self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+            self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
+            self.proj_out_2 = nn.Linear(
+                self.inner_dim, self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels
+            )
+        elif self.config.norm_type == "ada_norm_single":
+            self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+            self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5)
+            self.proj_out = nn.Linear(
+                self.inner_dim, self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels
+            )
+
+        # PixArt-Alpha blocks.
+        self.adaln_single = None
+        if self.config.norm_type == "ada_norm_single":
+            # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use
+            # additional conditions until we find better name
+            self.adaln_single = AdaLayerNormSingle(
+                self.inner_dim, use_additional_conditions=self.use_additional_conditions
+            )
+
+        self.caption_projection = None
+        if self.caption_channels is not None:
+            self.caption_projection = PixArtAlphaTextProjection(
+                in_features=self.caption_channels, hidden_size=self.inner_dim
+            )
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: Optional[torch.LongTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        added_cond_kwargs: Dict[str, torch.Tensor] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        use_image_num: Optional[int] = 0,
+        return_dict: bool = True,
+    ):
+        """
+        The [`Transformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
+                Input `hidden_states`.
+            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+                self-attention.
+            timestep ( `torch.LongTensor`, *optional*):
+                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
+            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
+                Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
+                `AdaLayerZeroNorm`.
+            cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            attention_mask ( `torch.Tensor`, *optional*):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            encoder_attention_mask ( `torch.Tensor`, *optional*):
+                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+                    * Mask `(batch, sequence_length)` True = keep, False = discard.
+                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+                above. This bias will be added to the cross-attention scores.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        batch_size, c, frame, h, w = hidden_states.shape
+        # print('hidden_states.shape', hidden_states.shape)
+        frame = frame - use_image_num  # 21-4=17
+        if cross_attention_kwargs is not None:
+            if cross_attention_kwargs.get("scale", None) is not None:
+                print.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        attention_mask_vid, attention_mask_img = None, None
+        if attention_mask is not None and attention_mask.ndim == 4:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #   (keep = +0,     discard = -10000.0)
+            # b, frame+use_image_num, h, w -> a video with images
+            # b, 1, h, w -> only images
+            attention_mask = attention_mask.to(self.dtype)
+            if get_sequence_parallel_state():
+                if npu_config is not None:
+                    attention_mask_vid = attention_mask[:, :frame * hccl_info.world_size]  # b, frame, h, w
+                    attention_mask_img = attention_mask[:, frame * hccl_info.world_size:]  # b, use_image_num, h, w
+                else:
+                    # print('before attention_mask.shape', attention_mask.shape)
+                    attention_mask_vid = attention_mask[:, :frame * nccl_info.world_size]  # b, frame, h, w
+                    attention_mask_img = attention_mask[:, frame * nccl_info.world_size:]  # b, use_image_num, h, w
+                    # print('after attention_mask.shape', attention_mask_vid.shape)
+            else:
+                attention_mask_vid = attention_mask[:, :frame]  # b, frame, h, w
+                attention_mask_img = attention_mask[:, frame:]  # b, use_image_num, h, w
+
+            if attention_mask_vid.numel() > 0:
+                attention_mask_vid_first_frame = attention_mask_vid[:, :1].repeat(1, self.patch_size_t-1, 1, 1)
+                attention_mask_vid = torch.cat([attention_mask_vid_first_frame, attention_mask_vid], dim=1)
+                attention_mask_vid = attention_mask_vid.unsqueeze(1)  # b 1 t h w
+                attention_mask_vid = F.max_pool3d(attention_mask_vid, kernel_size=(self.patch_size_t, self.patch_size, self.patch_size), 
+                                                  stride=(self.patch_size_t, self.patch_size, self.patch_size))
+                attention_mask_vid = rearrange(attention_mask_vid, 'b 1 t h w -> (b 1) 1 (t h w)') 
+            if attention_mask_img.numel() > 0:
+                attention_mask_img = F.max_pool2d(attention_mask_img, kernel_size=(self.patch_size, self.patch_size), stride=(self.patch_size, self.patch_size))
+                attention_mask_img = rearrange(attention_mask_img, 'b i h w -> (b i) 1 (h w)') 
+
+            attention_mask_vid = (1 - attention_mask_vid.bool().to(self.dtype)) * -10000.0 if attention_mask_vid.numel() > 0 else None
+            attention_mask_img = (1 - attention_mask_img.bool().to(self.dtype)) * -10000.0 if attention_mask_img.numel() > 0 else None
+
+            if frame == 1 and use_image_num == 0 and not get_sequence_parallel_state():
+                attention_mask_img = attention_mask_vid
+                attention_mask_vid = None
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        # import ipdb;ipdb.set_trace()
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 3:  
+            # b, 1+use_image_num, l -> a video with images
+            # b, 1, l -> only images
+            encoder_attention_mask = (1 - encoder_attention_mask.to(self.dtype)) * -10000.0
+            in_t = encoder_attention_mask.shape[1]
+            encoder_attention_mask_vid = encoder_attention_mask[:, :in_t-use_image_num]  # b, 1, l
+            encoder_attention_mask_vid = rearrange(encoder_attention_mask_vid, 'b 1 l -> (b 1) 1 l') if encoder_attention_mask_vid.numel() > 0 else None
+
+            encoder_attention_mask_img = encoder_attention_mask[:, in_t-use_image_num:]  # b, use_image_num, l
+            encoder_attention_mask_img = rearrange(encoder_attention_mask_img, 'b i l -> (b i) 1 l') if encoder_attention_mask_img.numel() > 0 else None
+
+            if frame == 1 and use_image_num == 0 and not get_sequence_parallel_state():
+                encoder_attention_mask_img = encoder_attention_mask_vid
+                encoder_attention_mask_vid = None
+
+        if npu_config is not None and attention_mask_vid is not None:
+            attention_mask_vid = npu_config.get_attention_mask(attention_mask_vid, attention_mask_vid.shape[-1])
+            encoder_attention_mask_vid = npu_config.get_attention_mask(encoder_attention_mask_vid,
+                                                                       attention_mask_vid.shape[-2])
+        if npu_config is not None and attention_mask_img is not None:
+            attention_mask_img = npu_config.get_attention_mask(attention_mask_img, attention_mask_img.shape[-1])
+            encoder_attention_mask_img = npu_config.get_attention_mask(encoder_attention_mask_img,
+                                                                       attention_mask_img.shape[-2])
+
+
+        # 1. Input
+        frame = ((frame - 1) // self.patch_size_t + 1) if frame % 2 == 1 else frame // self.patch_size_t  # patchfy
+        # print('frame', frame)
+        height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size
+
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+        hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, \
+        timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img = self._operate_on_patched_inputs(
+            hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num
+        )
+        # 2. Blocks
+        # import ipdb;ipdb.set_trace()
+        if get_sequence_parallel_state():
+            if hidden_states_vid is not None:
+                # print(333333333333333)
+                hidden_states_vid = rearrange(hidden_states_vid, 'b s h -> s b h', b=batch_size).contiguous()
+                encoder_hidden_states_vid = rearrange(encoder_hidden_states_vid, 'b s h -> s b h',
+                                                      b=batch_size).contiguous()
+                timestep_vid = timestep_vid.view(batch_size, 6, -1).transpose(0, 1).contiguous()
+                # print('timestep_vid', timestep_vid.shape)
+
+                
+        for block in self.transformer_blocks:
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                # import ipdb;ipdb.set_trace()
+                if hidden_states_vid is not None:
+                    hidden_states_vid = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states_vid,
+                        attention_mask_vid,
+                        encoder_hidden_states_vid,
+                        encoder_attention_mask_vid,
+                        timestep_vid,
+                        cross_attention_kwargs,
+                        class_labels,
+                        frame, 
+                        height, 
+                        width, 
+                        **ckpt_kwargs,
+                    )
+                # import ipdb;ipdb.set_trace()
+                if hidden_states_img is not None:
+                    hidden_states_img = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        hidden_states_img,
+                        attention_mask_img,
+                        encoder_hidden_states_img,
+                        encoder_attention_mask_img,
+                        timestep_img,
+                        cross_attention_kwargs,
+                        class_labels,
+                        1, 
+                        height, 
+                        width, 
+                        **ckpt_kwargs,
+                    )
+            else:
+                if hidden_states_vid is not None:
+                    hidden_states_vid = block(
+                        hidden_states_vid,
+                        attention_mask=attention_mask_vid,
+                        encoder_hidden_states=encoder_hidden_states_vid,
+                        encoder_attention_mask=encoder_attention_mask_vid,
+                        timestep=timestep_vid,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        class_labels=class_labels,
+                        frame=frame, 
+                        height=height, 
+                        width=width, 
+                    )
+                if hidden_states_img is not None:
+                    hidden_states_img = block(
+                        hidden_states_img,
+                        attention_mask=attention_mask_img,
+                        encoder_hidden_states=encoder_hidden_states_img,
+                        encoder_attention_mask=encoder_attention_mask_img,
+                        timestep=timestep_img,
+                        cross_attention_kwargs=cross_attention_kwargs,
+                        class_labels=class_labels,
+                        frame=1, 
+                        height=height, 
+                        width=width, 
+                    )
+
+        if get_sequence_parallel_state():
+            if hidden_states_vid is not None:
+                hidden_states_vid = rearrange(hidden_states_vid, 's b h -> b s h', b=batch_size).contiguous()
+
+        # 3. Output
+        output_vid, output_img = None, None 
+        if hidden_states_vid is not None:
+            output_vid = self._get_output_for_patched_inputs(
+                hidden_states=hidden_states_vid,
+                timestep=timestep_vid,
+                class_labels=class_labels,
+                embedded_timestep=embedded_timestep_vid,
+                num_frames=frame, 
+                height=height,
+                width=width,
+            )  # b c t h w
+        if hidden_states_img is not None:
+            output_img = self._get_output_for_patched_inputs(
+                hidden_states=hidden_states_img,
+                timestep=timestep_img,
+                class_labels=class_labels,
+                embedded_timestep=embedded_timestep_img,
+                num_frames=1, 
+                height=height,
+                width=width,
+            )  # b c 1 h w
+            if use_image_num != 0:
+                output_img = rearrange(output_img, '(b i) c 1 h w -> b c i h w', i=use_image_num)
+
+        if output_vid is not None and output_img is not None:
+            output = torch.cat([output_vid, output_img], dim=2)
+        elif output_vid is not None:
+            output = output_vid
+        elif output_img is not None:
+            output = output_img
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+
+    def _operate_on_patched_inputs(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num):
+        # batch_size = hidden_states.shape[0]
+        hidden_states_vid, hidden_states_img = self.pos_embed(hidden_states.to(self.dtype), frame)
+        timestep_vid, timestep_img = None, None
+        embedded_timestep_vid, embedded_timestep_img = None, None
+        encoder_hidden_states_vid, encoder_hidden_states_img = None, None
+
+        if self.adaln_single is not None:
+            if self.use_additional_conditions and added_cond_kwargs is None:
+                raise ValueError(
+                    "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`."
+                )
+            timestep, embedded_timestep = self.adaln_single(
+                timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
+            )  # b 6d, b d
+            if hidden_states_vid is None:
+                timestep_img = timestep
+                embedded_timestep_img = embedded_timestep
+            else:
+                timestep_vid = timestep
+                embedded_timestep_vid = embedded_timestep
+                if hidden_states_img is not None:
+                    timestep_img = repeat(timestep, 'b d -> (b i) d', i=use_image_num).contiguous()
+                    embedded_timestep_img = repeat(embedded_timestep, 'b d -> (b i) d', i=use_image_num).contiguous()
+
+        if self.caption_projection is not None:
+            encoder_hidden_states = self.caption_projection(encoder_hidden_states)  # b, 1+use_image_num, l, d or b, 1, l, d
+            if hidden_states_vid is None:
+                encoder_hidden_states_img = rearrange(encoder_hidden_states, 'b 1 l d -> (b 1) l d')
+            else:
+                encoder_hidden_states_vid = rearrange(encoder_hidden_states[:, :1], 'b 1 l d -> (b 1) l d')
+                if hidden_states_img is not None:
+                    encoder_hidden_states_img = rearrange(encoder_hidden_states[:, 1:], 'b i l d -> (b i) l d')
+
+
+        return hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img
+
+    
+    
+    def _get_output_for_patched_inputs(
+        self, hidden_states, timestep, class_labels, embedded_timestep, num_frames, height=None, width=None
+    ):  
+        # import ipdb;ipdb.set_trace()
+        if self.config.norm_type != "ada_norm_single":
+            conditioning = self.transformer_blocks[0].norm1.emb(
+                timestep, class_labels, hidden_dtype=self.dtype
+            )
+            shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+            hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+            hidden_states = self.proj_out_2(hidden_states)
+        elif self.config.norm_type == "ada_norm_single":
+            shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1)
+            hidden_states = self.norm_out(hidden_states)
+            # Modulation
+            hidden_states = hidden_states * (1 + scale) + shift
+            hidden_states = self.proj_out(hidden_states)
+            hidden_states = hidden_states.squeeze(1)
+
+        # unpatchify
+        if self.adaln_single is None:
+            height = width = int(hidden_states.shape[1] ** 0.5)
+        hidden_states = hidden_states.reshape(
+            shape=(-1, num_frames, height, width, self.patch_size_t, self.patch_size, self.patch_size, self.out_channels)
+        )
+        hidden_states = torch.einsum("nthwopqc->nctohpwq", hidden_states)
+        output = hidden_states.reshape(
+            shape=(-1, self.out_channels, num_frames * self.patch_size_t, height * self.patch_size, width * self.patch_size)
+        )
+        # import ipdb;ipdb.set_trace()
+        # if output.shape[2] % 2 == 0:
+        #     output = output[:, :, 1:]
+        return output
+     
+def OpenSoraT2V_S_122(**kwargs):
+    return OpenSoraT2V(num_layers=28, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1536, **kwargs)
+
+def OpenSoraT2V_B_122(**kwargs):
+    return OpenSoraT2V(num_layers=32, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1920, **kwargs)
+
+def OpenSoraT2V_L_122(**kwargs):
+    return OpenSoraT2V(num_layers=40, attention_head_dim=128, num_attention_heads=16, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2048, **kwargs)
+
+def OpenSoraT2V_ROPE_L_122(**kwargs):
+    return OpenSoraT2V(num_layers=32, attention_head_dim=96, num_attention_heads=24, patch_size_t=1, patch_size=2,
+                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2304, **kwargs)
+
+OpenSora_models = {
+    "OpenSoraT2V-S/122": OpenSoraT2V_S_122,  #       1.1B
+    "OpenSoraT2V-B/122": OpenSoraT2V_B_122,
+    "OpenSoraT2V-L/122": OpenSoraT2V_L_122,
+    "OpenSoraT2V-ROPE-L/122": OpenSoraT2V_ROPE_L_122,
+}
+
+OpenSora_models_class = {
+    "OpenSoraT2V-S/122": OpenSoraT2V,
+    "OpenSoraT2V-B/122": OpenSoraT2V,
+    "OpenSoraT2V-L/122": OpenSoraT2V,
+    "OpenSoraT2V-ROPE-L/122": OpenSoraT2V,
+}
+
+if __name__ == '__main__':
+    from opensora.models.causalvideovae import ae_channel_config, ae_stride_config
+    from opensora.models.causalvideovae import getae, getae_wrapper
+    from opensora.models.causalvideovae import CausalVAEModelWrapper
+
+    args = type('args', (), 
+    {
+        'ae': 'CausalVAEModel_4x8x8', 
+        'attention_mode': 'xformers', 
+        'use_rope': True, 
+        'model_max_length': 300, 
+        'max_height': 320,
+        'max_width': 240,
+        'num_frames': 1,
+        'use_image_num': 0, 
+        'compress_kv_factor': 1, 
+        'interpolation_scale_t': 1,
+        'interpolation_scale_h': 1,
+        'interpolation_scale_w': 1,
+    }
+    )
+    b = 16
+    c = 8
+    cond_c = 4096
+    num_timesteps = 1000
+    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
+    latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w)
+    num_frames = (args.num_frames - 1) // ae_stride_t + 1
+
+    device = torch.device('cuda:0')
+    model = OpenSoraT2V_ROPE_L_122(in_channels=c, 
+                              out_channels=c, 
+                              sample_size=latent_size, 
+                              sample_size_t=num_frames, 
+                              activation_fn="gelu-approximate",
+                            attention_bias=True,
+                            attention_type="default",
+                            double_self_attention=False,
+                            norm_elementwise_affine=False,
+                            norm_eps=1e-06,
+                            norm_num_groups=32,
+                            num_vector_embeds=None,
+                            only_cross_attention=False,
+                            upcast_attention=False,
+                            use_linear_projection=False,
+                            use_additional_conditions=False, 
+                            downsampler=None,
+                            interpolation_scale_t=args.interpolation_scale_t, 
+                            interpolation_scale_h=args.interpolation_scale_h, 
+                            interpolation_scale_w=args.interpolation_scale_w, 
+                            use_rope=args.use_rope, 
+                            ).to(device)
+    
+    try:
+        path = "PixArt-Alpha-XL-2-512.safetensors"
+        from safetensors.torch import load_file as safe_load
+        ckpt = safe_load(path, device="cpu")
+        # import ipdb;ipdb.set_trace()
+        if ckpt['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and ckpt['pos_embed.proj.weight'].ndim == 4:
+            repeat = model.pos_embed.proj.weight.shape[2]
+            ckpt['pos_embed.proj.weight'] = ckpt['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat)
+            del ckpt['proj_out.weight'], ckpt['proj_out.bias']
+        msg = model.load_state_dict(ckpt, strict=False)
+        print(msg)
+    except Exception as e:
+        print(e)
+    print(model)
+    print(f'{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B')
+    # import sys;sys.exit()
+    x = torch.randn(b, c,  1+(args.num_frames-1)//ae_stride_t+args.use_image_num, args.max_height//ae_stride_h, args.max_width//ae_stride_w).to(device)
+    cond = torch.randn(b, 1+args.use_image_num, args.model_max_length, cond_c).to(device)
+    attn_mask = torch.randint(0, 2, (b, 1+(args.num_frames-1)//ae_stride_t+args.use_image_num, args.max_height//ae_stride_h, args.max_width//ae_stride_w)).to(device)  # B L or B 1+num_images L
+    cond_mask = torch.randint(0, 2, (b, 1+args.use_image_num, args.model_max_length)).to(device)  # B L or B 1+num_images L
+    timestep = torch.randint(0, 1000, (b,), device=device)
+    model_kwargs = dict(hidden_states=x, encoder_hidden_states=cond, attention_mask=attn_mask, 
+                        encoder_attention_mask=cond_mask, use_image_num=args.use_image_num, timestep=timestep)
+    with torch.no_grad():
+        output = model(**model_kwargs)
+    print(output[0].shape)
+
+
+
+
diff --git a/opensora/models/diffusion/opensora/modules.py b/opensora/models/diffusion/opensora/modules.py
new file mode 100644
index 000000000..8904bc0c5
--- /dev/null
+++ b/opensora/models/diffusion/opensora/modules.py
@@ -0,0 +1,1383 @@
+from einops import rearrange
+from torch import nn
+import torch
+import numpy as np
+
+from einops import rearrange, repeat
+from typing import Any, Dict, Optional, Tuple
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from typing import Any, Dict, Optional
+import re
+import torch
+import torch.nn.functional as F
+from torch import nn
+import diffusers
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import FeedForward, GatedSelfAttentionDense
+from diffusers.models.attention_processor import Attention as Attention_
+from diffusers.models.embeddings import SinusoidalPositionalEmbedding
+from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm
+from .rope import PositionGetter3D, RoPE3D
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config, set_run_dtype
+    from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info
+    from opensora.acceleration.communications import all_to_all_SBH
+except:
+    torch_npu = None
+    npu_config = None
+    set_run_dtype = None
+    from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info
+    from opensora.utils.communications import all_to_all_SBH
+logger = logging.get_logger(__name__)
+
+def get_3d_sincos_pos_embed(
+    embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, 
+):
+    """
+    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # if isinstance(grid_size, int):
+    #     grid_size = (grid_size, grid_size)
+    grid_t = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0]
+    grid_h = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1]
+    grid_w = np.arange(grid_size[2], dtype=np.float32) / (grid_size[2] / base_size[2]) / interpolation_scale[2]
+    grid = np.meshgrid(grid_w, grid_h, grid_t)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([3, 1, grid_size[2], grid_size[1], grid_size[0]])
+    pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid)
+    # import ipdb;ipdb.set_trace()
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_3d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 3 != 0:
+        raise ValueError("embed_dim must be divisible by 3")
+
+    # import ipdb;ipdb.set_trace()
+    # use 1/3 of dimensions to encode grid_t/h/w
+    emb_t = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0])  # (T*H*W, D/3)
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1])  # (T*H*W, D/3)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2])  # (T*H*W, D/3)
+
+    emb = np.concatenate([emb_t, emb_h, emb_w], axis=1)  # (T*H*W, D)
+    return emb
+
+
+def get_2d_sincos_pos_embed(
+    embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, 
+):
+    """
+    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # if isinstance(grid_size, int):
+    #     grid_size = (grid_size, grid_size)
+
+    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0]
+    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1]
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    # use 1/3 of dimensions to encode grid_t/h/w
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+def get_1d_sincos_pos_embed(
+    embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, 
+):
+    """
+    grid_size: int of the grid return: pos_embed: [grid_size, embed_dim] or
+    [1+grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # if isinstance(grid_size, int):
+    #     grid_size = (grid_size, grid_size)
+
+    grid = np.arange(grid_size, dtype=np.float32) / (grid_size / base_size) / interpolation_scale
+    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid)  # (H*W, D/2)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+class PatchEmbed2D(nn.Module):
+    """2D Image to Patch Embedding but with 3D position embedding"""
+
+    def __init__(
+        self,
+        num_frames=1, 
+        height=224,
+        width=224,
+        patch_size_t=1,
+        patch_size=16,
+        in_channels=3,
+        embed_dim=768,
+        layer_norm=False,
+        flatten=True,
+        bias=True,
+        interpolation_scale=(1, 1),
+        interpolation_scale_t=1,
+        use_abs_pos=True, 
+    ):
+        super().__init__()
+        # assert num_frames == 1
+        self.use_abs_pos = use_abs_pos
+        self.flatten = flatten
+        self.layer_norm = layer_norm
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias
+        )
+        if layer_norm:
+            self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm = None
+
+        self.patch_size_t = patch_size_t
+        self.patch_size = patch_size
+        # See:
+        # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+        self.height, self.width = height // patch_size, width // patch_size
+        self.base_size = (height // patch_size, width // patch_size)
+        self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+        pos_embed = get_2d_sincos_pos_embed(
+            embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+        )
+        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+        self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.interpolation_scale_t = interpolation_scale_t
+        temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t)
+        self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+        # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+    def forward(self, latent, num_frames):
+        b, _, _, _, _ = latent.shape
+        video_latent, image_latent = None, None
+        # b c 1 h w
+        # assert latent.shape[-3] == 1 and num_frames == 1
+        height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+        latent = rearrange(latent, 'b c t h w -> (b t) c h w')
+        latent = self.proj(latent)
+
+        if self.flatten:
+            latent = latent.flatten(2).transpose(1, 2)  # BT C H W -> BT N C
+        if self.layer_norm:
+            latent = self.norm(latent)
+
+        if self.use_abs_pos:
+            # Interpolate positional embeddings if needed.
+            # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+            if self.height != height or self.width != width:
+                # raise NotImplementedError
+                pos_embed = get_2d_sincos_pos_embed(
+                    embed_dim=self.pos_embed.shape[-1],
+                    grid_size=(height, width),
+                    base_size=self.base_size,
+                    interpolation_scale=self.interpolation_scale,
+                )
+                pos_embed = torch.from_numpy(pos_embed)
+                pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                pos_embed = self.pos_embed
+
+
+            if self.num_frames != num_frames:
+                # import ipdb;ipdb.set_trace()
+                # raise NotImplementedError
+                if get_sequence_parallel_state():
+                    if npu_config is not None:
+                        sp_size = hccl_info.world_size
+                        temp_pos_embed = get_1d_sincos_pos_embed(
+                            embed_dim=self.temp_pos_embed.shape[-1],
+                            grid_size=num_frames * sp_size,
+                            base_size=self.base_size_t,
+                            interpolation_scale=self.interpolation_scale_t,
+                        )
+                        rank = hccl_info.rank % sp_size
+                        st_frame = rank * num_frames
+                        ed_frame = st_frame + num_frames
+                        temp_pos_embed = temp_pos_embed[st_frame: ed_frame]
+                    else:
+                        sp_size = nccl_info.world_size
+                        temp_pos_embed = get_1d_sincos_pos_embed(
+                            embed_dim=self.temp_pos_embed.shape[-1],
+                            grid_size=num_frames * sp_size,
+                            base_size=self.base_size_t,
+                            interpolation_scale=self.interpolation_scale_t,
+                        )
+                        rank = nccl_info.rank % sp_size
+                        st_frame = rank * num_frames
+                        ed_frame = st_frame + num_frames
+                        temp_pos_embed = temp_pos_embed[st_frame: ed_frame]
+
+                else:
+                    temp_pos_embed = get_1d_sincos_pos_embed(
+                        embed_dim=self.temp_pos_embed.shape[-1],
+                        grid_size=num_frames,
+                        base_size=self.base_size_t,
+                        interpolation_scale=self.interpolation_scale_t,
+                    )
+                temp_pos_embed = torch.from_numpy(temp_pos_embed)
+                temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                temp_pos_embed = self.temp_pos_embed
+
+            latent = (latent + pos_embed).to(latent.dtype)
+        
+        latent = rearrange(latent, '(b t) n c -> b t n c', b=b)
+        video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:]
+
+        if self.use_abs_pos:
+            # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+            temp_pos_embed = temp_pos_embed.unsqueeze(2)
+            video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None
+            image_latent = (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None
+
+        video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None
+        image_latent = rearrange(image_latent, 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None
+
+        if num_frames == 1 and image_latent is None and not get_sequence_parallel_state():
+            image_latent = video_latent
+            video_latent = None
+        # print('video_latent is None, image_latent is None', video_latent is None, image_latent is None)
+        return video_latent, image_latent
+    
+
+
+class OverlapPatchEmbed3D(nn.Module):
+    """2D Image to Patch Embedding but with 3D position embedding"""
+
+    def __init__(
+        self,
+        num_frames=1, 
+        height=224,
+        width=224,
+        patch_size_t=1,
+        patch_size=16,
+        in_channels=3,
+        embed_dim=768,
+        layer_norm=False,
+        flatten=True,
+        bias=True,
+        interpolation_scale=(1, 1),
+        interpolation_scale_t=1,
+        use_abs_pos=True, 
+    ):
+        super().__init__()
+        # assert patch_size_t == 1 and patch_size == 1
+        self.use_abs_pos = use_abs_pos
+        self.flatten = flatten
+        self.layer_norm = layer_norm
+
+        self.proj = nn.Conv3d(
+            in_channels, embed_dim, kernel_size=(patch_size_t, patch_size, patch_size), stride=(patch_size_t, patch_size, patch_size), bias=bias
+        )
+        if layer_norm:
+            self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm = None
+
+        self.patch_size_t = patch_size_t
+        self.patch_size = patch_size
+        # See:
+        # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+        self.height, self.width = height // patch_size, width // patch_size
+        self.base_size = (height // patch_size, width // patch_size)
+        self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+        pos_embed = get_2d_sincos_pos_embed(
+            embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+        )
+        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+        self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.interpolation_scale_t = interpolation_scale_t
+        temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t)
+        self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+        # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+    def forward(self, latent, num_frames):
+        b, _, _, _, _ = latent.shape
+        video_latent, image_latent = None, None
+        # b c 1 h w
+        # assert latent.shape[-3] == 1 and num_frames == 1
+        height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+        # latent = rearrange(latent, 'b c t h w -> (b t) c h w')
+        latent = self.proj(latent)
+
+        if self.flatten:
+            # latent = latent.flatten(2).transpose(1, 2)  # BT C H W -> BT N C
+            latent = rearrange(latent, 'b c t h w -> (b t) (h w) c ')
+        if self.layer_norm:
+            latent = self.norm(latent)
+
+        if self.use_abs_pos:
+            # Interpolate positional embeddings if needed.
+            # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+            if self.height != height or self.width != width:
+                # raise NotImplementedError
+                pos_embed = get_2d_sincos_pos_embed(
+                    embed_dim=self.pos_embed.shape[-1],
+                    grid_size=(height, width),
+                    base_size=self.base_size,
+                    interpolation_scale=self.interpolation_scale,
+                )
+                pos_embed = torch.from_numpy(pos_embed)
+                pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                pos_embed = self.pos_embed
+
+
+            if self.num_frames != num_frames:
+                # import ipdb;ipdb.set_trace()
+                # raise NotImplementedError
+                temp_pos_embed = get_1d_sincos_pos_embed(
+                    embed_dim=self.temp_pos_embed.shape[-1],
+                    grid_size=num_frames,
+                    base_size=self.base_size_t,
+                    interpolation_scale=self.interpolation_scale_t,
+                )
+                temp_pos_embed = torch.from_numpy(temp_pos_embed)
+                temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                temp_pos_embed = self.temp_pos_embed
+
+            latent = (latent + pos_embed).to(latent.dtype)
+        
+        latent = rearrange(latent, '(b t) n c -> b t n c', b=b)
+        video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:]
+
+        if self.use_abs_pos:
+            # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+            temp_pos_embed = temp_pos_embed.unsqueeze(2)
+            video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None
+            image_latent = (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None
+
+
+        video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None
+        image_latent = rearrange(image_latent, 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None
+
+        if num_frames == 1 and image_latent is None:
+            image_latent = video_latent
+            video_latent = None
+        return video_latent, image_latent
+    
+
+
+class OverlapPatchEmbed2D(nn.Module):
+    """2D Image to Patch Embedding but with 3D position embedding"""
+
+    def __init__(
+        self,
+        num_frames=1, 
+        height=224,
+        width=224,
+        patch_size_t=1,
+        patch_size=16,
+        in_channels=3,
+        embed_dim=768,
+        layer_norm=False,
+        flatten=True,
+        bias=True,
+        interpolation_scale=(1, 1),
+        interpolation_scale_t=1,
+        use_abs_pos=True, 
+    ):
+        super().__init__()
+        assert patch_size_t == 1
+        self.use_abs_pos = use_abs_pos
+        self.flatten = flatten
+        self.layer_norm = layer_norm
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias
+        )
+        if layer_norm:
+            self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm = None
+
+        self.patch_size_t = patch_size_t
+        self.patch_size = patch_size
+        # See:
+        # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+        self.height, self.width = height // patch_size, width // patch_size
+        self.base_size = (height // patch_size, width // patch_size)
+        self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+        pos_embed = get_2d_sincos_pos_embed(
+            embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+        )
+        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+        self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.interpolation_scale_t = interpolation_scale_t
+        temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t)
+        self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+        # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+    def forward(self, latent, num_frames):
+        b, _, _, _, _ = latent.shape
+        video_latent, image_latent = None, None
+        # b c 1 h w
+        # assert latent.shape[-3] == 1 and num_frames == 1
+        height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+        latent = rearrange(latent, 'b c t h w -> (b t) c h w')
+        latent = self.proj(latent)
+
+        if self.flatten:
+            latent = latent.flatten(2).transpose(1, 2)  # BT C H W -> BT N C
+        if self.layer_norm:
+            latent = self.norm(latent)
+
+        if self.use_abs_pos:
+            # Interpolate positional embeddings if needed.
+            # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+            if self.height != height or self.width != width:
+                # raise NotImplementedError
+                pos_embed = get_2d_sincos_pos_embed(
+                    embed_dim=self.pos_embed.shape[-1],
+                    grid_size=(height, width),
+                    base_size=self.base_size,
+                    interpolation_scale=self.interpolation_scale,
+                )
+                pos_embed = torch.from_numpy(pos_embed)
+                pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                pos_embed = self.pos_embed
+
+
+            if self.num_frames != num_frames:
+                # import ipdb;ipdb.set_trace()
+                # raise NotImplementedError
+                temp_pos_embed = get_1d_sincos_pos_embed(
+                    embed_dim=self.temp_pos_embed.shape[-1],
+                    grid_size=num_frames,
+                    base_size=self.base_size_t,
+                    interpolation_scale=self.interpolation_scale_t,
+                )
+                temp_pos_embed = torch.from_numpy(temp_pos_embed)
+                temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                temp_pos_embed = self.temp_pos_embed
+
+            latent = (latent + pos_embed).to(latent.dtype)
+        
+        latent = rearrange(latent, '(b t) n c -> b t n c', b=b)
+        video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:]
+
+        if self.use_abs_pos:
+            # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+            temp_pos_embed = temp_pos_embed.unsqueeze(2)
+            video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None
+            image_latent = (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None
+
+
+        video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None
+        image_latent = rearrange(image_latent, 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None
+
+        if num_frames == 1 and image_latent is None:
+            image_latent = video_latent
+            video_latent = None
+        return video_latent, image_latent
+    
+class Attention(Attention_):
+    def __init__(self, downsampler, attention_mode, use_rope, interpolation_scale_thw, **kwags):
+        processor = AttnProcessor2_0(attention_mode=attention_mode, use_rope=use_rope, interpolation_scale_thw=interpolation_scale_thw)
+        super().__init__(processor=processor, **kwags)
+        self.downsampler = None
+        if downsampler: # downsampler  k155_s122
+            downsampler_ker_size = list(re.search(r'k(\d{2,3})', downsampler).group(1)) # 122
+            down_factor = list(re.search(r's(\d{2,3})', downsampler).group(1))
+            downsampler_ker_size = [int(i) for i in downsampler_ker_size]
+            downsampler_padding = [(i - 1) // 2 for i in downsampler_ker_size]
+            down_factor = [int(i) for i in down_factor]
+            
+            if len(downsampler_ker_size) == 2:
+                self.downsampler = DownSampler2d(kwags['query_dim'], kwags['query_dim'], kernel_size=downsampler_ker_size, stride=1,
+                                            padding=downsampler_padding, groups=kwags['query_dim'], down_factor=down_factor,
+                                            down_shortcut=True)
+            elif len(downsampler_ker_size) == 3:
+                self.downsampler = DownSampler3d(kwags['query_dim'], kwags['query_dim'], kernel_size=downsampler_ker_size, stride=1,
+                                            padding=downsampler_padding, groups=kwags['query_dim'], down_factor=down_factor,
+                                            down_shortcut=True)
+                
+        # self.q_norm = nn.LayerNorm(kwags['dim_head'], elementwise_affine=True, eps=1e-6)
+        # self.k_norm = nn.LayerNorm(kwags['dim_head'], elementwise_affine=True, eps=1e-6) 
+
+    def prepare_attention_mask(
+        self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3
+    ) -> torch.Tensor:
+        r"""
+        Prepare the attention mask for the attention computation.
+
+        Args:
+            attention_mask (`torch.Tensor`):
+                The attention mask to prepare.
+            target_length (`int`):
+                The target length of the attention mask. This is the length of the attention mask after padding.
+            batch_size (`int`):
+                The batch size, which is used to repeat the attention mask.
+            out_dim (`int`, *optional*, defaults to `3`):
+                The output dimension of the attention mask. Can be either `3` or `4`.
+
+        Returns:
+            `torch.Tensor`: The prepared attention mask.
+        """
+        head_size = self.heads
+        if get_sequence_parallel_state():
+            head_size = head_size // nccl_info.world_size
+        if attention_mask is None:
+            return attention_mask
+
+        current_length: int = attention_mask.shape[-1]
+        if current_length != target_length:
+            if attention_mask.device.type == "mps":
+                # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+                # Instead, we can manually construct the padding tensor.
+                padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)
+                padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)
+                attention_mask = torch.cat([attention_mask, padding], dim=2)
+            else:
+                # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+                #       we want to instead pad by (0, remaining_length), where remaining_length is:
+                #       remaining_length: int = target_length - current_length
+                # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+                attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+        if out_dim == 3:
+            if attention_mask.shape[0] < batch_size * head_size:
+                attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
+        elif out_dim == 4:
+            attention_mask = attention_mask.unsqueeze(1)
+            attention_mask = attention_mask.repeat_interleave(head_size, dim=1)
+
+        return attention_mask
+
+class DownSampler3d(nn.Module):
+    def __init__(self, *args, **kwargs):
+        ''' Required kwargs: down_factor, downsampler'''
+        super().__init__()
+        self.down_factor = kwargs.pop('down_factor')
+        self.down_shortcut = kwargs.pop('down_shortcut')
+        self.layer = nn.Conv3d(*args, **kwargs)
+
+    def forward(self, x, attention_mask, t, h, w):
+        b = x.shape[0]
+        x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w)
+        if npu_config is None:
+            x = self.layer(x) + (x if self.down_shortcut else 0)
+        else:
+            x_dtype = x.dtype
+            x = npu_config.run_conv3d(self.layer, x, x_dtype) + (x if self.down_shortcut else 0)
+        
+        self.t = t//self.down_factor[0]
+        self.h = h//self.down_factor[1]
+        self.w = w//self.down_factor[2]
+        x = rearrange(x, 'b d (t dt) (h dh) (w dw) -> (b dt dh dw) (t h w) d', 
+                      t=t//self.down_factor[0], h=h//self.down_factor[1], w=w//self.down_factor[2], 
+                         dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2])
+    
+
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=t, h=h, w=w)
+        attention_mask = rearrange(attention_mask, 'b 1 (t dt) (h dh) (w dw) -> (b dt dh dw) 1 (t h w)',
+                      t=t//self.down_factor[0], h=h//self.down_factor[1], w=w//self.down_factor[2], 
+                         dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2])
+        return x, attention_mask
+        
+    def reverse(self, x, t, h, w):
+        x = rearrange(x, '(b dt dh dw) (t h w) d -> b (t dt h dh w dw) d', 
+                      t=t, h=h, w=w, 
+                         dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2])
+        return x
+
+
+class DownSampler2d(nn.Module):
+    def __init__(self, *args, **kwargs):
+        ''' Required kwargs: down_factor, downsampler'''
+        super().__init__()
+        self.down_factor = kwargs.pop('down_factor')
+        self.down_shortcut = kwargs.pop('down_shortcut')
+        self.layer = nn.Conv2d(*args, **kwargs)
+
+    def forward(self, x, attention_mask, t, h, w):
+        b = x.shape[0]
+        x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w)
+        x = self.layer(x) + (x if self.down_shortcut else 0)
+
+        self.t = 1
+        self.h = h//self.down_factor[0]
+        self.w = w//self.down_factor[1]
+
+        x = rearrange(x, 'b d (h dh) (w dw) -> (b dh dw) (h w) d', 
+                      h=h//self.down_factor[0], w=w//self.down_factor[1], 
+                      dh=self.down_factor[0], dw=self.down_factor[1])
+    
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> (b t) 1 h w', h=h, w=w)
+        attention_mask = rearrange(attention_mask, 'b 1 (h dh) (w dw) -> (b dh dw) 1 (h w)',
+                      h=h//self.down_factor[0], w=w//self.down_factor[1], 
+                         dh=self.down_factor[0], dw=self.down_factor[1])
+        return x, attention_mask
+    
+    def reverse(self, x, t, h, w):
+        x = rearrange(x, '(b t dh dw) (h w) d -> b (t h dh w dw) d', 
+                      t=t, h=h, w=w, 
+                      dh=self.down_factor[0], dw=self.down_factor[1])
+        return x
+    
+class AttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self, attention_mode='xformers', use_rope=False, interpolation_scale_thw=(1, 1, 1)):
+        self.use_rope = use_rope
+        self.interpolation_scale_thw = interpolation_scale_thw
+        if self.use_rope:
+            self._init_rope(interpolation_scale_thw)
+        self.attention_mode = attention_mode
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+
+    def _init_rope(self, interpolation_scale_thw):
+        self.rope = RoPE3D(interpolation_scale_thw=interpolation_scale_thw)
+        self.position_getter = PositionGetter3D()
+    
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        frame: int = 8, 
+        height: int = 16, 
+        width: int = 16, 
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+
+
+        if attn.downsampler is not None:
+            hidden_states, attention_mask = attn.downsampler(hidden_states, attention_mask, t=frame, h=height, w=width)
+            frame, height, width = attn.downsampler.t, attn.downsampler.h, attn.downsampler.w
+
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+        
+        if get_sequence_parallel_state():
+            if npu_config is not None:
+                sequence_length, batch_size, _ = (
+                    hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+                )
+            else:
+                sequence_length, batch_size, _ = (
+                    hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+                )
+        else:
+            batch_size, sequence_length, _ = (
+                hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+            )
+
+        if attention_mask is not None:
+            if npu_config is None:
+                attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length * nccl_info.world_size, batch_size)
+                # scaled_dot_product_attention expects attention_mask shape to be
+                # (batch, heads, source_length, target_length)
+                if get_sequence_parallel_state():
+                    attention_mask = attention_mask.view(batch_size, attn.heads // nccl_info.world_size, -1, attention_mask.shape[-1])
+                else:
+                    attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+            else:
+                attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1])
+        
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        if npu_config is not None and npu_config.on_npu:
+            if get_sequence_parallel_state():
+                query = query.view(-1, attn.heads, head_dim)  # [s // sp, b, h * d] -> [s // sp * b, h, d]
+                key = key.view(-1, attn.heads, head_dim)
+                value = value.view(-1, attn.heads, head_dim)
+                # query = attn.q_norm(query)
+                # key = attn.k_norm(key)
+                h_size = attn.heads * head_dim
+                sp_size = hccl_info.world_size
+                h_size_sp = h_size // sp_size
+                # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d]
+                query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp)
+                key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp)
+                value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).view(-1, batch_size, h_size_sp)
+                if self.use_rope:
+                    query = query.view(-1, batch_size, attn.heads // sp_size, head_dim)
+                    key = key.view(-1, batch_size, attn.heads // sp_size, head_dim)
+                    # require the shape of (batch_size x nheads x ntokens x dim)
+                    pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device)
+                    query = self.rope(query, pos_thw)
+                    key = self.rope(key, pos_thw)
+                query = query.view(-1, batch_size, h_size_sp)
+                key = key.view(-1, batch_size, h_size_sp)
+                value = value.view(-1, batch_size, h_size_sp)
+                hidden_states = npu_config.run_attention(query, key, value, attention_mask, "SBH",
+                                                         head_dim, attn.heads // sp_size)
+
+                hidden_states = hidden_states.view(-1, attn.heads // sp_size, head_dim)
+
+                # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d]
+                hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).view(-1, batch_size, h_size)
+            else:
+                if npu_config.enable_FA and query.dtype == torch.float32:
+                    dtype = torch.bfloat16
+                else:
+                    dtype = None
+
+                query = query.view(batch_size, -1, attn.heads, head_dim)
+                key = key.view(batch_size, -1, attn.heads, head_dim)
+                # query = attn.q_norm(query)
+                # key = attn.k_norm(key)
+                if self.use_rope:
+                    # require the shape of (batch_size x nheads x ntokens x dim)
+                    pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device)
+                    query = self.rope(query, pos_thw)
+                    key = self.rope(key, pos_thw)
+                query = query.view(batch_size, -1, attn.heads * head_dim)
+                key = key.view(batch_size, -1, attn.heads * head_dim)
+
+                with set_run_dtype(query, dtype):
+                    query, key, value = npu_config.set_current_run_dtype([query, key, value])
+                    hidden_states = npu_config.run_attention(query, key, value, attention_mask, "BSH",
+                                                             head_dim, attn.heads)
+
+                    hidden_states = npu_config.restore_dtype(hidden_states)
+        else:
+            if get_sequence_parallel_state():
+                query = query.reshape(-1, attn.heads, head_dim)  # [s // sp, b, h * d] -> [s // sp * b, h, d]
+                key = key.reshape(-1, attn.heads, head_dim)
+                value = value.reshape(-1, attn.heads, head_dim)
+                # query = attn.q_norm(query)
+                # key = attn.k_norm(key)
+                h_size = attn.heads * head_dim
+                sp_size = nccl_info.world_size
+                h_size_sp = h_size // sp_size
+                # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d]
+                query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp)
+                key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp)
+                value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp)
+                query = query.reshape(-1, batch_size, attn.heads // sp_size, head_dim)
+                key = key.reshape(-1, batch_size, attn.heads // sp_size, head_dim)
+                value = value.reshape(-1, batch_size, attn.heads // sp_size, head_dim)
+                # print('query', query.shape, 'key', key.shape, 'value', value.shape)
+                if self.use_rope:
+                    # require the shape of (batch_size x nheads x ntokens x dim)
+                    pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device)
+                    query = self.rope(query, pos_thw)
+                    key = self.rope(key, pos_thw)
+
+                # print('after rope query', query.shape, 'key', key.shape, 'value', value.shape)
+                query = rearrange(query, 's b h d -> b h s d')
+                key = rearrange(key, 's b h d -> b h s d')
+                value = rearrange(value, 's b h d -> b h s d')
+                # print('rearrange query', query.shape, 'key', key.shape, 'value', value.shape)
+
+                # 0, -10000 ->(bool) False, True ->(any) True ->(not) False
+                # 0, 0 ->(bool) False, False ->(any) False ->(not) True
+                if attention_mask is None or not torch.any(attention_mask.bool()):  # 0 mean visible
+                    attention_mask = None
+                # the output of sdp = (batch, num_heads, seq_len, head_dim)
+                # TODO: add support for attn.scale when we move to Torch 2.1
+                # import ipdb;ipdb.set_trace()
+                # print(attention_mask)
+                if self.attention_mode == 'flash':
+                    assert attention_mask is None, 'flash-attn do not support attention_mask'
+                    with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
+                        hidden_states = F.scaled_dot_product_attention(
+                            query, key, value, dropout_p=0.0, is_causal=False
+                        )
+                elif self.attention_mode == 'xformers':
+                    with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
+                        hidden_states = F.scaled_dot_product_attention(
+                            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+                        )
+                elif self.attention_mode == 'math':
+                    hidden_states = F.scaled_dot_product_attention(
+                        query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+                    )
+                else:
+                    raise NotImplementedError(f'Found attention_mode: {self.attention_mode}')
+                
+                hidden_states = rearrange(hidden_states, 'b h s d -> s b h d')
+
+                hidden_states = hidden_states.reshape(-1, attn.heads // sp_size, head_dim)
+
+                # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d]
+                hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).reshape(-1, batch_size, h_size)
+            else:
+                query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+                key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+                
+                # qk norm
+                # query = attn.q_norm(query)
+                # key = attn.k_norm(key)
+
+                if self.use_rope:
+                    # require the shape of (batch_size x nheads x ntokens x dim)
+                    pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device)
+                    query = self.rope(query, pos_thw)
+                    key = self.rope(key, pos_thw)
+                    
+                value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+                # 0, -10000 ->(bool) False, True ->(any) True ->(not) False
+                # 0, 0 ->(bool) False, False ->(any) False ->(not) True
+                if attention_mask is None or not torch.any(attention_mask.bool()):  # 0 mean visible
+                    attention_mask = None
+                # the output of sdp = (batch, num_heads, seq_len, head_dim)
+                # TODO: add support for attn.scale when we move to Torch 2.1
+                # import ipdb;ipdb.set_trace()
+                # print(attention_mask)
+                if self.attention_mode == 'flash':
+                    assert attention_mask is None, 'flash-attn do not support attention_mask'
+                    with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
+                        hidden_states = F.scaled_dot_product_attention(
+                            query, key, value, dropout_p=0.0, is_causal=False
+                        )
+                elif self.attention_mode == 'xformers':
+                    with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
+                        hidden_states = F.scaled_dot_product_attention(
+                            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+                        )
+                elif self.attention_mode == 'math':
+                    hidden_states = F.scaled_dot_product_attention(
+                        query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+                    )
+                else:
+                    raise NotImplementedError(f'Found attention_mode: {self.attention_mode}')
+                hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        if attn.downsampler is not None:
+            hidden_states = attn.downsampler.reverse(hidden_states, t=frame, h=height, w=width)
+        return hidden_states
+
+
+
+class FeedForward_Conv3d(nn.Module):
+    def __init__(self, downsampler, dim, hidden_features, bias=True):
+        super(FeedForward_Conv3d, self).__init__()
+        
+        self.bias = bias
+
+        self.project_in = nn.Linear(dim, hidden_features, bias=bias)
+
+        self.dwconv = nn.ModuleList([
+            nn.Conv3d(hidden_features, hidden_features, kernel_size=(5, 5, 5), stride=1, padding=(2, 2, 2), dilation=1,
+                        groups=hidden_features, bias=bias),
+            nn.Conv3d(hidden_features, hidden_features, kernel_size=(3, 3, 3), stride=1, padding=(1, 1, 1), dilation=1,
+                        groups=hidden_features, bias=bias),
+            nn.Conv3d(hidden_features, hidden_features, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0), dilation=1,
+                        groups=hidden_features, bias=bias)
+        ])
+
+        self.project_out = nn.Linear(hidden_features, dim, bias=bias)
+
+
+    def forward(self, x, t, h, w):
+        # import ipdb;ipdb.set_trace()
+        if npu_config is None:
+            x = self.project_in(x)
+            x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w)
+            x = F.gelu(x)
+            out = x
+            for module in self.dwconv:
+                out = out + module(x)
+            out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w)
+            x = self.project_out(out)
+        else:
+            x_dtype = x.dtype
+            x = npu_config.run_conv3d(self.project_in, x, npu_config.replaced_type)
+            x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w)
+            x = F.gelu(x)
+            out = x
+            for module in self.dwconv:
+                out = out + npu_config.run_conv3d(module, x, npu_config.replaced_type)
+            out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w)
+            x = npu_config.run_conv3d(self.project_out, out, x_dtype)
+        return x
+
+
+class FeedForward_Conv2d(nn.Module):
+    def __init__(self, downsampler, dim, hidden_features, bias=True):
+        super(FeedForward_Conv2d, self).__init__()
+        
+        self.bias = bias
+
+        self.project_in = nn.Linear(dim, hidden_features, bias=bias)
+
+        self.dwconv = nn.ModuleList([
+            nn.Conv2d(hidden_features, hidden_features, kernel_size=(5, 5), stride=1, padding=(2, 2), dilation=1,
+                        groups=hidden_features, bias=bias),
+            nn.Conv2d(hidden_features, hidden_features, kernel_size=(3, 3), stride=1, padding=(1, 1), dilation=1,
+                        groups=hidden_features, bias=bias),
+            nn.Conv2d(hidden_features, hidden_features, kernel_size=(1, 1), stride=1, padding=(0, 0), dilation=1,
+                        groups=hidden_features, bias=bias)
+        ])
+
+        self.project_out = nn.Linear(hidden_features, dim, bias=bias)
+
+
+    def forward(self, x, t, h, w):
+        # import ipdb;ipdb.set_trace()
+        x = self.project_in(x)
+        x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w)
+        x = F.gelu(x)
+        out = x
+        for module in self.dwconv:
+            out = out + module(x)
+        out = rearrange(out, '(b t) d h w -> b (t h w) d', t=t, h=h, w=w)
+        x = self.project_out(out)
+        return x
+
+@maybe_allow_in_graph
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        num_embeds_ada_norm (:
+            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        upcast_attention (`bool`, *optional*):
+            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+        final_dropout (`bool` *optional*, defaults to False):
+            Whether to apply a final dropout after the last feed-forward layer.
+        attention_type (`str`, *optional*, defaults to `"default"`):
+            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
+        positional_embeddings (`str`, *optional*, defaults to `None`):
+            The type of positional embeddings to apply to.
+        num_positional_embeddings (`int`, *optional*, defaults to `None`):
+            The maximum number of positional embeddings to apply.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+        norm_eps: float = 1e-5,
+        final_dropout: bool = False,
+        attention_type: str = "default",
+        positional_embeddings: Optional[str] = None,
+        num_positional_embeddings: Optional[int] = None,
+        ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
+        ada_norm_bias: Optional[int] = None,
+        ff_inner_dim: Optional[int] = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+        attention_mode: str = "xformers", 
+        downsampler: str = None, 
+        use_rope: bool = False, 
+        interpolation_scale_thw: Tuple[int] = (1, 1, 1), 
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+        self.downsampler = downsampler
+
+        # We keep these boolean flags for backward-compatibility.
+        self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
+        self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
+        self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
+        self.use_layer_norm = norm_type == "layer_norm"
+        self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous"
+
+        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
+                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
+            )
+
+        self.norm_type = norm_type
+        self.num_embeds_ada_norm = num_embeds_ada_norm
+
+        if positional_embeddings and (num_positional_embeddings is None):
+            raise ValueError(
+                "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
+            )
+
+        if positional_embeddings == "sinusoidal":
+            self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
+        else:
+            self.pos_embed = None
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        if norm_type == "ada_norm":
+            self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
+        elif norm_type == "ada_norm_zero":
+            self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        elif norm_type == "ada_norm_continuous":
+            self.norm1 = AdaLayerNormContinuous(
+                dim,
+                ada_norm_continous_conditioning_embedding_dim,
+                norm_elementwise_affine,
+                norm_eps,
+                ada_norm_bias,
+                "rms_norm",
+            )
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+            out_bias=attention_out_bias,
+            attention_mode=attention_mode, 
+            downsampler=downsampler, 
+            use_rope=use_rope, 
+            interpolation_scale_thw=interpolation_scale_thw, 
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            if norm_type == "ada_norm":
+                self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)
+            elif norm_type == "ada_norm_continuous":
+                self.norm2 = AdaLayerNormContinuous(
+                    dim,
+                    ada_norm_continous_conditioning_embedding_dim,
+                    norm_elementwise_affine,
+                    norm_eps,
+                    ada_norm_bias,
+                    "rms_norm",
+                )
+            else:
+                self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim if not double_self_attention else None,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+                out_bias=attention_out_bias,
+                attention_mode=attention_mode, 
+                downsampler=False, 
+                use_rope=False, 
+                interpolation_scale_thw=interpolation_scale_thw, 
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        if norm_type == "ada_norm_continuous":
+            self.norm3 = AdaLayerNormContinuous(
+                dim,
+                ada_norm_continous_conditioning_embedding_dim,
+                norm_elementwise_affine,
+                norm_eps,
+                ada_norm_bias,
+                "layer_norm",
+            )
+
+        elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm", "ada_norm_continuous"]:
+            self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        elif norm_type == "layer_norm_i2vgen":
+            self.norm3 = None
+
+        if downsampler:
+            downsampler_ker_size = list(re.search(r'k(\d{2,3})', downsampler).group(1)) # 122
+            # if len(downsampler_ker_size) == 3:
+            #     self.ff = FeedForward_Conv3d(
+            #         downsampler, 
+            #         dim,
+            #         2 * dim,
+            #         bias=ff_bias,
+            #     )
+            # elif len(downsampler_ker_size) == 2:
+            self.ff = FeedForward_Conv2d(
+                downsampler, 
+                dim,
+                2 * dim,
+                bias=ff_bias,
+            )
+        else:
+            self.ff = FeedForward(
+                dim,
+                dropout=dropout,
+                activation_fn=activation_fn,
+                final_dropout=final_dropout,
+                inner_dim=ff_inner_dim,
+                bias=ff_bias,
+            )
+
+        # 4. Fuser
+        if attention_type == "gated" or attention_type == "gated-text-image":
+            self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim)
+
+        # 5. Scale-shift for PixArt-Alpha.
+        if norm_type == "ada_norm_single":
+            self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        frame: int = None, 
+        height: int = None, 
+        width: int = None, 
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> torch.FloatTensor:
+        if cross_attention_kwargs is not None:
+            if cross_attention_kwargs.get("scale", None) is not None:
+                logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 0. Self-Attention
+        batch_size = hidden_states.shape[0]
+
+        # import ipdb;ipdb.set_trace()
+        if self.norm_type == "ada_norm":
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.norm_type == "ada_norm_zero":
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]:
+            norm_hidden_states = self.norm1(hidden_states)
+        elif self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif self.norm_type == "ada_norm_single":
+            # import ipdb;ipdb.set_trace()
+            if get_sequence_parallel_state():
+                batch_size = hidden_states.shape[1]
+                # print('hidden_states', hidden_states.shape)
+                # print('timestep', timestep.shape)
+                shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                        self.scale_shift_table[:, None] + timestep.reshape(6, batch_size, -1)
+                ).chunk(6, dim=0)
+            else:
+                shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                        self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+                ).chunk(6, dim=1)
+            norm_hidden_states = self.norm1(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+            # norm_hidden_states = norm_hidden_states.squeeze(1)
+        else:
+            raise ValueError("Incorrect norm used")
+
+        if self.pos_embed is not None:
+            norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+        # 1. Prepare GLIGEN inputs
+        cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
+        gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask, frame=frame, height=height, width=width, 
+            **cross_attention_kwargs,
+        )
+        if self.norm_type == "ada_norm_zero":
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        elif self.norm_type == "ada_norm_single":
+            attn_output = gate_msa * attn_output
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        # 1.2 GLIGEN Control
+        if gligen_kwargs is not None:
+            hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
+
+        # 3. Cross-Attention
+        if self.attn2 is not None:
+            if self.norm_type == "ada_norm":
+                norm_hidden_states = self.norm2(hidden_states, timestep)
+            elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]:
+                norm_hidden_states = self.norm2(hidden_states)
+            elif self.norm_type == "ada_norm_single":
+                # For PixArt norm2 isn't applied here:
+                # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
+                norm_hidden_states = hidden_states
+            elif self.norm_type == "ada_norm_continuous":
+                norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"])
+            else:
+                raise ValueError("Incorrect norm")
+
+            if self.pos_embed is not None and self.norm_type != "ada_norm_single":
+                norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                **cross_attention_kwargs,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 4. Feed-forward
+        # i2vgen doesn't have this norm 🤷‍♂️
+        if self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif not self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm3(hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm2(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+        # if self._chunk_size is not None:
+        #     # "feed_forward_chunk_size" can be used to save memory
+        #     ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        # else:
+
+        if self.downsampler:
+            ff_output = self.ff(norm_hidden_states, t=frame, h=height, w=width)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+        elif self.norm_type == "ada_norm_single":
+            ff_output = gate_mlp * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
diff --git a/opensora/models/diffusion/opensora/rope.py b/opensora/models/diffusion/opensora/rope.py
new file mode 100644
index 000000000..8ead2b4b4
--- /dev/null
+++ b/opensora/models/diffusion/opensora/rope.py
@@ -0,0 +1,98 @@
+import torch
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config, set_run_dtype
+    from opensora.acceleration.parallel_states import get_sequence_parallel_state
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import get_sequence_parallel_state
+
+class PositionGetter3D(object):
+    """ return positions of patches """
+
+    def __init__(self, ):
+        self.cache_positions = {}
+        
+    def __call__(self, b, t, h, w, device):
+        if not (b, t,h,w) in self.cache_positions:
+            x = torch.arange(w, device=device)
+            y = torch.arange(h, device=device)
+            z = torch.arange(t, device=device)
+            pos = torch.cartesian_prod(z, y, x)
+            if get_sequence_parallel_state():
+                # print('PositionGetter3D', PositionGetter3D)
+                pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, -1, 1).contiguous().expand(3, -1, b).clone()
+            else:
+                pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, 1, -1).contiguous().expand(3, b, -1).clone()
+            poses = (pos[0].contiguous(), pos[1].contiguous(), pos[2].contiguous())
+            max_poses = (int(poses[0].max()), int(poses[1].max()), int(poses[2].max()))
+
+            self.cache_positions[b, t, h, w] = (poses, max_poses)
+        pos = self.cache_positions[b, t, h, w]
+
+        return pos
+    
+
+class RoPE3D(torch.nn.Module):
+
+    def __init__(self, freq=10000.0, F0=1.0, interpolation_scale_thw=(1, 1, 1)):
+        super().__init__()
+        self.base = freq
+        self.F0 = F0
+        self.interpolation_scale_t = interpolation_scale_thw[0]
+        self.interpolation_scale_h = interpolation_scale_thw[1]
+        self.interpolation_scale_w = interpolation_scale_thw[2]
+        self.cache = {}
+
+    def get_cos_sin(self, D, seq_len, device, dtype, interpolation_scale=1):
+        if (D, seq_len, device, dtype) not in self.cache:
+            inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
+            t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) / interpolation_scale
+            freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
+            freqs = torch.cat((freqs, freqs), dim=-1)
+            cos = freqs.cos()  # (Seq, Dim)
+            sin = freqs.sin()
+            self.cache[D, seq_len, device, dtype] = (cos, sin)
+        return self.cache[D, seq_len, device, dtype]
+
+    @staticmethod
+    def rotate_half(x):
+        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def apply_rope1d(self, tokens, pos1d, cos, sin):
+        assert pos1d.ndim == 2
+        if npu_config is None and not get_sequence_parallel_state():
+            # for (batch_size x nheads x ntokens x dim)
+            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
+            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
+        else:
+            # for (batch_size x ntokens x nheads x dim)
+            cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :]
+            sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :]
+
+        return (tokens * cos) + (self.rotate_half(tokens) * sin)
+
+    def forward(self, tokens, positions):
+        """
+        input:
+            * tokens: batch_size x nheads x ntokens x dim
+            * positions: batch_size x ntokens x 3 (t, y and x position of each token)
+        output:
+            * tokens after appplying RoPE3D (batch_size x nheads x ntokens x x dim)
+        """
+        assert tokens.size(3) % 3 == 0, "number of dimensions should be a multiple of three"
+        D = tokens.size(3) // 3
+        poses, max_poses = positions
+        assert len(poses) == 3 and poses[0].ndim == 2# Batch, Seq, 3
+        cos_t, sin_t = self.get_cos_sin(D, max_poses[0] + 1, tokens.device, tokens.dtype, self.interpolation_scale_t)
+        cos_y, sin_y = self.get_cos_sin(D, max_poses[1] + 1, tokens.device, tokens.dtype, self.interpolation_scale_h)
+        cos_x, sin_x = self.get_cos_sin(D, max_poses[2] + 1, tokens.device, tokens.dtype, self.interpolation_scale_w)
+        # split features into three along the feature dimension, and apply rope1d on each half
+        t, y, x = tokens.chunk(3, dim=-1)
+        t = self.apply_rope1d(t, poses[0], cos_t, sin_t)
+        y = self.apply_rope1d(y, poses[1], cos_y, sin_y)
+        x = self.apply_rope1d(x, poses[2], cos_x, sin_x)
+        tokens = torch.cat((t, y, x), dim=-1)
+        return tokens
\ No newline at end of file
diff --git a/opensora/models/diffusion/udit/modeling_udit.py b/opensora/models/diffusion/udit/modeling_udit.py
new file mode 100644
index 000000000..2d5b87984
--- /dev/null
+++ b/opensora/models/diffusion/udit/modeling_udit.py
@@ -0,0 +1,1073 @@
+import os
+import numpy as np
+from torch import nn
+import torch
+from einops import rearrange, repeat
+from typing import Any, Dict, Optional, Tuple
+from torch.nn import functional as F
+from diffusers.loaders import PeftAdapterMixin
+from diffusers.models.transformer_2d import Transformer2DModelOutput
+from diffusers.utils import is_torch_version, deprecate
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormSingle
+from diffusers.models.embeddings import PixArtAlphaTextProjection
+from opensora.models.diffusion.udit.modules import Upsample2d, Downsample2d, PatchEmbed2D, BasicTransformerBlock, \
+    FP32_GELU, FP32_SiLU, FP32_Layernorm
+from opensora.utils.utils import to_2tuple
+import math
+import re
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+
+class UDiTT2V(ModelMixin, ConfigMixin, PeftAdapterMixin):
+    """
+    A 2D Transformer model for image-like data.
+
+    Parameters:
+        num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
+        in_channels (`int`, *optional*):
+            The number of channels in the input and output (specify if the input is **continuous**).
+        num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
+            This is fixed during training since it is used to learn a number of position embeddings.
+        num_vector_embeds (`int`, *optional*):
+            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
+            Includes the class for the masked latent pixel.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
+        num_embeds_ada_norm ( `int`, *optional*):
+            The number of diffusion steps used during training. Pass if at least one of the norm_layers is
+            `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
+            added to the hidden states.
+
+            During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
+        attention_bias (`bool`, *optional*):
+            Configure if the `TransformerBlocks` attention should contain a bias parameter.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        cross_attention_dim: Optional[int] = None,
+        attention_bias: bool = False,
+        sample_size: Optional[int] = None,
+        sample_size_t: Optional[int] = None,
+        patch_size: Optional[int] = None,
+        patch_size_t: Optional[int] = None,
+        num_vector_embeds: Optional[int] = None,
+        mlp_ratio: int = 4, 
+        depth: Optional[list] = [2, 5, 8, 5, 2],
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        attention_type: str = "default",
+        caption_channels: int = None,
+        interpolation_scale_h: float = None,
+        interpolation_scale_w: float = None,
+        interpolation_scale_t: float = None,
+        use_additional_conditions: Optional[bool] = None,
+        attention_mode: str = 'xformers', 
+        downsampler: str = 'k333_s222', 
+        use_rope: bool = False, 
+        use_stable_fp32: bool = False, 
+    ):
+        super().__init__()
+
+        # Set some common variables used across the board.
+        self.use_stable_fp32 = use_stable_fp32
+        self.use_rope = use_rope
+        self.downsampler = downsampler
+        self.use_linear_projection = use_linear_projection
+        self.interpolation_scale_t = interpolation_scale_t
+        self.interpolation_scale_h = interpolation_scale_h
+        self.interpolation_scale_w = interpolation_scale_w
+        self.caption_channels = caption_channels
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.gradient_checkpointing = False
+        use_additional_conditions = False
+        # if use_additional_conditions is None:
+            # if norm_type == "ada_norm_single" and sample_size == 128:
+            #     use_additional_conditions = True
+            # else:
+            # use_additional_conditions = False
+        self.use_additional_conditions = use_additional_conditions
+
+        # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
+        # Define whether input is continuous or discrete depending on configuration
+        assert in_channels is not None
+
+        if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
+            deprecation_message = (
+                f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
+                " incorrectly set to `'layer_norm'`. Make sure to set `norm_type` to `'ada_norm'` in the config."
+                " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
+                " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
+                " would be very nice if you could open a Pull request for the `transformer/config.json` file"
+            )
+            deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
+            norm_type = "ada_norm"
+
+        # 2. Initialize the right blocks.
+        # Initialize the output blocks and other projection blocks when necessary.
+        self._init_patched_inputs(norm_type=norm_type)
+        if self.use_stable_fp32:
+            self._replace_fp32_layers()
+
+    def _init_patched_inputs(self, norm_type):
+        assert self.config.sample_size_t is not None, "OpenSoraT2V over patched input must provide sample_size_t"
+        assert self.config.sample_size is not None, "OpenSoraT2V over patched input must provide sample_size"
+
+        self.config.sample_size = to_2tuple(self.config.sample_size)
+        # self.num_frames = self.config.sample_size_t
+        # self.height = self.config.sample_size[0]
+        # self.width = self.config.sample_size[1]
+        interpolation_scale_t = ((self.config.sample_size_t - 1) // 16 + 1) if self.config.sample_size_t % 2 == 1 else self.config.sample_size_t / 16
+        interpolation_scale_t = (
+            self.config.interpolation_scale_t if self.config.interpolation_scale_t is not None else interpolation_scale_t
+        )
+        interpolation_scale = (
+            self.config.interpolation_scale_h if self.config.interpolation_scale_h is not None else self.config.sample_size[0] / 30, 
+            self.config.interpolation_scale_w if self.config.interpolation_scale_w is not None else self.config.sample_size[1] / 40, 
+        )
+
+        # down_factor = list(re.search(r's(\d{2,3})', self.downsampler).group(1))
+        # down_factor = [int(i) for i in down_factor]
+        # down_factor = down_factor if isinstance(self.config.down_factor, list) else [self.config.down_factor] * 5
+        # down_factor = [2] * len(self.config.depth)
+        is_video_model = False
+        # if self.config.downsampler is not None and len(self.config.downsampler) == 9:
+        #     is_video_model = True  # to init weight from image
+            # self.pos_embed = OverlapPatchEmbed3D(
+            #     num_frames=self.config.sample_size_t,
+            #     height=self.config.sample_size[0],
+            #     width=self.config.sample_size[1],
+            #     patch_size_t=self.config.patch_size_t,
+            #     patch_size=self.config.patch_size,
+            #     in_channels=self.in_channels,
+            #     embed_dim=self.inner_dim,
+            #     interpolation_scale=interpolation_scale, 
+            #     interpolation_scale_t=interpolation_scale_t,
+            #     use_abs_pos=not self.config.use_rope, 
+            # )
+        # elif self.config.downsampler is not None and len(self.config.downsampler) == 7:
+        #     is_video_model = False
+        self.pos_embed = PatchEmbed2D(
+            num_frames=self.config.sample_size_t,
+            height=self.config.sample_size[0],
+            width=self.config.sample_size[1],
+            patch_size_t=self.config.patch_size_t,
+            patch_size=self.config.patch_size,
+            in_channels=self.in_channels,
+            embed_dim=self.inner_dim,
+            interpolation_scale=interpolation_scale, 
+            interpolation_scale_t=interpolation_scale_t,
+            use_abs_pos=not self.config.use_rope, 
+        )
+        # layer_thw = [[self.config.sample_size_t//self.config.patch_size_t, 
+        #               (self.config.sample_size[0] + self.config.sample_size[0] % (self.config.patch_size*2))//self.config.patch_size, 
+        #               (self.config.sample_size[1] + self.config.sample_size[1] % (self.config.patch_size*2))//self.config.patch_size]]
+        interpolation_scale_thw = (interpolation_scale_t, *interpolation_scale)
+        # for i in range((len(self.config.depth)-1)//2):
+        #     t = layer_thw[i][0] // 2 if layer_thw[i][0] != 1 else 1
+        #     h = (layer_thw[i][1] + layer_thw[i][1] % 4) // 2  # why mod 4, because downsample and downsampler in attention
+        #     w = (layer_thw[i][2] + layer_thw[i][2] % 4) // 2
+        #     layer_thw.append([t, h, w])
+        # self.layer_thw = layer_thw
+        self.encoder_level_1 = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim,
+                    self.config.num_attention_heads,
+                    self.config.attention_head_dim,
+                    downsampler=self.config.downsampler, 
+                    mlp_ratio=self.config.mlp_ratio, 
+                    dropout=self.config.dropout,
+                    cross_attention_dim=self.inner_dim,
+                    activation_fn=self.config.activation_fn,
+                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+                    attention_bias=self.config.attention_bias,
+                    only_cross_attention=self.config.only_cross_attention,
+                    double_self_attention=self.config.double_self_attention,
+                    upcast_attention=self.config.upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=self.config.norm_elementwise_affine,
+                    norm_eps=self.config.norm_eps,
+                    attention_type=self.config.attention_type,
+                    attention_mode=self.config.attention_mode, 
+                    use_rope=self.config.use_rope, 
+                    interpolation_scale_thw=interpolation_scale_thw, 
+                )
+                for _ in range(self.config.depth[0])
+            ]
+        )
+        # self.down1_2 = Downsample3d(self.inner_dim) if is_video_model else Downsample2d(self.inner_dim)
+        self.down1_2 = Downsample2d(self.inner_dim)
+        
+        self.encoder_level_2 = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim * 2,
+                    self.config.num_attention_heads,
+                    self.config.attention_head_dim * 2,
+                    downsampler=self.config.downsampler, 
+                    mlp_ratio=self.config.mlp_ratio, 
+                    dropout=self.config.dropout,
+                    cross_attention_dim=self.inner_dim * 2,
+                    activation_fn=self.config.activation_fn,
+                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+                    attention_bias=self.config.attention_bias,
+                    only_cross_attention=self.config.only_cross_attention,
+                    double_self_attention=self.config.double_self_attention,
+                    upcast_attention=self.config.upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=self.config.norm_elementwise_affine,
+                    norm_eps=self.config.norm_eps,
+                    attention_type=self.config.attention_type,
+                    attention_mode=self.config.attention_mode, 
+                    use_rope=self.config.use_rope, 
+                    interpolation_scale_thw=interpolation_scale_thw, 
+                )
+                for _ in range(self.config.depth[1])
+            ]
+        )
+        # self.down2_3 = Downsample3d(self.inner_dim * 2) if is_video_model else Downsample2d(self.inner_dim * 2)
+        self.down2_3 = Downsample2d(self.inner_dim * 2)
+
+        self.latent = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim * 4,
+                    self.config.num_attention_heads,
+                    self.config.attention_head_dim * 4,
+                    downsampler=self.config.downsampler, 
+                    mlp_ratio=self.config.mlp_ratio, 
+                    dropout=self.config.dropout,
+                    cross_attention_dim=self.inner_dim * 4,
+                    activation_fn=self.config.activation_fn,
+                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+                    attention_bias=self.config.attention_bias,
+                    only_cross_attention=self.config.only_cross_attention,
+                    double_self_attention=self.config.double_self_attention,
+                    upcast_attention=self.config.upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=self.config.norm_elementwise_affine,
+                    norm_eps=self.config.norm_eps,
+                    attention_type=self.config.attention_type,
+                    attention_mode=self.config.attention_mode, 
+                    use_rope=self.config.use_rope, 
+                    interpolation_scale_thw=interpolation_scale_thw, 
+                )
+                for _ in range(self.config.depth[2])
+            ]
+        )
+
+        # self.up3_2 = Upsample3d(int(self.inner_dim * 4)) if is_video_model else Upsample2d(self.inner_dim * 4)  ## From Level 4 to Level 3
+        self.up3_2 = Upsample2d(self.inner_dim * 4)  ## From Level 4 to Level 3
+        
+        self.reduce_chan_level2_norm = nn.LayerNorm(int(self.inner_dim * 4), elementwise_affine=True, eps=1e-6)
+        self.reduce_chan_level2 = nn.Linear(int(self.inner_dim * 4), int(self.inner_dim * 2), bias=True)
+        self.decoder_level_2 = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim * 2,
+                    self.config.num_attention_heads,
+                    self.config.attention_head_dim * 2,
+                    downsampler=self.config.downsampler, 
+                    mlp_ratio=self.config.mlp_ratio, 
+                    dropout=self.config.dropout,
+                    cross_attention_dim=self.inner_dim * 2,
+                    activation_fn=self.config.activation_fn,
+                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+                    attention_bias=self.config.attention_bias,
+                    only_cross_attention=self.config.only_cross_attention,
+                    double_self_attention=self.config.double_self_attention,
+                    upcast_attention=self.config.upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=self.config.norm_elementwise_affine,
+                    norm_eps=self.config.norm_eps,
+                    attention_type=self.config.attention_type,
+                    attention_mode=self.config.attention_mode, 
+                    use_rope=self.config.use_rope, 
+                    interpolation_scale_thw=interpolation_scale_thw, 
+                )
+                for _ in range(self.config.depth[3])
+            ]
+        )
+
+        # self.up2_1 = Upsample3d(int(self.inner_dim * 2)) if is_video_model else Upsample2d(self.inner_dim * 2)  ## From Level 4 to Level 3
+        self.up2_1 = Upsample2d(self.inner_dim * 2)  ## From Level 4 to Level 3
+        
+        self.reduce_chan_level1_norm = nn.LayerNorm(int(self.inner_dim * 2), elementwise_affine=True, eps=1e-6)
+        self.reduce_chan_level1 = nn.Linear(int(self.inner_dim * 2), int(self.inner_dim * 1), bias=True)
+        self.decoder_level_1 = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    self.inner_dim,
+                    self.config.num_attention_heads,
+                    self.config.attention_head_dim,
+                    downsampler=self.config.downsampler, 
+                    mlp_ratio=self.config.mlp_ratio, 
+                    dropout=self.config.dropout,
+                    cross_attention_dim=self.inner_dim,
+                    activation_fn=self.config.activation_fn,
+                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
+                    attention_bias=self.config.attention_bias,
+                    only_cross_attention=self.config.only_cross_attention,
+                    double_self_attention=self.config.double_self_attention,
+                    upcast_attention=self.config.upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=self.config.norm_elementwise_affine,
+                    norm_eps=self.config.norm_eps,
+                    attention_type=self.config.attention_type,
+                    attention_mode=self.config.attention_mode, 
+                    use_rope=self.config.use_rope, 
+                    interpolation_scale_thw=interpolation_scale_thw, 
+                )
+                for _ in range(self.config.depth[4])
+            ]
+        )
+
+        if self.config.norm_type != "ada_norm_single":
+            self.norm_out = nn.LayerNorm(2 * self.inner_dim, elementwise_affine=False, eps=1e-6)
+            self.proj_out_1 = nn.Linear(2 * self.inner_dim, 2 * self.inner_dim)
+            self.proj_out_2 = nn.Linear(
+                2 * self.inner_dim, self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels
+            )
+        elif self.config.norm_type == "ada_norm_single":
+            self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
+            self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / (self.inner_dim)**0.5)
+            self.proj_out = nn.Linear(
+                self.inner_dim, self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels
+            )
+
+        # PixArt-Alpha blocks.
+        # self.adaln_single = None
+        # if self.config.norm_type == "ada_norm_single":
+        # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use
+        # additional conditions until we find better name
+        self.adaln_single_1 = AdaLayerNormSingle(
+            self.inner_dim, use_additional_conditions=self.use_additional_conditions
+        )
+        self.adaln_single_2 = AdaLayerNormSingle(
+            self.inner_dim * 2, use_additional_conditions=self.use_additional_conditions
+        )
+        self.adaln_single_3 = AdaLayerNormSingle(
+            self.inner_dim * 4, use_additional_conditions=self.use_additional_conditions
+        )
+
+        # self.caption_projection = None
+        # if self.caption_channels is not None:
+        self.caption_projection_1 = PixArtAlphaTextProjection(
+            in_features=self.caption_channels, hidden_size=self.inner_dim
+        )
+        self.caption_projection_2 = PixArtAlphaTextProjection(
+            in_features=self.caption_channels, hidden_size=self.inner_dim * 2
+        )
+        self.caption_projection_3 = PixArtAlphaTextProjection(
+            in_features=self.caption_channels, hidden_size=self.inner_dim * 4
+        )
+
+    def _replace_fp32_layers(self, module=None):
+        if module is None:
+            module = self
+        for name, submodule in module.named_children():
+            if isinstance(submodule, nn.LayerNorm):
+                # print(f"Replacing LayerNorm in {name}")
+                new_layer = FP32_Layernorm(submodule.normalized_shape, submodule.eps, submodule.elementwise_affine)
+                if submodule.elementwise_affine:
+                    new_layer.weight.data.copy_(submodule.weight.data.float())
+                    if submodule.bias is not None:
+                        new_layer.bias.data.copy_(submodule.bias.data.float()) 
+                setattr(module, name, new_layer)
+            elif isinstance(submodule, nn.SiLU):
+                # print(f"Replacing SiLU in {name}")
+                setattr(module, name, FP32_SiLU(submodule.inplace))
+            elif isinstance(submodule, nn.GELU):
+                # print(f"Replacing GELU in {name}")
+                setattr(module, name, FP32_GELU(submodule.approximate))
+            else:
+                self._replace_fp32_layers(submodule) 
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: Optional[torch.LongTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        added_cond_kwargs: Dict[str, torch.Tensor] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        use_image_num: Optional[int] = 0,
+        return_dict: bool = True,
+    ):
+        """
+        The [`Transformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
+                Input `hidden_states`.
+            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+                self-attention.
+            timestep ( `torch.LongTensor`, *optional*):
+                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
+            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
+                Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
+                `AdaLayerZeroNorm`.
+            cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            attention_mask ( `torch.Tensor`, *optional*):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            encoder_attention_mask ( `torch.Tensor`, *optional*):
+                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+                    * Mask `(batch, sequence_length)` True = keep, False = discard.
+                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+                above. This bias will be added to the cross-attention scores.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unets.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        batch_size, c, frame, height, width = hidden_states.shape
+        frame = frame - use_image_num  # 21-4=17
+        if cross_attention_kwargs is not None:
+            if cross_attention_kwargs.get("scale", None) is not None:
+                print.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        
+        if attention_mask is not None and attention_mask.ndim == 4:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #   (keep = +0,     discard = -10000.0)
+            # b, frame+use_image_num, h, w -> a video with images
+            # b, 1, h, w -> only images
+            pad_h_0, pad_w_0 = height % (self.config.patch_size * 2), width % (self.config.patch_size * 2)
+
+            hidden_states = F.pad(hidden_states, (0, pad_w_0, 0, pad_h_0, 0, 0), mode='reflect')
+            attention_mask = attention_mask.to(self.dtype)
+            attention_mask = attention_mask.unsqueeze(1)  # b 1 t h w
+            attention_mask = F.pad(attention_mask, (0, pad_w_0, 0, pad_h_0, 0, 0))
+            attention_mask = F.max_pool3d(attention_mask, kernel_size=(self.config.patch_size_t, self.config.patch_size, self.config.patch_size), 
+                                                  stride=(self.config.patch_size_t, self.config.patch_size, self.config.patch_size))
+            attention_mask = rearrange(attention_mask, 'b 1 t h w -> b 1 (t h w)') 
+
+            attention_bias = (1 - attention_mask.bool().to(self.dtype)) * -10000.0
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 3:  
+            # b, 1, l -> only video
+            encoder_attention_mask = (1 - encoder_attention_mask.to(self.dtype)) * -10000.0
+
+
+        # 1. Input
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+        hidden_states, encoder_hidden_states_1, encoder_hidden_states_2, encoder_hidden_states_3, \
+            timestep_1, timestep_2, timestep_3, \
+                embedded_timestep_1, embedded_timestep_2, embedded_timestep_3 = self._operate_on_patched_inputs(
+            hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num
+        )
+        frame, height, width = frame // self.config.patch_size_t, \
+            (height + pad_h_0) // (self.config.patch_size), (width + pad_w_0) // (self.config.patch_size)
+
+
+        assert not torch.any(torch.isnan(hidden_states)), 'after _operate_on_patched_inputs'
+        def create_custom_forward(module, return_dict=None):
+            def custom_forward(*inputs):
+                if return_dict is not None:
+                    return module(*inputs, return_dict=return_dict)
+                else:
+                    return module(*inputs)
+
+            return custom_forward
+        
+        # encoder_1
+        out_enc_level1 = hidden_states
+        # import ipdb;ipdb.set_trace()
+        if self.training and self.gradient_checkpointing:
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            for block in self.encoder_level_1:
+                out_enc_level1 = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    out_enc_level1,
+                    attention_bias,
+                    encoder_hidden_states_1,
+                    encoder_attention_mask,
+                    timestep_1,
+                    cross_attention_kwargs,
+                    class_labels,
+                    frame, 
+                    height, 
+                    width, 
+                    **ckpt_kwargs,
+                )
+        else:
+            for block in self.encoder_level_1:
+                out_enc_level1 = block(
+                    out_enc_level1,
+                    attention_mask=attention_bias,
+                    encoder_hidden_states=encoder_hidden_states_1,
+                    encoder_attention_mask=encoder_attention_mask,
+                    timestep=timestep_1,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    frame=frame, 
+                    height=height, 
+                    width=width, 
+                )
+        pad_h_1, pad_w_1 = height % 4, width % 4
+        
+        inp_enc_level2, attention_bias, attention_mask = self.down1_2(out_enc_level1, attention_mask, frame, height, width, pad_h=pad_h_1, pad_w=pad_w_1)
+        # frame, height, width = frame // 2 if frame != 1 else frame, (height + pad_h_1) // 2, (width + pad_w_1) // 2
+        height, width = (height + pad_h_1) // 2, (width + pad_w_1) // 2
+
+        # encoder_2
+        out_enc_level2 = inp_enc_level2
+
+        if self.training and self.gradient_checkpointing:
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            
+            for block in self.encoder_level_2:
+                out_enc_level2 = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    out_enc_level2,
+                    attention_bias,
+                    encoder_hidden_states_2,
+                    encoder_attention_mask,
+                    timestep_2,
+                    cross_attention_kwargs,
+                    class_labels,
+                    frame, 
+                    height, 
+                    width, 
+                    **ckpt_kwargs,
+                )
+        else:
+            for block in self.encoder_level_2:
+                out_enc_level2 = block(
+                    out_enc_level2,
+                    attention_mask=attention_bias,
+                    encoder_hidden_states=encoder_hidden_states_2,
+                    encoder_attention_mask=encoder_attention_mask,
+                    timestep=timestep_2,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    frame=frame, 
+                    height=height, 
+                    width=width, 
+                )
+        pad_h_2, pad_w_2 = height % 4, width % 4
+        
+        # import ipdb;ipdb.set_trace()
+        inp_enc_level3, attention_bias, attention_mask = self.down2_3(out_enc_level2, attention_mask, frame, height, width, pad_h=pad_h_2, pad_w=pad_w_2)
+        # frame, height, width = frame // 2 if frame != 1 else frame, (height + pad_h_2) // 2, (width + pad_w_2) // 2
+        height, width = (height + pad_h_2) // 2, (width + pad_w_2) // 2
+
+        # latent
+        latent = inp_enc_level3
+        if self.training and self.gradient_checkpointing:
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            
+            for block in self.latent:
+                latent = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    latent,
+                    attention_bias,
+                    encoder_hidden_states_3,
+                    encoder_attention_mask,
+                    timestep_3,
+                    cross_attention_kwargs,
+                    class_labels,
+                    frame, 
+                    height, 
+                    width, 
+                    **ckpt_kwargs,
+                )
+        else:
+            for block in self.latent:
+                latent = block(
+                    latent,
+                    attention_mask=attention_bias,
+                    encoder_hidden_states=encoder_hidden_states_3,
+                    encoder_attention_mask=encoder_attention_mask,
+                    timestep=timestep_3,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    frame=frame, 
+                    height=height, 
+                    width=width, 
+                )
+
+        # decoder_2
+        
+        # import ipdb;ipdb.set_trace()
+        inp_dec_level2, attention_bias, attention_mask = self.up3_2(latent, attention_mask, frame, height, width, pad_h=pad_h_2, pad_w=pad_w_2)
+        # frame, height, width = frame * 2 if frame != 1 else frame, height * 2 - pad_h_2, width * 2 - pad_w_2
+        height, width = height * 2 - pad_h_2, width * 2 - pad_w_2
+        inp_dec_level2 = torch.cat([inp_dec_level2, out_enc_level2], 2)
+        inp_dec_level2 = self.reduce_chan_level2_norm(inp_dec_level2)
+        inp_dec_level2 = self.reduce_chan_level2(inp_dec_level2)
+        out_dec_level2 = inp_dec_level2
+
+        if self.training and self.gradient_checkpointing:
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            
+            for block in self.decoder_level_2:
+                out_dec_level2 = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    out_dec_level2,
+                    attention_bias,
+                    encoder_hidden_states_2,
+                    encoder_attention_mask,
+                    timestep_2,
+                    cross_attention_kwargs,
+                    class_labels,
+                    frame, 
+                    height, 
+                    width, 
+                    **ckpt_kwargs,
+                )
+        else:
+            for block in self.decoder_level_2:
+                out_dec_level2 = block(
+                    out_dec_level2,
+                    attention_mask=attention_bias,
+                    encoder_hidden_states=encoder_hidden_states_2,
+                    encoder_attention_mask=encoder_attention_mask,
+                    timestep=timestep_2,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    frame=frame, 
+                    height=height, 
+                    width=width, 
+                )
+
+        # decoder_1
+        
+        # import ipdb;ipdb.set_trace()
+        inp_dec_level1, attention_bias, attention_mask = self.up2_1(out_dec_level2, attention_mask, frame, height, width, pad_h=pad_h_1, pad_w=pad_w_1)
+        # frame, height, width = frame * 2 if frame != 1 else frame, height * 2 - pad_h_1, width * 2 - pad_w_1
+        height, width = height * 2 - pad_h_1, width * 2 - pad_w_1
+        inp_dec_level1 = torch.cat([inp_dec_level1, out_enc_level1], 2)
+        inp_dec_level1 = self.reduce_chan_level1_norm(inp_dec_level1)
+        inp_dec_level1 = self.reduce_chan_level1(inp_dec_level1)
+        out_dec_level1 = inp_dec_level1
+
+        if self.training and self.gradient_checkpointing:
+
+            ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            
+            for block in self.decoder_level_1:
+                out_dec_level1 = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    out_dec_level1,
+                    attention_bias,
+                    encoder_hidden_states_1,
+                    encoder_attention_mask,
+                    timestep_1,
+                    cross_attention_kwargs,
+                    class_labels,
+                    frame, 
+                    height, 
+                    width, 
+                    **ckpt_kwargs,
+                )
+        else:
+            for block in self.decoder_level_1:
+                out_dec_level1 = block(
+                    out_dec_level1,
+                    attention_mask=attention_bias,
+                    encoder_hidden_states=encoder_hidden_states_1,
+                    encoder_attention_mask=encoder_attention_mask,
+                    timestep=timestep_1,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    class_labels=class_labels,
+                    frame=frame, 
+                    height=height, 
+                    width=width, 
+                )
+
+        assert not torch.any(torch.isnan(out_dec_level1)), 'after out_dec_level1'
+        # 3. Output
+        output = self._get_output_for_patched_inputs(
+            hidden_states=out_dec_level1,
+            timestep=timestep_1,
+            class_labels=class_labels,
+            embedded_timestep=embedded_timestep_1,
+            num_frames=frame, 
+            height=height,
+            width=width,
+        )  # b c t h w
+        
+        assert not torch.any(torch.isnan(output)), 'after output'
+        frame, height, width = frame * self.config.patch_size_t, height * self.config.patch_size - pad_h_0, width * self.config.patch_size - pad_w_0
+        output = output[:, :, :frame, :height, :width]
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+
+    def _operate_on_patched_inputs(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num):
+        # batch_size = hidden_states.shape[0]
+        hidden_states = self.pos_embed(hidden_states.to(self.dtype), frame)
+
+        if self.use_additional_conditions and added_cond_kwargs is None:
+            raise ValueError(
+                "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`."
+            )
+        timestep_1, embedded_timestep_1 = self.adaln_single_1(
+            timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
+        )  # b 6d, b d
+        timestep_2, embedded_timestep_2 = self.adaln_single_2(
+            timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
+        )  # b 6d, b d
+        timestep_3, embedded_timestep_3 = self.adaln_single_3(
+            timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
+        )  # b 6d, b d
+
+        encoder_hidden_states_1 = self.caption_projection_1(encoder_hidden_states)  # b, 1+use_image_num, l, d or b, 1, l, d
+        encoder_hidden_states_1 = rearrange(encoder_hidden_states_1[:, :1], 'b 1 l d -> (b 1) l d')
+        encoder_hidden_states_2 = self.caption_projection_2(encoder_hidden_states)  # b, 1+use_image_num, l, d or b, 1, l, d
+        encoder_hidden_states_2 = rearrange(encoder_hidden_states_2[:, :1], 'b 1 l d -> (b 1) l d')
+        encoder_hidden_states_3 = self.caption_projection_3(encoder_hidden_states)  # b, 1+use_image_num, l, d or b, 1, l, d
+        encoder_hidden_states_3 = rearrange(encoder_hidden_states_3[:, :1], 'b 1 l d -> (b 1) l d')
+
+
+        return hidden_states, encoder_hidden_states_1, encoder_hidden_states_2, encoder_hidden_states_3, \
+            timestep_1, timestep_2, timestep_3, embedded_timestep_1, embedded_timestep_2, embedded_timestep_3
+
+    
+    
+    def _get_output_for_patched_inputs(
+        self, hidden_states, timestep, class_labels, embedded_timestep, num_frames, height=None, width=None
+    ):  
+        if self.config.norm_type != "ada_norm_single":
+            conditioning = self.transformer_blocks[0].norm1.emb(
+                timestep, class_labels, hidden_dtype=self.dtype
+            )
+            shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+            hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+            hidden_states = self.proj_out_2(hidden_states)
+        elif self.config.norm_type == "ada_norm_single":
+            shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1)
+            hidden_states = self.norm_out(hidden_states)
+            # Modulation
+            hidden_states = hidden_states * (1 + scale) + shift
+            hidden_states = self.proj_out(hidden_states)
+            hidden_states = hidden_states.squeeze(1)
+
+        # # unpatchify
+        # hidden_states = hidden_states.reshape(
+        #     shape=(-1, num_frames, height, width, self.out_channels)
+        # )
+        # output = torch.einsum("nthwc->ncthw", hidden_states)
+        # return output
+            # unpatchify
+        hidden_states = hidden_states.reshape(
+            shape=(-1, num_frames, height, width, self.config.patch_size_t, self.config.patch_size, self.config.patch_size, self.config.out_channels)
+        )
+        hidden_states = torch.einsum("nthwopqc->nctohpwq", hidden_states)
+        output = hidden_states.reshape(
+            shape=(-1, self.config.out_channels, num_frames * self.config.patch_size_t, height * self.config.patch_size, width * self.config.patch_size)
+        )
+        # import ipdb;ipdb.set_trace()
+        # if output.shape[2] % 2 == 0:
+        #     output = output[:, :, 1:]
+        return output
+
+
+def UDiTT2V_S_111(**kwargs):
+    return UDiTT2V(depth=[2, 5, 8, 5, 2], attention_head_dim=24, num_attention_heads=16, patch_size_t=1, patch_size=1, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_S_122(**kwargs):
+    return UDiTT2V(depth=[2, 5, 8, 5, 2], attention_head_dim=24, num_attention_heads=16, patch_size_t=1, patch_size=2, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_B_111(**kwargs):
+    return UDiTT2V(depth=[2, 5, 10, 5, 2], attention_head_dim=24, num_attention_heads=24, patch_size_t=1, patch_size=1, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_B_122(**kwargs):
+    return UDiTT2V(depth=[2, 5, 10, 5, 2], attention_head_dim=24, num_attention_heads=24, patch_size_t=1, patch_size=2, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_L_111(**kwargs):
+    return UDiTT2V(depth=[4, 8, 12, 8, 4], attention_head_dim=24, num_attention_heads=32, patch_size_t=1, patch_size=1, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_L_211(**kwargs):
+    return UDiTT2V(depth=[4, 8, 12, 8, 4], attention_head_dim=24, num_attention_heads=32, patch_size_t=2, patch_size=1, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_L_122(**kwargs):
+    return UDiTT2V(depth=[4, 8, 12, 8, 4], attention_head_dim=24, num_attention_heads=32, patch_size_t=1, patch_size=2, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_L_222(**kwargs):
+    return UDiTT2V(depth=[4, 8, 12, 8, 4], attention_head_dim=24, num_attention_heads=32, patch_size_t=2, patch_size=2, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+def UDiTT2V_XL_111(**kwargs):
+    return UDiTT2V(depth=[4, 8, 12, 8, 4], attention_head_dim=24, num_attention_heads=48, patch_size_t=1, patch_size=1, 
+                   mlp_ratio=2, norm_type="ada_norm_single", caption_channels=4096, **kwargs)
+
+
+UDiT_models = {
+    "UDiTT2V-S/111": UDiTT2V_S_111,  # 0.4B    0.7B if video
+    "UDiTT2V-S/122": UDiTT2V_S_122,  # 0.4B    0.7B if video
+    "UDiTT2V-B/111": UDiTT2V_B_111,  # 1.0B    1.6B if video
+    "UDiTT2V-B/122": UDiTT2V_B_122,  # 1.0B    1.6B if video
+    "UDiTT2V-L/111": UDiTT2V_L_111,  # 2.2B    3.3B if video
+    "UDiTT2V-L/211": UDiTT2V_L_211,  # 2.2B    3.3B if video
+    "UDiTT2V-L/122": UDiTT2V_L_122,  # 2.2B    3.3B if video
+    "UDiTT2V-L/222": UDiTT2V_L_222,  # 2.2B    3.3B if video
+    "UDiTT2V-XL/111": UDiTT2V_XL_111,  # 5.0B    7.4B if video
+}
+
+UDiT_models_class = {
+    "UDiTT2V-S/111": UDiTT2V,
+    "UDiTT2V-S/122": UDiTT2V,
+    "UDiTT2V-B/111": UDiTT2V,
+    "UDiTT2V-B/122": UDiTT2V,
+    "UDiTT2V-L/111": UDiTT2V,
+    "UDiTT2V-L/211": UDiTT2V,
+    "UDiTT2V-L/122": UDiTT2V,
+    "UDiTT2V-L/222": UDiTT2V,
+    "UDiTT2V-XL/111": UDiTT2V,
+}
+
+
+if __name__ == '__main__':
+    import sys
+    from copy import deepcopy
+    from opensora.models.causalvideovae import ae_channel_config, ae_stride_config
+    from opensora.models.causalvideovae import getae, getae_wrapper
+    from opensora.models.causalvideovae import CausalVAEModelWrapper
+
+    args = type('args', (), 
+    {
+        'ae': 'CausalVAEModel_4x8x8', 
+        'attention_mode': 'xformers', 
+        'use_rope': True, 
+        'model_max_length': 300, 
+        'max_height': 240, 
+        'max_width': 320, 
+        'num_frames': 1, 
+        'use_image_num': 0, 
+        'compress_kv_factor': 1, 
+        'interpolation_scale_t': 1,
+        'interpolation_scale_h': 1,
+        'interpolation_scale_w': 1,
+    }
+    )
+    b = 16
+    c = 4
+    cond_c = 4096
+    num_timesteps = 1000
+    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
+    latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w)
+    num_frames = (args.num_frames - 1) // ae_stride_t + 1
+
+    device = torch.device('cuda:0')
+
+
+    
+    model = UDiTT2V_L_122(in_channels=c, 
+                              out_channels=c, 
+                              sample_size=latent_size, 
+                              sample_size_t=num_frames, 
+                              activation_fn="gelu-approximate",
+                            attention_bias=True,
+                            attention_type="default",
+                            double_self_attention=False,
+                            norm_elementwise_affine=False,
+                            norm_eps=1e-06,
+                            norm_num_groups=32,
+                            num_vector_embeds=None,
+                            only_cross_attention=False,
+                            upcast_attention=False,
+                            use_linear_projection=False,
+                            use_additional_conditions=False, 
+                            downsampler=None, 
+                            interpolation_scale_t=args.interpolation_scale_t, 
+                            interpolation_scale_h=args.interpolation_scale_h, 
+                            interpolation_scale_w=args.interpolation_scale_w, 
+                            use_rope=args.use_rope).to(device)
+
+    print(model)
+    print(f'{sum(p.numel() for p in model.parameters() if p.requires_grad)/1e9} B')
+
+
+    model_state_dict = model.state_dict()
+    pretrained = "/storage/ongoing/new/Open-Sora-Plan/bs2_20node_73000k_480p_61x480p_lr5e-5_snr5_noioff0.02_ema_rope_uditultra122_qknorm_ds222_mt5xxl_sucai288w/checkpoint-11500/model_ema/diffusion_pytorch_model.safetensors"
+    try:
+        if 'safetensors' in pretrained:  # pixart series
+            from safetensors.torch import load_file as safe_load
+            # import ipdb;ipdb.set_trace()
+            pretrained_checkpoint = safe_load(pretrained, device="cpu")
+            pretrained_keys = set(list(pretrained_checkpoint.keys()))
+            model_keys = set(list(model_state_dict.keys()))
+            common_keys = list(pretrained_keys & model_keys)
+            checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()}
+            # if checkpoint['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and checkpoint['pos_embed.proj.weight'].ndim == 4:
+            #     logger.info(f"Resize pos_embed, {checkpoint['pos_embed.proj.weight'].shape} -> {model.pos_embed.proj.weight.shape}")
+            #     repeat = model.pos_embed.proj.weight.shape[2]
+            #     checkpoint['pos_embed.proj.weight'] = checkpoint['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat)
+                # del checkpoint['proj_out.weight'], checkpoint['proj_out.bias']
+        else:  # latest stage training weight
+            checkpoint = torch.load(pretrained, map_location='cpu')
+            if 'model' in checkpoint:
+                checkpoint = checkpoint['model']
+        # import ipdb;ipdb.set_trace()
+        missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+        print(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}')
+        print(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!')
+    except Exception as e:
+        print(e)
+
+    # import sys;sys.exit()
+    # try:
+    #     path = "bs32_1node_480p_lr1e-4_snr5_noioff0.02_ema_uditultra22_ds22_mt5xxl/checkpoint-500/model/diffusion_pytorch_model.safetensors"
+    #     from safetensors.torch import load_file as safe_load
+    #     ckpt = safe_load(path, device="cpu")
+    #     new_ckpt = {}
+    #     k_size = 3
+    #     t_stride = 1
+    #     for k, v in ckpt.items():
+    #         if 'pos_embed.proj.weight' in k:
+    #             new_v = v.unsqueeze(-3).repeat(1, 1, k_size, 1, 1)  # 768, 4, 3, 3 -> 768, 4, 3, 3, 3
+    #         elif 'attn1.downsampler.layer.weight' in k:
+    #             new_v = v.unsqueeze(-3).repeat(1, 1, k_size, 1, 1)  # 768, 4, 3, 3 -> 768, 4, 3, 3, 3
+    #         elif 'body.0.weight' in k and 'down' in k:
+    #             in_c = v.shape[0]
+    #             new_v = v[:in_c//2].unsqueeze(-3).repeat(1, 1, k_size, 1, 1)  # 384, 768, 3, 3 -> 192, 768, 3, 3, 3
+    #         elif 'body.0.weight' in k and 'up' in k:
+    #             new_v = v.unsqueeze(-3).repeat(2, 1, k_size, 1, 1)  # 6144, 3072, 3, 3 -> 12288, 3072, 3, 3, 3
+    #         elif 'proj_out' in k:
+    #             if 'weight' in k:
+    #                 new_v = v.repeat(t_stride, 1)  # 16, 768 -> 32, 768
+    #             elif 'bias' in k:
+    #                 new_v = v.repeat(t_stride)  # 16 -> 32
+    #         else:
+    #             new_v = v
+    #         new_ckpt[k] = new_v
+    #     msg = model.load_state_dict(new_ckpt, strict=False)
+    #     # print(msg)
+    # except Exception as e:
+    #     print(e)
+    x = torch.randn(b, c,  1+(args.num_frames-1)//ae_stride_t+args.use_image_num, args.max_height//ae_stride_h, args.max_width//ae_stride_w).to(device)
+    cond = torch.randn(b, 1+args.use_image_num, args.model_max_length, cond_c).to(device)
+    attn_mask = torch.randint(0, 2, (b, 1+(args.num_frames-1)//ae_stride_t+args.use_image_num, args.max_height//ae_stride_h, args.max_width//ae_stride_w)).to(device)  # B L or B 1+num_images L
+    cond_mask = torch.randint(0, 2, (b, 1+args.use_image_num, args.model_max_length)).to(device)  # B L or B 1+num_images L
+    timestep = torch.randint(0, 1000, (b,), device=device)
+    model_kwargs = dict(hidden_states=x, encoder_hidden_states=cond, attention_mask=attn_mask, 
+                        encoder_attention_mask=cond_mask, use_image_num=args.use_image_num, timestep=timestep)
+    with torch.no_grad():
+        output = model(**model_kwargs)[0]
+    import ipdb;ipdb.set_trace()
+    print(output.shape)
+    
+
+
+
+    # from peft import LoraConfig, PeftModel, get_peft_model
+    # from opensora.utils.lora_utils import EMAModel_LoRA, maybe_zero_3, get_peft_state_maybe_zero_3
+    # lora_save_path = '/storage/ongoing/new/Open-Sora-Plan/debug_lora/model_lora'
+    # ema_lora_save_path = '/storage/ongoing/new/Open-Sora-Plan/debug_lora/ema_model_lora'
+    # origin_model_path = '/storage/ongoing/new/Open-Sora-Plan/bs16_4node_240p_lr1e-4_snr5_noioff0.02_ema_rope_uditultra22_ds22_mt5xxl/checkpoint-500/model_ema'
+    # model = UDiTUltraT2V.from_pretrained(origin_model_path)
+    # lora_config = LoraConfig(
+    #     r=64,
+    #     lora_alpha=64,
+    #     init_lora_weights="gaussian",
+    #     target_modules=["to_k", "to_q", "to_v", "to_out.0"],
+    # )
+    # model_lora = get_peft_model(model, lora_config)
+    # # --------------------ema lora_model----------------------------------
+    # # create ema lora_model
+    # ema_model = deepcopy(model_lora)
+    # ema_model_lora = EMAModel_LoRA(lora_config, parameters=ema_model.parameters(), update_after_step=0, 
+    #                                model_cls=UDiTUltraT2V, model_config=ema_model.config)
+    # ema_model_lora.save_pretrained(ema_lora_save_path)
+    # ema_model_load_lora = EMAModel_LoRA.from_pretrained(ema_lora_save_path, UDiTUltraT2V, lora_config, origin_model_path)
+    # ema_model_lora.load_state_dict(ema_model_load_lora.state_dict())
+    # ema_model_lora.to(device)
+
+    # # -----------------lora model---------------------------------
+    # # get lora weight
+    # model_lora.save_pretrained(lora_save_path)
+    # # ----------------load lora model------------------------------
+    # # load lora weight
+    # model = UDiTUltraT2V.from_pretrained(origin_model_path)
+    # import ipdb;ipdb.set_trace()
+    # model_load_lora = PeftModel.from_pretrained(model, lora_save_path)
+    # for k, v in model_load_lora.state_dict().items():
+    #     assert torch.allclose(v, model_lora.state_dict()[k])
+    # # for k, v in zip(ema_model_lora.shadow_params, model_lora.parameters()):
+    # #     assert torch.allclose(v, k)
+    # print('Merging LoRA weights...')
+    # import ipdb;ipdb.set_trace()
+    # model_load_lora_merge = model_load_lora.merge_and_unload()
+    # with torch.no_grad():
+    #     output = model_load_lora_merge(**model_kwargs)
+    # print(output[0].shape)
+
+
+
+
diff --git a/opensora/models/diffusion/udit/modules.py b/opensora/models/diffusion/udit/modules.py
new file mode 100644
index 000000000..631848900
--- /dev/null
+++ b/opensora/models/diffusion/udit/modules.py
@@ -0,0 +1,1173 @@
+from einops import rearrange
+from torch import nn
+import torch
+import numpy as np
+
+from einops import rearrange, repeat
+from typing import Any, Dict, Optional, Tuple
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+import diffusers
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import FeedForward, GatedSelfAttentionDense
+from diffusers.models.attention_processor import Attention as Attention_
+from diffusers.models.embeddings import SinusoidalPositionalEmbedding
+from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm
+from .rope import PositionGetter3D, RoPE3D
+import re
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config, set_run_dtype
+except:
+    torch_npu = None
+    npu_config = None
+    set_run_dtype = None
+logger = logging.get_logger(__name__)
+
+def get_3d_sincos_pos_embed(
+    embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, 
+):
+    """
+    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # if isinstance(grid_size, int):
+    #     grid_size = (grid_size, grid_size)
+    grid_t = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0]
+    grid_h = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1]
+    grid_w = np.arange(grid_size[2], dtype=np.float32) / (grid_size[2] / base_size[2]) / interpolation_scale[2]
+    grid = np.meshgrid(grid_w, grid_h, grid_t)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([3, 1, grid_size[2], grid_size[1], grid_size[0]])
+    pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_3d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 3 != 0:
+        raise ValueError("embed_dim must be divisible by 3")
+
+    # use 1/3 of dimensions to encode grid_t/h/w
+    emb_t = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0])  # (T*H*W, D/3)
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1])  # (T*H*W, D/3)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2])  # (T*H*W, D/3)
+
+    emb = np.concatenate([emb_t, emb_h, emb_w], axis=1)  # (T*H*W, D)
+    return emb
+
+
+def get_2d_sincos_pos_embed(
+    embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, 
+):
+    """
+    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # if isinstance(grid_size, int):
+    #     grid_size = (grid_size, grid_size)
+
+    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0]
+    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1]
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    # use 1/3 of dimensions to encode grid_t/h/w
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+def get_1d_sincos_pos_embed(
+    embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, 
+):
+    """
+    grid_size: int of the grid return: pos_embed: [grid_size, embed_dim] or
+    [1+grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    # if isinstance(grid_size, int):
+    #     grid_size = (grid_size, grid_size)
+
+    grid = np.arange(grid_size, dtype=np.float32) / (grid_size / base_size) / interpolation_scale
+    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid)  # (H*W, D/2)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+class FP32_Layernorm(nn.LayerNorm):
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        origin_dtype = inputs.dtype
+        return F.layer_norm(inputs.float(), self.normalized_shape, self.weight.float() if self.weight is not None else None, 
+                            self.bias.float() if self.bias is not None else None, self.eps).to(origin_dtype)
+
+
+class FP32_SiLU(nn.SiLU):
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.silu(inputs.float(), inplace=self.inplace).to(inputs.dtype)
+
+
+class FP32_GELU(nn.GELU):
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.gelu(inputs.float(), approximate=self.approximate).to(inputs.dtype)
+
+
+
+class OverlapPatchEmbed3D(nn.Module):
+    """2D Image to Patch Embedding but with 3D position embedding"""
+
+    def __init__(
+        self,
+        num_frames=1, 
+        height=224,
+        width=224,
+        patch_size_t=1,
+        patch_size=16,
+        in_channels=3,
+        embed_dim=768,
+        layer_norm=False,
+        flatten=True,
+        bias=True,
+        interpolation_scale=(1, 1),
+        interpolation_scale_t=1,
+        use_abs_pos=False, 
+    ):
+        super().__init__()
+        # assert patch_size_t == 1 and patch_size == 1
+        self.use_abs_pos = use_abs_pos
+        self.flatten = flatten
+        self.layer_norm = layer_norm
+        self.proj = nn.Conv3d(
+            in_channels, embed_dim, kernel_size=3, padding=1, stride=(patch_size_t, patch_size, patch_size), bias=bias
+        )
+        if layer_norm:
+            self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm = None
+
+        self.patch_size_t = patch_size_t
+        self.patch_size = patch_size
+        # See:
+        # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+        self.height, self.width = height // patch_size, width // patch_size
+        self.base_size = (height // patch_size, width // patch_size)
+        self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+        pos_embed = get_2d_sincos_pos_embed(
+            embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+        )
+        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+        self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.interpolation_scale_t = interpolation_scale_t
+        temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t)
+        self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+        # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+    def forward(self, latent, num_frames):
+        b, _, _, _, _ = latent.shape
+        video_latent, image_latent = None, None
+        # b c 1 h w
+        # assert latent.shape[-3] == 1 and num_frames == 1
+        
+        num_frames = latent.shape[-3] // self.patch_size_t
+        height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+        # latent = rearrange(latent, 'b c t h w -> (b t) c h w')
+        latent = self.proj(latent)
+
+        if self.flatten:
+            # latent = latent.flatten(2).transpose(1, 2)  # BT C H W -> BT N C
+            latent = rearrange(latent, 'b c t h w -> (b t) (h w) c ')
+        if self.layer_norm:
+            latent = self.norm(latent)
+
+        if self.use_abs_pos:
+            # Interpolate positional embeddings if needed.
+            # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+            if self.height != height or self.width != width:
+                # raise NotImplementedError
+                pos_embed = get_2d_sincos_pos_embed(
+                    embed_dim=self.pos_embed.shape[-1],
+                    grid_size=(height, width),
+                    base_size=self.base_size,
+                    interpolation_scale=self.interpolation_scale,
+                )
+                pos_embed = torch.from_numpy(pos_embed)
+                pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                pos_embed = self.pos_embed
+
+
+            if self.num_frames != num_frames:
+                # import ipdb;ipdb.set_trace()
+                # raise NotImplementedError
+                temp_pos_embed = get_1d_sincos_pos_embed(
+                    embed_dim=self.temp_pos_embed.shape[-1],
+                    grid_size=num_frames,
+                    base_size=self.base_size_t,
+                    interpolation_scale=self.interpolation_scale_t,
+                )
+                temp_pos_embed = torch.from_numpy(temp_pos_embed)
+                temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                temp_pos_embed = self.temp_pos_embed
+
+            latent = (latent + pos_embed).to(latent.dtype)
+        
+        latent = rearrange(latent, '(b t) n c -> b t n c', b=b)
+        video_latent = latent
+
+        if self.use_abs_pos:
+            # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+            temp_pos_embed = temp_pos_embed.unsqueeze(2)
+            video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype)
+
+        video_latent = rearrange(video_latent, 'b t n c -> b (t n) c')
+        return video_latent
+    
+
+
+class PatchEmbed2D(nn.Module):
+    """2D Image to Patch Embedding but with 3D position embedding"""
+
+    def __init__(
+        self,
+        num_frames=1, 
+        height=224,
+        width=224,
+        patch_size_t=1,
+        patch_size=16,
+        in_channels=3,
+        embed_dim=768,
+        layer_norm=False,
+        flatten=True,
+        bias=True,
+        interpolation_scale=(1, 1),
+        interpolation_scale_t=1,
+        use_abs_pos=True, 
+    ):
+        super().__init__()
+        # assert num_frames == 1
+        self.use_abs_pos = use_abs_pos
+        self.flatten = flatten
+        self.layer_norm = layer_norm
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias
+        )
+        if layer_norm:
+            self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
+        else:
+            self.norm = None
+
+        self.patch_size_t = patch_size_t
+        self.patch_size = patch_size
+        # See:
+        # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161
+
+        self.height, self.width = height // patch_size, width // patch_size
+        self.base_size = (height // patch_size, width // patch_size)
+        self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1])
+        pos_embed = get_2d_sincos_pos_embed(
+            embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale
+        )
+        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False)
+
+        self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t
+        self.interpolation_scale_t = interpolation_scale_t
+        temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t)
+        self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+        # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0]))
+
+    def forward(self, latent, num_frames):
+        b, _, _, _, _ = latent.shape
+        video_latent, image_latent = None, None
+        # b c 1 h w
+        # assert latent.shape[-3] == 1 and num_frames == 1
+        height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size
+        latent = rearrange(latent, 'b c t h w -> (b t) c h w')
+        latent = self.proj(latent)
+
+        if self.flatten:
+            latent = latent.flatten(2).transpose(1, 2)  # BT C H W -> BT N C
+        if self.layer_norm:
+            latent = self.norm(latent)
+
+        if self.use_abs_pos:
+            # Interpolate positional embeddings if needed.
+            # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160)
+            if self.height != height or self.width != width:
+                # raise NotImplementedError
+                pos_embed = get_2d_sincos_pos_embed(
+                    embed_dim=self.pos_embed.shape[-1],
+                    grid_size=(height, width),
+                    base_size=self.base_size,
+                    interpolation_scale=self.interpolation_scale,
+                )
+                pos_embed = torch.from_numpy(pos_embed)
+                pos_embed = pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                pos_embed = self.pos_embed
+
+
+            if self.num_frames != num_frames:
+                # import ipdb;ipdb.set_trace()
+                # raise NotImplementedError
+                if get_sequence_parallel_state():
+                    if npu_config is not None:
+                        sp_size = hccl_info.world_size
+                        temp_pos_embed = get_1d_sincos_pos_embed(
+                            embed_dim=self.temp_pos_embed.shape[-1],
+                            grid_size=num_frames * sp_size,
+                            base_size=self.base_size_t,
+                            interpolation_scale=self.interpolation_scale_t,
+                        )
+                        rank = hccl_info.rank % sp_size
+                        st_frame = rank * num_frames
+                        ed_frame = st_frame + num_frames
+                        temp_pos_embed = temp_pos_embed[st_frame: ed_frame]
+                    else:
+                        sp_size = nccl_info.world_size
+                        temp_pos_embed = get_1d_sincos_pos_embed(
+                            embed_dim=self.temp_pos_embed.shape[-1],
+                            grid_size=num_frames * sp_size,
+                            base_size=self.base_size_t,
+                            interpolation_scale=self.interpolation_scale_t,
+                        )
+                        rank = nccl_info.rank % sp_size
+                        st_frame = rank * num_frames
+                        ed_frame = st_frame + num_frames
+                        temp_pos_embed = temp_pos_embed[st_frame: ed_frame]
+
+                else:
+                    temp_pos_embed = get_1d_sincos_pos_embed(
+                        embed_dim=self.temp_pos_embed.shape[-1],
+                        grid_size=num_frames,
+                        base_size=self.base_size_t,
+                        interpolation_scale=self.interpolation_scale_t,
+                    )
+                temp_pos_embed = torch.from_numpy(temp_pos_embed)
+                temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device)
+            else:
+                temp_pos_embed = self.temp_pos_embed
+
+            latent = (latent + pos_embed).to(latent.dtype)
+        
+        latent = rearrange(latent, '(b t) n c -> b t n c', b=b)
+        video_latent = latent
+
+        if self.use_abs_pos:
+            # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh()
+            temp_pos_embed = temp_pos_embed.unsqueeze(2)
+            video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) 
+        video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') 
+
+        return video_latent
+    
+    
+
+class Attention(Attention_):
+    def __init__(self, downsampler, attention_mode, use_rope, interpolation_scale_thw, **kwags):
+        processor = AttnProcessor2_0(attention_mode=attention_mode, use_rope=use_rope, interpolation_scale_thw=interpolation_scale_thw)
+        super().__init__(processor=processor, **kwags)
+        self.downsampler = None
+        # if downsampler: # downsampler  k155_s122
+        #     downsampler_ker_size = list(re.search(r'k(\d{2,3})', downsampler).group(1)) # 122
+        #     down_factor = list(re.search(r's(\d{2,3})', downsampler).group(1))
+        #     downsampler_ker_size = [int(i) for i in downsampler_ker_size]
+        #     downsampler_padding = [(i - 1) // 2 for i in downsampler_ker_size]
+        #     down_factor = [int(i) for i in down_factor]
+            
+        #     if len(downsampler_ker_size) == 2:
+        #         self.downsampler = DownSampler2d(kwags['query_dim'], kwags['query_dim'], kernel_size=downsampler_ker_size, stride=1,
+        #                                     padding=downsampler_padding, groups=kwags['query_dim'], down_factor=down_factor,
+        #                                     down_shortcut=True)
+        #     elif len(downsampler_ker_size) == 3:
+        #         self.downsampler = DownSampler3d(kwags['query_dim'], kwags['query_dim'], kernel_size=downsampler_ker_size, stride=1,
+        #                                     padding=downsampler_padding, groups=kwags['query_dim'], down_factor=down_factor,
+        #                                     down_shortcut=True)
+        
+        # self.q_norm = nn.LayerNorm(kwags['dim_head'], elementwise_affine=True, eps=1e-6)
+        # self.k_norm = nn.LayerNorm(kwags['dim_head'], elementwise_affine=True, eps=1e-6) 
+        
+
+class DownSampler3d(nn.Module):
+    def __init__(self, *args, **kwargs):
+        ''' Required kwargs: down_factor, downsampler'''
+        super().__init__()
+        self.down_factor = kwargs.pop('down_factor')
+        self.down_shortcut = kwargs.pop('down_shortcut')
+        self.layer = nn.Conv3d(*args, **kwargs)
+
+    def forward(self, x, attention_mask, t, h, w):
+        b = x.shape[0]
+        x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w)
+        if npu_config is None:
+            x = self.layer(x) + (x if self.down_shortcut else 0)
+        else:
+            x_dtype = x.dtype
+            x = npu_config.run_conv3d(self.layer, x, x_dtype) + (x if self.down_shortcut else 0)
+        
+        self.t = t//self.down_factor[0]
+        self.h = h//self.down_factor[1]
+        self.w = w//self.down_factor[2]
+        # import ipdb;ipdb.set_trace()
+        x = rearrange(x, 'b d (t dt) (h dh) (w dw) -> (b dt dh dw) (t h w) d', 
+                      t=t//self.down_factor[0], h=h//self.down_factor[1], w=w//self.down_factor[2], 
+                         dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2])
+    
+
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=t, h=h, w=w)
+        attention_mask = rearrange(attention_mask, 'b 1 (t dt) (h dh) (w dw) -> (b dt dh dw) 1 (t h w)',
+                      t=t//self.down_factor[0], h=h//self.down_factor[1], w=w//self.down_factor[2], 
+                         dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2])
+        return x, attention_mask
+    
+    def reverse(self, x, t, h, w):
+        # import ipdb;ipdb.set_trace()
+        x = rearrange(x, '(b dt dh dw) (t h w) d -> b (t dt h dh w dw) d', 
+                      t=t, h=h, w=w, 
+                         dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2])
+        return x
+
+
+
+class DownSampler2d(nn.Module):
+    def __init__(self, *args, **kwargs):
+        ''' Required kwargs: down_factor, downsampler'''
+        super().__init__()
+        self.down_factor = kwargs.pop('down_factor')
+        self.down_shortcut = kwargs.pop('down_shortcut')
+        self.layer = nn.Conv2d(*args, **kwargs)
+
+    def forward(self, x, attention_mask, t, h, w):
+        # import ipdb;ipdb.set_trace()
+        b = x.shape[0]
+        x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w)
+        x = self.layer(x) + (x if self.down_shortcut else 0)
+
+        self.t = 1
+        self.h = h//self.down_factor[0]
+        self.w = w//self.down_factor[1]
+
+        x = rearrange(x, 'b d (h dh) (w dw) -> (b dh dw) (h w) d', 
+                      h=h//self.down_factor[0], w=w//self.down_factor[1], 
+                      dh=self.down_factor[0], dw=self.down_factor[1])
+    
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> (b t) 1 h w', h=h, w=w)
+        attention_mask = rearrange(attention_mask, 'b 1 (h dh) (w dw) -> (b dh dw) 1 (h w)',
+                      h=h//self.down_factor[0], w=w//self.down_factor[1], 
+                         dh=self.down_factor[0], dw=self.down_factor[1])
+        return x, attention_mask
+    
+    def reverse(self, x, t, h, w):
+        x = rearrange(x, '(b t dh dw) (h w) d -> b (t h dh w dw) d', 
+                      t=t, h=h, w=w, 
+                      dh=self.down_factor[0], dw=self.down_factor[1])
+        return x
+    
+
+class AttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self, attention_mode='xformers', use_rope=False, interpolation_scale_thw=(1, 1, 1)):
+        self.use_rope = use_rope
+        self.interpolation_scale_thw = interpolation_scale_thw
+        if self.use_rope:
+            self._init_rope(interpolation_scale_thw)
+        self.attention_mode = attention_mode
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def _init_rope(self, interpolation_scale_thw):
+        self.rope = RoPE3D(interpolation_scale_thw=interpolation_scale_thw)
+        self.position_getter = PositionGetter3D()
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        frame: int = 8, 
+        height: int = 16, 
+        width: int = 16, 
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+        
+        if attn.downsampler is not None:
+            hidden_states, attention_mask = attn.downsampler(hidden_states, attention_mask, t=frame, h=height, w=width)
+            frame, height, width = attn.downsampler.t, attn.downsampler.h, attn.downsampler.w
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        # if input_ndim == 4:
+        #     batch_size, channel, height, width = hidden_states.shape
+        #     hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            if npu_config is None:
+                attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+                # scaled_dot_product_attention expects attention_mask shape to be
+                # (batch, heads, source_length, target_length)
+                attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+            else:
+                if npu_config.enable_FA:
+                    attention_mask = attention_mask.to(torch.bool)
+                attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1])
+                attention_mask = attention_mask.repeat(1, 1, hidden_states.shape[1], 1)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+            
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        if npu_config is not None and npu_config.on_npu:
+            if npu_config.enable_FA and query.dtype == torch.float32:
+                dtype = torch.bfloat16
+            else:
+                dtype = None
+
+            if self.use_rope:
+                query = query.view(batch_size, -1, attn.heads, head_dim)
+                key = key.view(batch_size, -1, attn.heads, head_dim)
+                # require the shape of (batch_size x nheads x ntokens x dim)
+                pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device)
+                query = self.rope(query, pos_thw)
+                key = self.rope(key, pos_thw)
+                query = query.view(batch_size, -1, attn.heads * head_dim)
+                key = key.view(batch_size, -1, attn.heads * head_dim)
+
+            with set_run_dtype(query, dtype):
+                query, key, value = npu_config.set_current_run_dtype([query, key, value])
+                hidden_states = npu_config.run_attention(query, key, value, attention_mask, "BSH",
+                                                         head_dim, attn.heads)
+
+                hidden_states = npu_config.restore_dtype(hidden_states)
+        else:
+            query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            # qk norm
+            # query = attn.q_norm(query)
+            # key = attn.k_norm(key)
+
+            if self.use_rope:
+                # require the shape of (batch_size x nheads x ntokens x dim)
+                pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device)
+                query = self.rope(query, pos_thw)
+                key = self.rope(key, pos_thw)
+
+            value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            # 0, -10000 ->(bool) False, True ->(any) True ->(not) False
+            # 0, 0 ->(bool) False, False ->(any) False ->(not) True
+            if attention_mask is None or not torch.any(attention_mask.bool()):  # 0 mean visible
+                attention_mask = None
+            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+            # TODO: add support for attn.scale when we move to Torch 2.1
+            
+            if self.attention_mode == 'flash':
+                assert attention_mask is None, 'flash-attn do not support attention_mask'
+                with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
+                    hidden_states = F.scaled_dot_product_attention(
+                        query, key, value, dropout_p=0.0, is_causal=False
+                    )
+            elif self.attention_mode == 'xformers':
+                with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
+                    hidden_states = F.scaled_dot_product_attention(
+                        query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+                    )
+            elif self.attention_mode == 'math':
+                hidden_states = F.scaled_dot_product_attention(
+                    query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+                )
+            else:
+                raise NotImplementedError(f'Found attention_mode: {self.attention_mode}')
+            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        if attn.downsampler is not None:
+            hidden_states = attn.downsampler.reverse(hidden_states, t=frame, h=height, w=width)
+
+        return hidden_states
+
+
+class PixelUnshuffle(nn.Module):
+    def __init__(self, ratio, ratio_t=None):
+        super().__init__()
+        self.r = ratio
+        self.r_t = ratio_t if ratio_t else 1
+
+    def forward(self, x):
+        # if self.r_t is not None and self.r_t != 1:
+        if x.ndim == 5:
+            b, c, t, h, w = x.shape
+            # import ipdb;ipdb.set_trace()
+            assert t % self.r_t == 0 and h % self.r == 0 and w % self.r == 0
+            x = rearrange(x, 'b c (t r1) (h r2) (w r3) -> b (c r1 r2 r3) t h w', r1=self.r_t, r2=self.r, r3=self.r)
+        else:
+            b, c, h, w = x.shape
+            assert h % self.r == 0 and w % self.r == 0
+            x = rearrange(x, 'b c (h r2) (w r3) -> b (c r2 r3) h w', r2=self.r, r3=self.r)
+        return x
+    
+class PixelShuffle(nn.Module):
+    def __init__(self, ratio, ratio_t=None):
+        super().__init__()
+        self.r = ratio
+        self.r_t = ratio_t if ratio_t else 1
+
+    def forward(self, x):
+        if x.ndim == 5:
+            b, c, t, h, w = x.shape
+            # import ipdb;ipdb.set_trace()
+            assert c % (self.r_t*self.r*self.r) == 0
+            x = rearrange(x, 'b (c r1 r2 r3) t h w -> b c (t r1) (h r2) (w r3)', r1=self.r_t, r2=self.r, r3=self.r)
+        else:
+            b, c, h, w = x.shape
+            assert c % (self.r*self.r) == 0
+            x = rearrange(x, 'b (c r2 r3) h w -> b c (h r2) (w r3)', r2=self.r, r3=self.r)
+        return x
+
+class Downsample3d(nn.Module):
+    def __init__(self, n_feat):
+        super(Downsample3d, self).__init__()
+
+        self.body = nn.Sequential(nn.Conv3d(n_feat, n_feat // 4, kernel_size=3, stride=1, padding=1, bias=False),
+                                  PixelUnshuffle(2, 2))
+
+    def forward(self, x, attention_mask, frames, height, width, pad_h=0, pad_w=0):
+        x = rearrange(x, 'b (t h w) d -> b d t h w', t=frames, h=height, w=width)
+        # x = F.pad(x, (0, pad_w, 0, pad_h, 0, 0), mode='reflect')
+        x = F.pad(x, (0, pad_w, 0, pad_h, 0, 0))
+        if npu_config is None:
+            x = self.body(x)
+        else:
+            x_dtype = x.dtype
+            x = npu_config.run_conv3d(self.body, x, x_dtype)
+        x = rearrange(x, 'b d t h w -> b (t h w) d')
+        
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=frames, h=height, w=width)
+        attention_mask = F.pad(attention_mask, (0, pad_w, 0, pad_h, 0, 0))
+        attention_mask = F.max_pool3d(attention_mask, kernel_size=2, stride=2)
+        attention_mask = rearrange(attention_mask, 'b 1 t h w -> b 1 (t h w)')
+        attention_bias = (1 - attention_mask.bool().to(x.dtype)) * -10000.0
+
+        return x, attention_bias, attention_mask
+    
+class Upsample3d(nn.Module):
+    def __init__(self, n_feat):
+        super(Upsample3d, self).__init__()
+
+        self.body = nn.Sequential(nn.Conv3d(n_feat, n_feat * 4, kernel_size=3, stride=1, padding=1, bias=False),
+                                  PixelShuffle(2, 2))
+
+    def forward(self, x, attention_mask, frames, height, width, pad_h=0, pad_w=0):
+        x = rearrange(x, 'b (t h w) d -> b d t h w', t=frames, h=height, w=width)
+        if npu_config is None:
+            x = self.body(x)
+        else:
+            x_dtype = x.dtype
+            x = npu_config.run_conv3d(self.body, x, x_dtype)
+        x = x[:, :, :, :height*2-pad_h, :width*2-pad_w]
+        x = rearrange(x, 'b d t h w -> b (t h w) d')
+
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=frames, h=height, w=width)
+        attention_mask = attention_mask.repeat_interleave(2, -1).repeat_interleave(2, -2).repeat_interleave(2, -3)
+        attention_mask = attention_mask[:, :, :, :height*2-pad_h, :width*2-pad_w]
+        attention_mask = rearrange(attention_mask, 'b 1 t h w -> b 1 (t h w)')
+        attention_bias = (1 - attention_mask.bool().to(x.dtype)) * -10000.0
+
+        return x, attention_bias, attention_mask
+
+
+class Downsample2d(nn.Module):
+    def __init__(self, n_feat, is_video_model=False):
+        super(Downsample2d, self).__init__()
+        self.body = nn.Conv2d(n_feat, 2*n_feat, kernel_size=3, stride=2, padding=1, bias=False)
+
+    def forward(self, x, attention_mask, frames, height, width, pad_h=0, pad_w=0):
+        # import ipdb;ipdb.set_trace()
+        x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=frames, h=height, w=width)
+        if npu_config is None:
+            x = F.pad(x, (0, pad_w, 0, pad_h))
+        else:
+            x = npu_config.run_pad_2d(F.pad, x, pad=(0, pad_w, 0, pad_h))
+        x = self.body(x)
+        x = rearrange(x, '(b t) d h w -> b (t h w) d', t=frames)
+        
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> (b t) 1 h w', t=frames, h=height, w=width)
+        attention_mask = F.pad(attention_mask, (0, pad_w, 0, pad_h))
+        attention_mask = F.max_pool2d(attention_mask.float(), kernel_size=2, stride=2)
+        attention_mask = rearrange(attention_mask, '(b t) 1 h w -> b 1 (t h w)', t=frames)
+        attention_bias = (1 - attention_mask.bool().to(x.dtype)) * -10000.0
+
+        return x, attention_bias, attention_mask
+    
+class Upsample2d(nn.Module):
+    def __init__(self, n_feat, is_video_model=False):
+        super(Upsample2d, self).__init__()
+        self.body = nn.Conv2d(n_feat, n_feat // 2, kernel_size=3, stride=1, padding=1, bias=False)
+
+    def forward(self, x, attention_mask, frames, height, width, pad_h=0, pad_w=0):
+        x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=frames, h=height, w=width)
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.body(x)
+        x = x[:, :, :height*2-pad_h, :width*2-pad_w]
+        x = rearrange(x, '(b t) d h w -> b (t h w) d', t=frames)
+
+        attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> (b t) 1 h w', t=frames, h=height, w=width)
+        attention_mask = attention_mask.repeat_interleave(2, -1).repeat_interleave(2, -2)
+        attention_mask = attention_mask[:, :, :height*2-pad_h, :width*2-pad_w]
+        attention_mask = rearrange(attention_mask, '(b t) 1 h w -> b 1 (t h w)', t=frames)
+        attention_bias = (1 - attention_mask.bool().to(x.dtype)) * -10000.0
+
+        return x, attention_bias, attention_mask
+
+
+class FeedForward_Conv2d(nn.Module):
+    def __init__(self, downsampler, dim, hidden_features, bias=True):
+        super(FeedForward_Conv2d, self).__init__()
+        
+        self.bias = bias
+
+        self.project_in = nn.Linear(dim, hidden_features, bias=bias)
+        self.act = nn.GELU()
+        self.dwconv = nn.ModuleList([
+            nn.Conv2d(hidden_features, hidden_features, kernel_size=(5, 5), stride=1, padding=(2, 2), dilation=1,
+                        groups=hidden_features, bias=bias),
+            nn.Conv2d(hidden_features, hidden_features, kernel_size=(3, 3), stride=1, padding=(1, 1), dilation=1,
+                        groups=hidden_features, bias=bias),
+            nn.Conv2d(hidden_features, hidden_features, kernel_size=(1, 1), stride=1, padding=(0, 0), dilation=1,
+                        groups=hidden_features, bias=bias)
+        ])
+
+        self.project_out = nn.Linear(hidden_features, dim, bias=bias)
+
+
+    def forward(self, x, t, h, w):
+        # import ipdb;ipdb.set_trace()
+        x = self.project_in(x)
+        x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w)
+        
+        x = self.act(x)
+        out = x
+        for module in self.dwconv:
+            out = out + module(x)
+        out = rearrange(out, '(b t) d h w -> b (t h w) d', t=t, h=h, w=w)
+        x = self.project_out(out)
+        return x
+
+@maybe_allow_in_graph
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        num_embeds_ada_norm (:
+            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        upcast_attention (`bool`, *optional*):
+            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+        final_dropout (`bool` *optional*, defaults to False):
+            Whether to apply a final dropout after the last feed-forward layer.
+        attention_type (`str`, *optional*, defaults to `"default"`):
+            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
+        positional_embeddings (`str`, *optional*, defaults to `None`):
+            The type of positional embeddings to apply to.
+        num_positional_embeddings (`int`, *optional*, defaults to `None`):
+            The maximum number of positional embeddings to apply.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
+        norm_eps: float = 1e-5,
+        final_dropout: bool = False,
+        attention_type: str = "default",
+        positional_embeddings: Optional[str] = None,
+        num_positional_embeddings: Optional[int] = None,
+        ada_norm_continous_conditioning_embedding_dim: Optional[int] = None,
+        ada_norm_bias: Optional[int] = None,
+        ff_inner_dim: Optional[int] = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+        attention_mode: str = "xformers", 
+        downsampler: str = None, 
+        mlp_ratio: int = 4, 
+        use_rope: bool = False,
+        interpolation_scale_thw: Tuple[int] = (1, 1, 1), 
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+
+        # We keep these boolean flags for backward-compatibility.
+        self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
+        self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
+        self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
+        self.use_layer_norm = norm_type == "layer_norm"
+        self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous"
+
+        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
+                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
+            )
+
+        self.norm_type = norm_type
+        self.num_embeds_ada_norm = num_embeds_ada_norm
+
+        if positional_embeddings and (num_positional_embeddings is None):
+            raise ValueError(
+                "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
+            )
+
+        if positional_embeddings == "sinusoidal":
+            self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings)
+        else:
+            self.pos_embed = None
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        if norm_type == "ada_norm":
+            self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
+        elif norm_type == "ada_norm_zero":
+            self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        elif norm_type == "ada_norm_continuous":
+            self.norm1 = AdaLayerNormContinuous(
+                dim,
+                ada_norm_continous_conditioning_embedding_dim,
+                norm_elementwise_affine,
+                norm_eps,
+                ada_norm_bias,
+                "rms_norm",
+            )
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps)
+
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+            out_bias=attention_out_bias,
+            downsampler=downsampler, 
+            attention_mode=attention_mode, 
+            use_rope=use_rope, 
+            interpolation_scale_thw=interpolation_scale_thw, 
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            if norm_type == "ada_norm":
+                self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)
+            elif norm_type == "ada_norm_continuous":
+                self.norm2 = AdaLayerNormContinuous(
+                    dim,
+                    ada_norm_continous_conditioning_embedding_dim,
+                    norm_elementwise_affine,
+                    norm_eps,
+                    ada_norm_bias,
+                    "rms_norm",
+                )
+            else:
+                self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim if not double_self_attention else None,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+                out_bias=attention_out_bias,
+                downsampler=None, 
+                attention_mode=attention_mode, 
+                use_rope=False, 
+                interpolation_scale_thw=interpolation_scale_thw, 
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        if norm_type == "ada_norm_continuous":
+            self.norm3 = AdaLayerNormContinuous(
+                dim,
+                ada_norm_continous_conditioning_embedding_dim,
+                norm_elementwise_affine,
+                norm_eps,
+                ada_norm_bias,
+                "layer_norm",
+            )
+
+        elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm", "ada_norm_continuous"]:
+            self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        elif norm_type == "layer_norm_i2vgen":
+            self.norm3 = None
+
+        self.ff = FeedForward(
+                dim,
+                dropout=dropout,
+                activation_fn=activation_fn,
+                final_dropout=final_dropout,
+                inner_dim=ff_inner_dim,
+                bias=ff_bias,
+            )
+
+        # 4. Fuser
+        if attention_type == "gated" or attention_type == "gated-text-image":
+            self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim)
+
+        # 5. Scale-shift for PixArt-Alpha.
+        if norm_type == "ada_norm_single":
+            self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        frame: int = None, 
+        height: int = None, 
+        width: int = None, 
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> torch.FloatTensor:
+        if cross_attention_kwargs is not None:
+            if cross_attention_kwargs.get("scale", None) is not None:
+                logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
+
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 0. Self-Attention
+        batch_size = hidden_states.shape[0]
+
+        if self.norm_type == "ada_norm":
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.norm_type == "ada_norm_zero":
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]:
+            norm_hidden_states = self.norm1(hidden_states)
+        elif self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif self.norm_type == "ada_norm_single":
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+            ).chunk(6, dim=1)
+            norm_hidden_states = self.norm1(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+            norm_hidden_states = norm_hidden_states.squeeze(1)
+        else:
+            raise ValueError("Incorrect norm used")
+
+        if self.pos_embed is not None:
+            norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+        # 1. Prepare GLIGEN inputs
+        cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
+        gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
+        
+        # 0, -10000 ->(bool) False, True ->(any) True ->(not) False
+        # 0, 0 ->(bool) False, False ->(any) False ->(not) True
+        # assert attention_mask.bool().float().sum() / attention_mask.bool().float().numel() <= 1/16, 'must ~all visible'
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask, frame=frame, height=height, width=width, 
+            **cross_attention_kwargs,
+        )
+        if self.norm_type == "ada_norm_zero":
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        elif self.norm_type == "ada_norm_single":
+            attn_output = gate_msa * attn_output
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        # 1.2 GLIGEN Control
+        if gligen_kwargs is not None:
+            hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
+
+        # 3. Cross-Attention
+        if self.attn2 is not None:
+            if self.norm_type == "ada_norm":
+                norm_hidden_states = self.norm2(hidden_states, timestep)
+            elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]:
+                norm_hidden_states = self.norm2(hidden_states)
+            elif self.norm_type == "ada_norm_single":
+                # For PixArt norm2 isn't applied here:
+                # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
+                norm_hidden_states = hidden_states
+            elif self.norm_type == "ada_norm_continuous":
+                norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"])
+            else:
+                raise ValueError("Incorrect norm")
+
+            if self.pos_embed is not None and self.norm_type != "ada_norm_single":
+                norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                **cross_attention_kwargs,
+            )
+            hidden_states = attn_output + hidden_states
+        # 4. Feed-forward
+        # i2vgen doesn't have this norm 🤷‍♂️
+        if self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif not self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm3(hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm2(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+        # if self.downsampler:
+        #     ff_output = self.ff(norm_hidden_states, t=frame, h=height, w=width)
+        # else:
+        ff_output = self.ff(norm_hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+        elif self.norm_type == "ada_norm_single":
+            ff_output = gate_mlp * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
diff --git a/opensora/models/diffusion/udit/rope.py b/opensora/models/diffusion/udit/rope.py
new file mode 100644
index 000000000..b42096b9d
--- /dev/null
+++ b/opensora/models/diffusion/udit/rope.py
@@ -0,0 +1,154 @@
+import torch
+
+class PositionGetter3D(object):
+    """ return positions of patches """
+
+    def __init__(self, ):
+        self.cache_positions = {}
+        
+    def __call__(self, b, t, h, w, device):
+        if not (t,h,w) in self.cache_positions:
+            x = torch.arange(w, device=device)
+            y = torch.arange(h, device=device)
+            z = torch.arange(t, device=device)
+            self.cache_positions[t,h,w] = torch.cartesian_prod(z, y, x) # (t, h, w, 3)
+        pos = self.cache_positions[t,h,w].view(1, t*h*w, 3).expand(b, -1, 3).clone()
+        return pos
+    
+
+class RoPE3D(torch.nn.Module):
+
+    def __init__(self, freq=10000.0, F0=1.0, interpolation_scale_thw=(1, 1, 1)):
+        super().__init__()
+        self.base = freq
+        self.F0 = F0
+        self.interpolation_scale_t = interpolation_scale_thw[0]
+        self.interpolation_scale_h = interpolation_scale_thw[1]
+        self.interpolation_scale_w = interpolation_scale_thw[2]
+        self.cache = {}
+
+    def get_cos_sin(self, D, seq_len, device, dtype, interpolation_scale=1):
+        if (D, seq_len, device, dtype) not in self.cache:
+            inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
+            t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) / interpolation_scale
+            freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
+            freqs = torch.cat((freqs, freqs), dim=-1)
+            cos = freqs.cos()  # (Seq, Dim)
+            sin = freqs.sin()
+            self.cache[D, seq_len, device, dtype] = (cos, sin)
+        return self.cache[D, seq_len, device, dtype]
+
+    @staticmethod
+    def rotate_half(x):
+        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def apply_rope1d(self, tokens, pos1d, cos, sin):
+        assert pos1d.ndim == 2
+        cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
+        sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
+        return (tokens * cos) + (self.rotate_half(tokens) * sin)
+
+    def forward(self, tokens, positions):
+        """
+        input:
+            * tokens: batch_size x nheads x ntokens x dim
+            * positions: batch_size x ntokens x 3 (t, y and x position of each token)
+        output:
+            * tokens after appplying RoPE3D (batch_size x nheads x ntokens x dim)
+        """
+        assert tokens.size(3) % 3 == 0, "number of dimensions should be a multiple of three"
+        D = tokens.size(3) // 3
+        assert positions.ndim == 3 and positions.shape[-1] == 3  # Batch, Seq, 3
+        cos_t, sin_t = self.get_cos_sin(D, int(positions[:, :, 0].max()) + 1, tokens.device, tokens.dtype, self.interpolation_scale_t)
+        cos_y, sin_y = self.get_cos_sin(D, int(positions[:, :, 1].max()) + 1, tokens.device, tokens.dtype, self.interpolation_scale_h)
+        cos_x, sin_x = self.get_cos_sin(D, int(positions[:, :, 2].max()) + 1, tokens.device, tokens.dtype, self.interpolation_scale_w)
+        # split features into three along the feature dimension, and apply rope1d on each half
+        t, y, x = tokens.chunk(3, dim=-1)
+        t = self.apply_rope1d(t, positions[:, :, 0], cos_t, sin_t)
+        y = self.apply_rope1d(y, positions[:, :, 1], cos_y, sin_y)
+        x = self.apply_rope1d(x, positions[:, :, 2], cos_x, sin_x)
+        tokens = torch.cat((t, y, x), dim=-1)
+        return tokens
+
+
+
+# import torch
+# from einops import rearrange, repeat
+
+# class PositionGetter3D(object):
+#     """ return positions of patches """
+
+#     def __init__(self, ):
+#         self.cache_positions = {}
+        
+#     def __call__(self, b, t, h, w, device):
+#         if not (t,h,w) in self.cache_positions:
+#             x = torch.arange(w, device=device)
+#             y = torch.arange(h, device=device)
+#             z = torch.arange(t, device=device)
+#             positions = torch.cartesian_prod(z, y, x) # (t, h, w, 3)
+#             positions = rearrange(positions, 'n d -> d 1 n')
+#             positions = repeat(positions, 'd 1 n -> d b n', b=b).clone()
+#             poses = (positions[0], positions[1], positions[2])
+#             max_pos = (int(poses[0].max()), int(poses[1].max()), int(poses[2].max()))
+#             self.cache_positions[t,h,w] = (poses, max_pos)
+#         pos = self.cache_positions[t,h,w]
+#         return pos
+    
+
+# class RoPE3D(torch.nn.Module):
+
+#     def __init__(self, freq=10000.0, F0=1.0, interpolation_scale_thw=(1, 1, 1)):
+#         super().__init__()
+#         self.base = freq
+#         self.F0 = F0
+#         self.interpolation_scale_t = interpolation_scale_thw[0]
+#         self.interpolation_scale_h = interpolation_scale_thw[1]
+#         self.interpolation_scale_w = interpolation_scale_thw[2]
+#         self.cache = {}
+
+#     def get_cos_sin(self, D, seq_len, device, dtype, interpolation_scale=1):
+#         if (D, seq_len, device, dtype) not in self.cache:
+#             inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
+#             t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) / interpolation_scale
+#             freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
+#             freqs = torch.cat((freqs, freqs), dim=-1)
+#             cos = freqs.cos()  # (Seq, Dim)
+#             sin = freqs.sin()
+#             self.cache[D, seq_len, device, dtype] = (cos, sin)
+#         return self.cache[D, seq_len, device, dtype]
+
+#     @staticmethod
+#     def rotate_half(x):
+#         x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+#         return torch.cat((-x2, x1), dim=-1)
+
+#     def apply_rope1d(self, tokens, pos1d, cos, sin):
+#         assert pos1d.ndim == 2
+#         cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
+#         sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
+#         return (tokens * cos) + (self.rotate_half(tokens) * sin)
+
+#     def forward(self, tokens, positions):
+#         """
+#         input:
+#             * tokens: batch_size x nheads x ntokens x dim
+#             * positions: batch_size x ntokens x 3 (t, y and x position of each token)
+#         output:
+#             * tokens after appplying RoPE3D (batch_size x nheads x ntokens x dim)
+#         """
+#         assert tokens.size(3) % 3 == 0, "number of dimensions should be a multiple of three"
+#         D = tokens.size(3) // 3
+#         poses, max_pos = positions
+#         assert len(poses) == 3 and poses[0].ndim == 2   # Batch, Seq, 3
+#         cos_t, sin_t = self.get_cos_sin(D, max_pos[0] + 1, tokens.device, tokens.dtype, self.interpolation_scale_t)
+#         cos_y, sin_y = self.get_cos_sin(D, max_pos[1] + 1, tokens.device, tokens.dtype, self.interpolation_scale_h)
+#         cos_x, sin_x = self.get_cos_sin(D, max_pos[2] + 1, tokens.device, tokens.dtype, self.interpolation_scale_w)
+#         # split features into three along the feature dimension, and apply rope1d on each half
+#         t, y, x = tokens.chunk(3, dim=-1)
+#         t = self.apply_rope1d(t, poses[0], cos_t, sin_t)
+#         y = self.apply_rope1d(y, poses[1], cos_y, sin_y)
+#         x = self.apply_rope1d(x, poses[2], cos_x, sin_x)
+#         tokens = torch.cat((t, y, x), dim=-1)
+#         return tokens
diff --git a/opensora/models/diffusion/utils/curope/__init__.py b/opensora/models/diffusion/utils/curope/__init__.py
new file mode 100644
index 000000000..25e3d48a1
--- /dev/null
+++ b/opensora/models/diffusion/utils/curope/__init__.py
@@ -0,0 +1,4 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+from .curope2d import cuRoPE2D
diff --git a/opensora/models/diffusion/utils/curope/curope.cpp b/opensora/models/diffusion/utils/curope/curope.cpp
new file mode 100644
index 000000000..8fe9058e0
--- /dev/null
+++ b/opensora/models/diffusion/utils/curope/curope.cpp
@@ -0,0 +1,69 @@
+/* 
+  Copyright (C) 2022-present Naver Corporation. All rights reserved.
+  Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+*/
+
+#include <torch/extension.h>
+
+// forward declaration
+void rope_2d_cuda( torch::Tensor tokens, const torch::Tensor pos, const float base, const float fwd );
+
+void rope_2d_cpu( torch::Tensor tokens, const torch::Tensor positions, const float base, const float fwd )
+{
+    const int B = tokens.size(0);
+    const int N = tokens.size(1);
+    const int H = tokens.size(2);
+    const int D = tokens.size(3) / 4;
+
+    auto tok = tokens.accessor<float, 4>();
+    auto pos = positions.accessor<int64_t, 3>();
+
+    for (int b = 0; b < B; b++) {
+      for (int x = 0; x < 2; x++) { // y and then x (2d)
+        for (int n = 0; n < N; n++) {
+        
+            // grab the token position
+            const int p = pos[b][n][x];
+
+            for (int h = 0; h < H; h++) {
+                for (int d = 0; d < D; d++) {
+                    // grab the two values
+                    float u = tok[b][n][h][d+0+x*2*D];
+                    float v = tok[b][n][h][d+D+x*2*D];
+
+                    // grab the cos,sin
+                    const float inv_freq = fwd * p / powf(base, d/float(D));
+                    float c = cosf(inv_freq);
+                    float s = sinf(inv_freq);
+
+                    // write the result
+                    tok[b][n][h][d+0+x*2*D] = u*c - v*s;
+                    tok[b][n][h][d+D+x*2*D] = v*c + u*s;
+                }
+            }
+        }
+      }
+    }
+}
+
+void rope_2d( torch::Tensor tokens,     // B,N,H,D
+        const torch::Tensor positions,  // B,N,2
+        const float base, 
+        const float fwd )
+{
+    TORCH_CHECK(tokens.dim() == 4, "tokens must have 4 dimensions");
+    TORCH_CHECK(positions.dim() == 3, "positions must have 3 dimensions");
+    TORCH_CHECK(tokens.size(0) == positions.size(0), "batch size differs between tokens & positions");
+    TORCH_CHECK(tokens.size(1) == positions.size(1), "seq_length differs between tokens & positions");
+    TORCH_CHECK(positions.size(2) == 2, "positions.shape[2] must be equal to 2");
+    TORCH_CHECK(tokens.is_cuda() == positions.is_cuda(), "tokens and positions are not on the same device" );
+
+    if (tokens.is_cuda())
+        rope_2d_cuda( tokens, positions, base, fwd );
+    else
+        rope_2d_cpu( tokens, positions, base, fwd );
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("rope_2d", &rope_2d, "RoPE 2d forward/backward");
+}
diff --git a/opensora/models/diffusion/utils/curope/curope2d.py b/opensora/models/diffusion/utils/curope/curope2d.py
new file mode 100644
index 000000000..a49c12f8c
--- /dev/null
+++ b/opensora/models/diffusion/utils/curope/curope2d.py
@@ -0,0 +1,40 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+import torch
+
+try:
+    import curope as _kernels # run `python setup.py install`
+except ModuleNotFoundError:
+    from . import curope as _kernels # run `python setup.py build_ext --inplace`
+
+
+class cuRoPE2D_func (torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, tokens, positions, base, F0=1):
+        ctx.save_for_backward(positions)
+        ctx.saved_base = base
+        ctx.saved_F0 = F0
+        # tokens = tokens.clone() # uncomment this if inplace doesn't work
+        _kernels.rope_2d( tokens, positions, base, F0 )
+        ctx.mark_dirty(tokens)
+        return tokens
+
+    @staticmethod
+    def backward(ctx, grad_res):
+        positions, base, F0 = ctx.saved_tensors[0], ctx.saved_base, ctx.saved_F0
+        _kernels.rope_2d( grad_res, positions, base, -F0 )
+        ctx.mark_dirty(grad_res)
+        return grad_res, None, None, None
+
+
+class cuRoPE2D(torch.nn.Module):
+    def __init__(self, freq=100.0, F0=1.0):
+        super().__init__()
+        self.base = freq 
+        self.F0 = F0
+
+    def forward(self, tokens, positions): 
+        cuRoPE2D_func.apply( tokens.transpose(1,2), positions, self.base, self.F0 )
+        return tokens
\ No newline at end of file
diff --git a/opensora/models/diffusion/utils/curope/kernels.cu b/opensora/models/diffusion/utils/curope/kernels.cu
new file mode 100644
index 000000000..7156cd1bb
--- /dev/null
+++ b/opensora/models/diffusion/utils/curope/kernels.cu
@@ -0,0 +1,108 @@
+/* 
+  Copyright (C) 2022-present Naver Corporation. All rights reserved.
+  Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+*/
+
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+
+#define CHECK_CUDA(tensor) {\
+    TORCH_CHECK((tensor).is_cuda(), #tensor " is not in cuda memory"); \
+    TORCH_CHECK((tensor).is_contiguous(), #tensor " is not contiguous"); }
+void CHECK_KERNEL() {auto error = cudaGetLastError(); TORCH_CHECK( error == cudaSuccess, cudaGetErrorString(error));}
+
+
+template < typename scalar_t  >
+__global__ void rope_2d_cuda_kernel( 
+        //scalar_t* __restrict__ tokens, 
+        torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> tokens,
+        const int64_t* __restrict__ pos, 
+        const float base, 
+        const float fwd )
+        // const int N, const int H, const int D )
+{
+    // tokens shape = (B, N, H, D)
+    const int N = tokens.size(1);
+    const int H = tokens.size(2);
+    const int D = tokens.size(3);
+    
+    // each block update a single token, for all heads
+    // each thread takes care of a single output
+    extern __shared__ float shared[];
+    float* shared_inv_freq = shared + D;
+
+    const int b = blockIdx.x / N;
+    const int n = blockIdx.x % N;
+
+    const int Q = D / 4; 
+    // one token = [0..Q : Q..2Q : 2Q..3Q : 3Q..D]
+    //              u_Y     v_Y     u_X      v_X
+
+    // shared memory: first, compute inv_freq
+    if (threadIdx.x < Q)
+        shared_inv_freq[threadIdx.x] = fwd / powf(base, threadIdx.x/float(Q));
+    __syncthreads();
+
+    // start of X or Y part
+    const int X = threadIdx.x < D/2 ? 0 : 1; 
+    const int m = (X*D/2) + (threadIdx.x % Q);   // index of u_Y or u_X
+
+    // grab the cos,sin appropriate for me
+    const float freq = pos[blockIdx.x*2+X] * shared_inv_freq[threadIdx.x % Q];
+    const float cos = cosf(freq);
+    const float sin = sinf(freq);
+    /*
+    float* shared_cos_sin = shared + D + D/4;
+    if ((threadIdx.x % (D/2)) < Q)
+        shared_cos_sin[m+0] = cosf(freq);
+    else
+        shared_cos_sin[m+Q] = sinf(freq);
+    __syncthreads();
+    const float cos = shared_cos_sin[m+0];
+    const float sin = shared_cos_sin[m+Q];
+    */
+
+    for (int h = 0; h < H; h++)
+    {
+        // then, load all the token for this head in shared memory
+        shared[threadIdx.x] = tokens[b][n][h][threadIdx.x];
+        __syncthreads();
+
+        const float u = shared[m];
+        const float v = shared[m+Q];
+        
+        // write output
+        if ((threadIdx.x % (D/2)) < Q)
+            tokens[b][n][h][threadIdx.x] = u*cos - v*sin;
+        else
+            tokens[b][n][h][threadIdx.x] = v*cos + u*sin;
+    }
+}
+
+void rope_2d_cuda( torch::Tensor tokens, const torch::Tensor pos, const float base, const float fwd ) 
+{
+    const int B = tokens.size(0); // batch size
+    const int N = tokens.size(1); // sequence length
+    const int H = tokens.size(2); // number of heads
+    const int D = tokens.size(3); // dimension per head
+
+    TORCH_CHECK(tokens.stride(3) == 1 && tokens.stride(2) == D, "tokens are not contiguous");
+    TORCH_CHECK(pos.is_contiguous(), "positions are not contiguous");
+    TORCH_CHECK(pos.size(0) == B && pos.size(1) == N && pos.size(2) == 2, "bad pos.shape");
+    TORCH_CHECK(D % 4 == 0, "token dim must be multiple of 4");
+
+    // one block for each layer, one thread per local-max
+    const int THREADS_PER_BLOCK = D;
+    const int N_BLOCKS = B * N; // each block takes care of H*D values
+    const int SHARED_MEM = sizeof(float) * (D + D/4);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(tokens.type(), "rope_2d_cuda", ([&] {
+        rope_2d_cuda_kernel<scalar_t> <<<N_BLOCKS, THREADS_PER_BLOCK, SHARED_MEM>>> (
+            //tokens.data_ptr<scalar_t>(), 
+            tokens.packed_accessor32<scalar_t,4,torch::RestrictPtrTraits>(),
+            pos.data_ptr<int64_t>(), 
+            base, fwd); //, N, H, D );
+    }));
+}
diff --git a/opensora/models/diffusion/utils/curope/setup.py b/opensora/models/diffusion/utils/curope/setup.py
new file mode 100644
index 000000000..230632ed0
--- /dev/null
+++ b/opensora/models/diffusion/utils/curope/setup.py
@@ -0,0 +1,34 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+from setuptools import setup
+from torch import cuda
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+# compile for all possible CUDA architectures
+all_cuda_archs = cuda.get_gencode_flags().replace('compute=','arch=').split()
+# alternatively, you can list cuda archs that you want, eg:
+# all_cuda_archs = [
+    # '-gencode', 'arch=compute_70,code=sm_70',
+    # '-gencode', 'arch=compute_75,code=sm_75',
+    # '-gencode', 'arch=compute_80,code=sm_80',
+    # '-gencode', 'arch=compute_86,code=sm_86'
+# ]
+
+setup(
+    name = 'curope',
+    ext_modules = [
+        CUDAExtension(
+                name='curope',
+                sources=[
+                    "curope.cpp",
+                    "kernels.cu",
+                ],
+                extra_compile_args = dict(
+                    nvcc=['-O3','--ptxas-options=-v',"--use_fast_math"]+all_cuda_archs, 
+                    cxx=['-O3'])
+                )
+    ],
+    cmdclass = {
+        'build_ext': BuildExtension
+    })
diff --git a/opensora/models/diffusion/utils/pos_embed.py b/opensora/models/diffusion/utils/pos_embed.py
new file mode 100644
index 000000000..61872bc97
--- /dev/null
+++ b/opensora/models/diffusion/utils/pos_embed.py
@@ -0,0 +1,243 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+# croco: https://github.com/naver/croco
+# diffusers: https://github.com/huggingface/diffusers
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+import numpy as np
+import torch
+
+
+def get_2d_sincos_pos_embed(
+        embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16
+):
+    """
+    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
+    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    if isinstance(grid_size, int):
+        grid_size = (grid_size, grid_size)
+
+    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale
+    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed(
+        embed_dim, length, interpolation_scale=1.0, base_size=16
+):
+    pos = torch.arange(0, length).unsqueeze(1) / interpolation_scale
+    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, pos)
+    return pos_embed
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000 ** omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# ----------------------------------------------------------
+# RoPE2D: RoPE implementation in 2D
+# ----------------------------------------------------------
+
+try:
+    from .curope import cuRoPE2D
+
+    RoPE2D = cuRoPE2D
+except ImportError:
+    print('Warning, cannot find cuda-compiled version of RoPE2D, using a slow pytorch version instead')
+
+
+    class RoPE2D(torch.nn.Module):
+
+        def __init__(self, freq=10000.0, F0=1.0, scaling_factor=1.0):
+            super().__init__()
+            self.base = freq
+            self.F0 = F0
+            self.scaling_factor = scaling_factor
+            self.cache = {}
+
+        def get_cos_sin(self, D, seq_len, device, dtype):
+            if (D, seq_len, device, dtype) not in self.cache:
+                inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
+                t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+                freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
+                freqs = torch.cat((freqs, freqs), dim=-1)
+                cos = freqs.cos()  # (Seq, Dim)
+                sin = freqs.sin()
+                self.cache[D, seq_len, device, dtype] = (cos, sin)
+            return self.cache[D, seq_len, device, dtype]
+
+        @staticmethod
+        def rotate_half(x):
+            x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+            return torch.cat((-x2, x1), dim=-1)
+
+        def apply_rope1d(self, tokens, pos1d, cos, sin):
+            assert pos1d.ndim == 2
+            
+            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
+            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
+            return (tokens * cos) + (self.rotate_half(tokens) * sin)
+
+        def forward(self, tokens, positions):
+            """
+            input:
+                * tokens: batch_size x nheads x ntokens x dim
+                * positions: batch_size x ntokens x 2 (y and x position of each token)
+            output:
+                * tokens after appplying RoPE2D (batch_size x nheads x ntokens x dim)
+            """
+            assert tokens.size(3) % 2 == 0, "number of dimensions should be a multiple of two"
+            D = tokens.size(3) // 2
+            assert positions.ndim == 3 and positions.shape[-1] == 2  # Batch, Seq, 2
+            cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype)
+            # split features into two along the feature dimension, and apply rope1d on each half
+            y, x = tokens.chunk(2, dim=-1)
+            y = self.apply_rope1d(y, positions[:, :, 0], cos, sin)
+            x = self.apply_rope1d(x, positions[:, :, 1], cos, sin)
+            tokens = torch.cat((y, x), dim=-1)
+            return tokens
+
+class LinearScalingRoPE2D(RoPE2D):
+    """Code from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L148"""
+
+    def forward(self, tokens, positions):
+        # difference to the original RoPE: a scaling factor is aplied to the position ids
+        dtype = positions.dtype
+        positions = positions.float() / self.scaling_factor
+        positions = positions.to(dtype)
+        tokens = super().forward(tokens, positions)
+        return tokens
+
+
+try:
+    from .curope import cuRoPE1D
+
+    RoPE1D = cuRoPE1D
+except ImportError:
+    print('Warning, cannot find cuda-compiled version of RoPE2D, using a slow pytorch version instead')
+
+
+    class RoPE1D(torch.nn.Module):
+
+        def __init__(self, freq=10000.0, F0=1.0, scaling_factor=1.0):
+            super().__init__()
+            self.base = freq
+            self.F0 = F0
+            self.scaling_factor = scaling_factor
+            self.cache = {}
+
+        def get_cos_sin(self, D, seq_len, device, dtype):
+            if (D, seq_len, device, dtype) not in self.cache:
+                inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
+                t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+                freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
+                freqs = torch.cat((freqs, freqs), dim=-1)
+                cos = freqs.cos()  # (Seq, Dim)
+                sin = freqs.sin()
+                self.cache[D, seq_len, device, dtype] = (cos, sin)
+            return self.cache[D, seq_len, device, dtype]
+
+        @staticmethod
+        def rotate_half(x):
+            x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+            return torch.cat((-x2, x1), dim=-1)
+
+        def apply_rope1d(self, tokens, pos1d, cos, sin):
+            assert pos1d.ndim == 2
+            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
+            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
+            return (tokens * cos) + (self.rotate_half(tokens) * sin)
+
+        def forward(self, tokens, positions):
+            """
+            input:
+                * tokens: batch_size x nheads x ntokens x dim
+                * positions: batch_size x ntokens (t position of each token)
+            output:
+                * tokens after appplying RoPE2D (batch_size x nheads x ntokens x dim)
+            """
+            D = tokens.size(3)
+            assert positions.ndim == 2  # Batch, Seq
+            cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype)
+            tokens = self.apply_rope1d(tokens, positions, cos, sin)
+            return tokens
+
+class LinearScalingRoPE1D(RoPE1D):
+    """Code from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L148"""
+    
+    def forward(self, tokens, positions):
+        # difference to the original RoPE: a scaling factor is aplied to the position ids
+        dtype = positions.dtype
+        positions = positions.float() / self.scaling_factor
+        positions = positions.to(dtype)
+        tokens = super().forward(tokens, positions)
+        return tokens
+
+
+class PositionGetter2D(object):
+    """ return positions of patches """
+
+    def __init__(self):
+        self.cache_positions = {}
+        
+    def __call__(self, b, h, w, device):
+        if not (h,w) in self.cache_positions:
+            x = torch.arange(w, device=device)
+            y = torch.arange(h, device=device)
+            self.cache_positions[h,w] = torch.cartesian_prod(y, x) # (h, w, 2)
+        pos = self.cache_positions[h,w].view(1, h*w, 2).expand(b, -1, 2).clone()
+        return pos
+
+
+
+class PositionGetter1D(object):
+    """ return positions of patches """
+
+    def __init__(self):
+        self.cache_positions = {}
+        
+    def __call__(self, b, l, device):
+        if not (l) in self.cache_positions:
+            x = torch.arange(l, device=device)
+            self.cache_positions[l] = x # (l, )
+        pos = self.cache_positions[l].view(1, l).expand(b, -1).clone()
+        return pos
\ No newline at end of file
diff --git a/opensora/models/diffusion/utils/transport/__init__.py b/opensora/models/diffusion/utils/transport/__init__.py
new file mode 100644
index 000000000..db68edd20
--- /dev/null
+++ b/opensora/models/diffusion/utils/transport/__init__.py
@@ -0,0 +1,63 @@
+from .transport import Transport, ModelType, WeightType, PathType, Sampler
+
+def create_transport(
+    path_type='Linear',
+    prediction="velocity",
+    loss_weight=None,
+    train_eps=None,
+    sample_eps=None,
+):
+    """function for creating Transport object
+    **Note**: model prediction defaults to velocity
+    Args:
+    - path_type: type of path to use; default to linear
+    - learn_score: set model prediction to score
+    - learn_noise: set model prediction to noise
+    - velocity_weighted: weight loss by velocity weight
+    - likelihood_weighted: weight loss by likelihood weight
+    - train_eps: small epsilon for avoiding instability during training
+    - sample_eps: small epsilon for avoiding instability during sampling
+    """
+
+    if prediction == "noise":
+        model_type = ModelType.NOISE
+    elif prediction == "score":
+        model_type = ModelType.SCORE
+    else:
+        model_type = ModelType.VELOCITY
+
+    if loss_weight == "velocity":
+        loss_type = WeightType.VELOCITY
+    elif loss_weight == "likelihood":
+        loss_type = WeightType.LIKELIHOOD
+    else:
+        loss_type = WeightType.NONE
+
+    path_choice = {
+        "Linear": PathType.LINEAR,
+        "GVP": PathType.GVP,
+        "VP": PathType.VP,
+    }
+
+    path_type = path_choice[path_type]
+
+    if (path_type in [PathType.VP]):
+        train_eps = 1e-5 if train_eps is None else train_eps
+        sample_eps = 1e-3 if train_eps is None else sample_eps
+    elif (path_type in [PathType.GVP, PathType.LINEAR] and model_type != ModelType.VELOCITY):
+        train_eps = 1e-3 if train_eps is None else train_eps
+        sample_eps = 1e-3 if train_eps is None else sample_eps
+    else: # velocity & [GVP, LINEAR] is stable everywhere
+        train_eps = 0
+        sample_eps = 0
+    
+    # create flow state
+    state = Transport(
+        model_type=model_type,
+        path_type=path_type,
+        loss_type=loss_type,
+        train_eps=train_eps,
+        sample_eps=sample_eps,
+    )
+    
+    return state
\ No newline at end of file
diff --git a/opensora/models/diffusion/utils/transport/integrators.py b/opensora/models/diffusion/utils/transport/integrators.py
new file mode 100644
index 000000000..adf7c7b4c
--- /dev/null
+++ b/opensora/models/diffusion/utils/transport/integrators.py
@@ -0,0 +1,117 @@
+import numpy as np
+import torch as th
+import torch.nn as nn
+from torchdiffeq import odeint
+from functools import partial
+from tqdm import tqdm
+
+class sde:
+    """SDE solver class"""
+    def __init__(
+        self, 
+        drift,
+        diffusion,
+        *,
+        t0,
+        t1,
+        num_steps,
+        sampler_type,
+    ):
+        assert t0 < t1, "SDE sampler has to be in forward time"
+
+        self.num_timesteps = num_steps
+        self.t = th.linspace(t0, t1, num_steps)
+        self.dt = self.t[1] - self.t[0]
+        self.drift = drift
+        self.diffusion = diffusion
+        self.sampler_type = sampler_type
+
+    def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs):
+        w_cur = th.randn(x.size()).to(x)
+        t = th.ones(x.size(0)).to(x) * t
+        dw = w_cur * th.sqrt(self.dt)
+        drift = self.drift(x, t, model, **model_kwargs)
+        diffusion = self.diffusion(x, t)
+        mean_x = x + drift * self.dt
+        x = mean_x + th.sqrt(2 * diffusion) * dw
+        return x, mean_x
+    
+    def __Heun_step(self, x, _, t, model, **model_kwargs):
+        w_cur = th.randn(x.size()).to(x)
+        dw = w_cur * th.sqrt(self.dt)
+        t_cur = th.ones(x.size(0)).to(x) * t
+        diffusion = self.diffusion(x, t_cur)
+        xhat = x + th.sqrt(2 * diffusion) * dw
+        K1 = self.drift(xhat, t_cur, model, **model_kwargs)
+        xp = xhat + self.dt * K1
+        K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs)
+        return xhat + 0.5 * self.dt * (K1 + K2), xhat # at last time point we do not perform the heun step
+
+    def __forward_fn(self):
+        """TODO: generalize here by adding all private functions ending with steps to it"""
+        sampler_dict = {
+            "Euler": self.__Euler_Maruyama_step,
+            "Heun": self.__Heun_step,
+        }
+
+        try:
+            sampler = sampler_dict[self.sampler_type]
+        except:
+            raise NotImplementedError("Smapler type not implemented.")
+    
+        return sampler
+
+    def sample(self, init, model, **model_kwargs):
+        """forward loop of sde"""
+        x = init
+        mean_x = init 
+        samples = []
+        sampler = self.__forward_fn()
+        for ti in self.t[:-1]:
+            with th.no_grad():
+                x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs)
+                samples.append(x)
+
+        return samples
+
+class ode:
+    """ODE solver class"""
+    def __init__(
+        self,
+        drift,
+        *,
+        t0,
+        t1,
+        sampler_type,
+        num_steps,
+        atol,
+        rtol,
+    ):
+        assert t0 < t1, "ODE sampler has to be in forward time"
+
+        self.drift = drift
+        self.t = th.linspace(t0, t1, num_steps)
+        self.atol = atol
+        self.rtol = rtol
+        self.sampler_type = sampler_type
+
+    def sample(self, x, model, **model_kwargs):
+        
+        device = x[0].device if isinstance(x, tuple) else x.device
+        def _fn(t, x):
+            t = th.ones(x[0].size(0)).to(device) * t if isinstance(x, tuple) else th.ones(x.size(0)).to(device) * t
+            model_output = self.drift(x, t, model, **model_kwargs)
+            return model_output
+
+        t = self.t.to(device)
+        atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol]
+        rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol]
+        samples = odeint(
+            _fn,
+            x,
+            t,
+            method=self.sampler_type,
+            atol=atol,
+            rtol=rtol
+        )
+        return samples
\ No newline at end of file
diff --git a/opensora/models/diffusion/utils/transport/path.py b/opensora/models/diffusion/utils/transport/path.py
new file mode 100644
index 000000000..156a7b0de
--- /dev/null
+++ b/opensora/models/diffusion/utils/transport/path.py
@@ -0,0 +1,192 @@
+import torch as th
+import numpy as np
+from functools import partial
+
+def expand_t_like_x(t, x):
+    """Function to reshape time t to broadcastable dimension of x
+    Args:
+      t: [batch_dim,], time vector
+      x: [batch_dim,...], data point
+    """
+    dims = [1] * (len(x.size()) - 1)
+    t = t.view(t.size(0), *dims)
+    return t
+
+
+#################### Coupling Plans ####################
+
+class ICPlan:
+    """Linear Coupling Plan"""
+    def __init__(self, sigma=0.0):
+        self.sigma = sigma
+
+    def compute_alpha_t(self, t):
+        """Compute the data coefficient along the path"""
+        return t, 1
+    
+    def compute_sigma_t(self, t):
+        """Compute the noise coefficient along the path"""
+        return 1 - t, -1
+    
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Compute the ratio between d_alpha and alpha"""
+        return 1 / t
+
+    def compute_drift(self, x, t):
+        """We always output sde according to score parametrization; """
+        t = expand_t_like_x(t, x)
+        alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        drift = alpha_ratio * x
+        diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
+
+        return -drift, diffusion
+
+    def compute_diffusion(self, x, t, form="constant", norm=1.0):
+        """Compute the diffusion term of the SDE
+        Args:
+          x: [batch_dim, ...], data point
+          t: [batch_dim,], time vector
+          form: str, form of the diffusion term
+          norm: float, norm of the diffusion term
+        """
+        t = expand_t_like_x(t, x)
+        choices = {
+            "constant": norm,
+            "SBDM": norm * self.compute_drift(x, t)[1],
+            "sigma": norm * self.compute_sigma_t(t)[0],
+            "linear": norm * (1 - t),
+            "decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2,
+            "inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2,
+        }
+
+        try:
+            diffusion = choices[form]
+        except KeyError:
+            raise NotImplementedError(f"Diffusion form {form} not implemented")
+        
+        return diffusion
+
+    def get_score_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to score
+        Args:
+            velocity: [batch_dim, ...] shaped tensor; velocity model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+        score = (reverse_alpha_ratio * velocity - mean) / var
+        return score
+    
+    def get_noise_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to denoiser
+        Args:
+            velocity: [batch_dim, ...] shaped tensor; velocity model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = reverse_alpha_ratio * d_sigma_t - sigma_t
+        noise = (reverse_alpha_ratio * velocity - mean) / var
+        return noise
+
+    def get_velocity_from_score(self, score, x, t):
+        """Wrapper function: transfrom score prediction model to velocity
+        Args:
+            score: [batch_dim, ...] shaped tensor; score model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        drift, var = self.compute_drift(x, t)
+        velocity = var * score - drift
+        return velocity
+
+    def compute_mu_t(self, t, x0, x1):
+        """Compute the mean of time-dependent density p_t"""
+        t = expand_t_like_x(t, x1)
+        alpha_t, _ = self.compute_alpha_t(t)
+        sigma_t, _ = self.compute_sigma_t(t)
+        return alpha_t * x1 + sigma_t * x0
+    
+    def compute_xt(self, t, x0, x1):
+        """Sample xt from time-dependent density p_t; rng is required"""
+        xt = self.compute_mu_t(t, x0, x1)
+        return xt
+    
+    def compute_ut(self, t, x0, x1, xt):
+        """Compute the vector field corresponding to p_t"""
+        t = expand_t_like_x(t, x1)
+        _, d_alpha_t = self.compute_alpha_t(t)
+        _, d_sigma_t = self.compute_sigma_t(t)
+        return d_alpha_t * x1 + d_sigma_t * x0
+    
+    def plan(self, t, x0, x1):
+        xt = self.compute_xt(t, x0, x1)
+        ut = self.compute_ut(t, x0, x1, xt)
+        return t, xt, ut
+    
+
+class VPCPlan(ICPlan):
+    """class for VP path flow matching"""
+
+    def __init__(self, sigma_min=0.1, sigma_max=20.0):
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.log_mean_coeff = lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min) - 0.5 * (1 - t) * self.sigma_min 
+        self.d_log_mean_coeff = lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min) + 0.5 * self.sigma_min
+
+
+    def compute_alpha_t(self, t):
+        """Compute coefficient of x1"""
+        alpha_t = self.log_mean_coeff(t)
+        alpha_t = th.exp(alpha_t)
+        d_alpha_t = alpha_t * self.d_log_mean_coeff(t)
+        return alpha_t, d_alpha_t
+    
+    def compute_sigma_t(self, t):
+        """Compute coefficient of x0"""
+        p_sigma_t = 2 * self.log_mean_coeff(t)
+        sigma_t = th.sqrt(1 - th.exp(p_sigma_t))
+        d_sigma_t = th.exp(p_sigma_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * sigma_t)
+        return sigma_t, d_sigma_t
+    
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
+        return self.d_log_mean_coeff(t)
+
+    def compute_drift(self, x, t):
+        """Compute the drift term of the SDE"""
+        t = expand_t_like_x(t, x)
+        beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min)
+        return -0.5 * beta_t * x, beta_t / 2
+    
+
+class GVPCPlan(ICPlan):
+    def __init__(self, sigma=0.0):
+        super().__init__(sigma)
+    
+    def compute_alpha_t(self, t):
+        """Compute coefficient of x1"""
+        alpha_t = th.sin(t * np.pi / 2)
+        d_alpha_t = np.pi / 2 * th.cos(t * np.pi / 2)
+        return alpha_t, d_alpha_t
+    
+    def compute_sigma_t(self, t):
+        """Compute coefficient of x0"""
+        sigma_t = th.cos(t * np.pi / 2)
+        d_sigma_t = -np.pi / 2 * th.sin(t * np.pi / 2)
+        return sigma_t, d_sigma_t
+    
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
+        return np.pi / (2 * th.tan(t * np.pi / 2))
\ No newline at end of file
diff --git a/opensora/models/diffusion/utils/transport/transport.py b/opensora/models/diffusion/utils/transport/transport.py
new file mode 100644
index 000000000..396c516cf
--- /dev/null
+++ b/opensora/models/diffusion/utils/transport/transport.py
@@ -0,0 +1,443 @@
+import torch as th
+import numpy as np
+import logging
+
+import enum
+
+from . import path
+from .utils import EasyDict, log_state, mean_flat
+from .integrators import ode, sde
+
+class ModelType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    NOISE = enum.auto()  # the model predicts epsilon
+    SCORE = enum.auto()  # the model predicts \nabla \log p(x)
+    VELOCITY = enum.auto()  # the model predicts v(x)
+
+class PathType(enum.Enum):
+    """
+    Which type of path to use.
+    """
+
+    LINEAR = enum.auto()
+    GVP = enum.auto()
+    VP = enum.auto()
+
+class WeightType(enum.Enum):
+    """
+    Which type of weighting to use.
+    """
+
+    NONE = enum.auto()
+    VELOCITY = enum.auto()
+    LIKELIHOOD = enum.auto()
+
+
+class Transport:
+
+    def __init__(
+        self,
+        *,
+        model_type,
+        path_type,
+        loss_type,
+        train_eps,
+        sample_eps,
+    ):
+        path_options = {
+            PathType.LINEAR: path.ICPlan,
+            PathType.GVP: path.GVPCPlan,
+            PathType.VP: path.VPCPlan,
+        }
+
+        self.loss_type = loss_type
+        self.model_type = model_type
+        self.path_sampler = path_options[path_type]()
+        self.train_eps = train_eps
+        self.sample_eps = sample_eps
+
+    def prior_logp(self, z):
+        '''
+            Standard multivariate normal prior
+            Assume z is batched
+        '''
+        shape = th.tensor(z.size())
+        N = th.prod(shape[1:])
+        _fn = lambda x: -N / 2. * np.log(2 * np.pi) - th.sum(x ** 2) / 2.
+        return th.vmap(_fn)(z)
+    
+
+    def check_interval(
+        self, 
+        train_eps, 
+        sample_eps, 
+        *, 
+        diffusion_form="SBDM",
+        sde=False, 
+        reverse=False, 
+        eval=False,
+        last_step_size=0.0,
+    ):
+        t0 = 0
+        t1 = 1
+        eps = train_eps if not eval else sample_eps
+        if (type(self.path_sampler) in [path.VPCPlan]):
+
+            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
+
+        elif (type(self.path_sampler) in [path.ICPlan, path.GVPCPlan]) \
+            and (self.model_type != ModelType.VELOCITY or sde): # avoid numerical issue by taking a first semi-implicit step
+
+            t0 = eps if (diffusion_form == "SBDM" and sde) or self.model_type != ModelType.VELOCITY else 0
+            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
+        
+        if reverse:
+            t0, t1 = 1 - t0, 1 - t1
+
+        return t0, t1
+
+
+    def sample(self, x1):
+        """Sampling x0 & t based on shape of x1 (if needed)
+          Args:
+            x1 - data point; [batch, *dim]
+        """
+        
+        x0 = th.randn_like(x1)
+        t0, t1 = self.check_interval(self.train_eps, self.sample_eps)
+        t = th.rand((x1.shape[0],)) * (t1 - t0) + t0
+        t = t.to(x1)
+        return t, x0, x1
+    
+
+    def training_losses(
+        self, 
+        model,  
+        x1, 
+        model_kwargs=None
+    ):
+        """Loss for training the score model
+        Args:
+        - model: backbone model; could be score, noise, or velocity
+        - x1: datapoint
+        - model_kwargs: additional arguments for the model
+        """
+        if model_kwargs == None:
+            model_kwargs = {}
+        
+        t, x0, x1 = self.sample(x1)
+        t, xt, ut = self.path_sampler.plan(t, x0, x1)
+        model_output = model(xt, t, **model_kwargs)
+        B, *_, C = xt.shape
+        assert model_output.size() == (B, *xt.size()[1:-1], C)
+
+        terms = {}
+        terms['pred'] = model_output
+        if self.model_type == ModelType.VELOCITY:
+            terms['loss'] = mean_flat(((model_output - ut) ** 2))
+        else: 
+            _, drift_var = self.path_sampler.compute_drift(xt, t)
+            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, xt))
+            if self.loss_type in [WeightType.VELOCITY]:
+                weight = (drift_var / sigma_t) ** 2
+            elif self.loss_type in [WeightType.LIKELIHOOD]:
+                weight = drift_var / (sigma_t ** 2)
+            elif self.loss_type in [WeightType.NONE]:
+                weight = 1
+            else:
+                raise NotImplementedError()
+            
+            if self.model_type == ModelType.NOISE:
+                terms['loss'] = mean_flat(weight * ((model_output - x0) ** 2))
+            else:
+                terms['loss'] = mean_flat(weight * ((model_output * sigma_t + x0) ** 2))
+                
+        return terms
+    
+
+    def get_drift(
+        self
+    ):
+        """member function for obtaining the drift of the probability flow ODE"""
+        def score_ode(x, t, model, **model_kwargs):
+            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
+            model_output = model(x, t, **model_kwargs)
+            return (-drift_mean + drift_var * model_output) # by change of variable
+        
+        def noise_ode(x, t, model, **model_kwargs):
+            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
+            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))
+            model_output = model(x, t, **model_kwargs)
+            score = model_output / -sigma_t
+            return (-drift_mean + drift_var * score)
+        
+        def velocity_ode(x, t, model, **model_kwargs):
+            model_output = model(x, t, **model_kwargs)
+            return model_output
+
+        if self.model_type == ModelType.NOISE:
+            drift_fn = noise_ode
+        elif self.model_type == ModelType.SCORE:
+            drift_fn = score_ode
+        else:
+            drift_fn = velocity_ode
+        
+        def body_fn(x, t, model, **model_kwargs):
+            model_output = drift_fn(x, t, model, **model_kwargs)
+            assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape"
+            return model_output
+
+        return body_fn
+    
+
+    def get_score(
+        self,
+    ):
+        """member function for obtaining score of 
+            x_t = alpha_t * x + sigma_t * eps"""
+        if self.model_type == ModelType.NOISE:
+            score_fn = lambda x, t, model, **kwargs: model(x, t, **kwargs) / -self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))[0]
+        elif self.model_type == ModelType.SCORE:
+            score_fn = lambda x, t, model, **kwagrs: model(x, t, **kwagrs)
+        elif self.model_type == ModelType.VELOCITY:
+            score_fn = lambda x, t, model, **kwargs: self.path_sampler.get_score_from_velocity(model(x, t, **kwargs), x, t)
+        else:
+            raise NotImplementedError()
+        
+        return score_fn
+
+
+class Sampler:
+    """Sampler class for the transport model"""
+    def __init__(
+        self,
+        transport,
+    ):
+        """Constructor for a general sampler; supporting different sampling methods
+        Args:
+        - transport: an tranport object specify model prediction & interpolant type
+        """
+        
+        self.transport = transport
+        self.drift = self.transport.get_drift()
+        self.score = self.transport.get_score()
+    
+    def __get_sde_diffusion_and_drift(
+        self,
+        *,
+        diffusion_form="SBDM",
+        diffusion_norm=1.0,
+    ):
+
+        def diffusion_fn(x, t):
+            diffusion = self.transport.path_sampler.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm)
+            return diffusion
+        
+        sde_drift = \
+            lambda x, t, model, **kwargs: \
+                self.drift(x, t, model, **kwargs) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs)
+    
+        sde_diffusion = diffusion_fn
+
+        return sde_drift, sde_diffusion
+    
+    def __get_last_step(
+        self,
+        sde_drift,
+        *,
+        last_step,
+        last_step_size,
+    ):
+        """Get the last step function of the SDE solver"""
+    
+        if last_step is None:
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x
+        elif last_step == "Mean":
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x + sde_drift(x, t, model, **model_kwargs) * last_step_size
+        elif last_step == "Tweedie":
+            alpha = self.transport.path_sampler.compute_alpha_t # simple aliasing; the original name was too long
+            sigma = self.transport.path_sampler.compute_sigma_t
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x / alpha(t)[0][0] + (sigma(t)[0][0] ** 2) / alpha(t)[0][0] * self.score(x, t, model, **model_kwargs)
+        elif last_step == "Euler":
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x + self.drift(x, t, model, **model_kwargs) * last_step_size
+        else:
+            raise NotImplementedError()
+
+        return last_step_fn
+
+    def sample_sde(
+        self,
+        *,
+        sampling_method="Euler",
+        diffusion_form="SBDM",
+        diffusion_norm=1.0,
+        last_step="Mean",
+        last_step_size=0.04,
+        num_steps=250,
+    ):
+        """returns a sampling function with given SDE settings
+        Args:
+        - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama
+        - diffusion_form: function form of diffusion coefficient; default to be matching SBDM
+        - diffusion_norm: function magnitude of diffusion coefficient; default to 1
+        - last_step: type of the last step; default to identity
+        - last_step_size: size of the last step; default to match the stride of 250 steps over [0,1]
+        - num_steps: total integration step of SDE
+        """
+
+        if last_step is None:
+            last_step_size = 0.0
+
+        sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift(
+            diffusion_form=diffusion_form,
+            diffusion_norm=diffusion_norm,
+        )
+
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            diffusion_form=diffusion_form,
+            sde=True,
+            eval=True,
+            reverse=False,
+            last_step_size=last_step_size,
+        )
+
+        _sde = sde(
+            sde_drift,
+            sde_diffusion,
+            t0=t0,
+            t1=t1,
+            num_steps=num_steps,
+            sampler_type=sampling_method
+        )
+
+        last_step_fn = self.__get_last_step(sde_drift, last_step=last_step, last_step_size=last_step_size)
+            
+
+        def _sample(init, model, **model_kwargs):
+            xs = _sde.sample(init, model, **model_kwargs)
+            ts = th.ones(init.size(0), device=init.device) * t1
+            x = last_step_fn(xs[-1], ts, model, **model_kwargs)
+            xs.append(x)
+
+            assert len(xs) == num_steps, "Samples does not match the number of steps"
+
+            return xs
+
+        return _sample
+    
+    def sample_ode(
+        self,
+        *,
+        sampling_method="dopri5",
+        num_steps=50,
+        atol=1e-6,
+        rtol=1e-3,
+        reverse=False,
+    ):
+        """returns a sampling function with given ODE settings
+        Args:
+        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
+        - num_steps: 
+            - fixed solver (Euler, Heun): the actual number of integration steps performed
+            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
+        - atol: absolute error tolerance for the solver
+        - rtol: relative error tolerance for the solver
+        - reverse: whether solving the ODE in reverse (data to noise); default to False
+        """
+        if reverse:
+            drift = lambda x, t, model, **kwargs: self.drift(x, th.ones_like(t) * (1 - t), model, **kwargs)
+        else:
+            drift = self.drift
+
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            sde=False,
+            eval=True,
+            reverse=reverse,
+            last_step_size=0.0,
+        )
+
+        _ode = ode(
+            drift=drift,
+            t0=t0,
+            t1=t1,
+            sampler_type=sampling_method,
+            num_steps=num_steps,
+            atol=atol,
+            rtol=rtol,
+        )
+        
+        return _ode.sample
+
+    def sample_ode_likelihood(
+        self,
+        *,
+        sampling_method="dopri5",
+        num_steps=50,
+        atol=1e-6,
+        rtol=1e-3,
+    ):
+        
+        """returns a sampling function for calculating likelihood with given ODE settings
+        Args:
+        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
+        - num_steps: 
+            - fixed solver (Euler, Heun): the actual number of integration steps performed
+            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
+        - atol: absolute error tolerance for the solver
+        - rtol: relative error tolerance for the solver
+        """
+        def _likelihood_drift(x, t, model, **model_kwargs):
+            x, _ = x
+            eps = th.randint(2, x.size(), dtype=th.float, device=x.device) * 2 - 1
+            t = th.ones_like(t) * (1 - t)
+            with th.enable_grad():
+                x.requires_grad = True
+                grad = th.autograd.grad(th.sum(self.drift(x, t, model, **model_kwargs) * eps), x)[0]
+                logp_grad = th.sum(grad * eps, dim=tuple(range(1, len(x.size()))))
+                drift = self.drift(x, t, model, **model_kwargs)
+            return (-drift, logp_grad)
+        
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            sde=False,
+            eval=True,
+            reverse=False,
+            last_step_size=0.0,
+        )
+
+        _ode = ode(
+            drift=_likelihood_drift,
+            t0=t0,
+            t1=t1,
+            sampler_type=sampling_method,
+            num_steps=num_steps,
+            atol=atol,
+            rtol=rtol,
+        )
+
+        def _sample_fn(x, model, **model_kwargs):
+            init_logp = th.zeros(x.size(0)).to(x)
+            input = (x, init_logp)
+            drift, delta_logp = _ode.sample(input, model, **model_kwargs)
+            drift, delta_logp = drift[-1], delta_logp[-1]
+            prior_logp = self.transport.prior_logp(drift)
+            logp = prior_logp - delta_logp
+            return logp, drift
+
+        return _sample_fn
\ No newline at end of file
diff --git a/opensora/models/diffusion/utils/transport/utils.py b/opensora/models/diffusion/utils/transport/utils.py
new file mode 100644
index 000000000..446460355
--- /dev/null
+++ b/opensora/models/diffusion/utils/transport/utils.py
@@ -0,0 +1,29 @@
+import torch as th
+
+class EasyDict:
+
+    def __init__(self, sub_dict):
+        for k, v in sub_dict.items():
+            setattr(self, k, v)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+def mean_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return th.mean(x, dim=list(range(1, len(x.size()))))
+
+def log_state(state):
+    result = []
+    
+    sorted_state = dict(sorted(state.items()))
+    for key, value in sorted_state.items():
+        # Check if the value is an instance of a class
+        if "<object" in str(value) or "object at" in str(value):
+            result.append(f"{key}: [{value.__class__.__name__}]")
+        else:
+            result.append(f"{key}: {value}")
+    
+    return '\n'.join(result)
\ No newline at end of file
diff --git a/opensora/models/frame_interpolation/cfgs/AMT-G.yaml b/opensora/models/frame_interpolation/cfgs/AMT-G.yaml
new file mode 100644
index 000000000..d259d4fe9
--- /dev/null
+++ b/opensora/models/frame_interpolation/cfgs/AMT-G.yaml
@@ -0,0 +1,9 @@
+
+seed: 2023
+
+network:
+  name: networks.AMT-G.Model
+  params:
+    corr_radius: 3
+    corr_lvls: 4
+    num_flows: 5
\ No newline at end of file
diff --git a/opensora/models/frame_interpolation/interpolation.py b/opensora/models/frame_interpolation/interpolation.py
new file mode 100644
index 000000000..0c60bed8a
--- /dev/null
+++ b/opensora/models/frame_interpolation/interpolation.py
@@ -0,0 +1,197 @@
+# this script is modified from https://github.com/MCG-NKU/AMT/blob/main/demos/demo_2x.py
+from json import load
+import os
+import cv2
+import sys
+import glob
+import torch
+import argparse
+import numpy as np
+import os.path as osp
+from warnings import warn
+from omegaconf import OmegaConf
+from torchvision.utils import make_grid
+sys.path.append('.')
+from utils.utils import (
+    read, write,
+    img2tensor, tensor2img,
+    check_dim_and_resize
+    )
+from utils.build_utils import build_from_cfg
+from utils.utils import InputPadder
+
+
+AMT_G = {
+    'name': 'networks.AMT-G.Model',
+    'params':{
+        'corr_radius': 3,
+        'corr_lvls': 4,
+        'num_flows': 5,
+    }
+}
+
+
+
+def init(device="cuda"):
+
+    '''
+        initialize the device and the anchor resolution.
+    '''
+
+    if device == 'cuda':
+        anchor_resolution = 1024 * 512
+        anchor_memory = 1500 * 1024**2
+        anchor_memory_bias = 2500 * 1024**2
+        vram_avail = torch.cuda.get_device_properties(device).total_memory
+        print("VRAM available: {:.1f} MB".format(vram_avail / 1024 ** 2))
+    else:
+        # Do not resize in cpu mode
+        anchor_resolution = 8192*8192
+        anchor_memory = 1
+        anchor_memory_bias = 0
+        vram_avail = 1
+    
+    return anchor_resolution, anchor_memory, anchor_memory_bias, vram_avail
+
+def get_input_video_from_path(input_path, device="cuda"):
+
+    '''
+        Get the input video from the input_path.
+
+        params:
+            input_path: str, the path of the input video.
+            devices: str, the device to run the model.
+        returns:
+            inputs: list, the list of the input frames.
+            scale: float, the scale of the input frames.
+            padder: InputPadder, the padder to pad the input frames.
+    '''
+
+    anchor_resolution, anchor_memory, anchor_memory_bias, vram_avail = init(device)
+
+    if osp.splitext(input_path)[-1] in ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', 
+                                        '.webm', '.MP4', '.AVI', '.MOV', '.MKV', '.FLV', 
+                                        '.WMV', '.WEBM']:
+
+        vcap = cv2.VideoCapture(input_path)
+
+        inputs = []
+        w = int(vcap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        h = int(vcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        scale = anchor_resolution / (h * w) * np.sqrt((vram_avail - anchor_memory_bias) / anchor_memory)
+        scale = 1 if scale > 1 else scale
+        scale = 1 / np.floor(1 / np.sqrt(scale) * 16) * 16
+        if scale < 1:
+            print(f"Due to the limited VRAM, the video will be scaled by {scale:.2f}")
+        padding = int(16 / scale)
+        padder = InputPadder((h, w), padding)
+        while True:
+            ret, frame = vcap.read()
+            if ret is False:
+                break
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_t = img2tensor(frame).to(device)
+            frame_t = padder.pad(frame_t)
+            inputs.append(frame_t)
+        print(f'Loading the [video] from {input_path}, the number of frames [{len(inputs)}]')
+    else:
+        raise TypeError("Input should be a video.")
+    
+    return inputs, scale, padder
+
+
+def load_model(ckpt_path, device="cuda"):
+
+    '''
+        load the frame interpolation model.
+    '''
+    network_cfg = AMT_G
+    network_name = network_cfg['name']
+    print(f'Loading [{network_name}] from [{ckpt_path}]...')
+    model = build_from_cfg(network_cfg)
+    ckpt = torch.load(ckpt_path)
+    model.load_state_dict(ckpt['state_dict'])
+    model = model.to(device)
+    model.eval()
+    return model
+
+def interpolater(model, inputs, scale, padder, iters=1):
+
+    '''
+        interpolating with the interpolation model.
+
+        params:
+            model: nn.Module, the frame interpolation model.
+            inputs: list, the list of the input frames.
+            scale: float, the scale of the input frames.
+            iters: int, the number of iterations of interpolation. The final frames model generating is 2 ** iters * (m - 1) + 1 and m is input frames. 
+        returns:
+            outputs: list, the list of the output frames.
+    '''
+
+    print(f'Start frame interpolation:')
+    embt = torch.tensor(1/2).float().view(1, 1, 1, 1).to(device)
+
+    for i in range(iters):
+        print(f'Iter {i+1}. input_frames={len(inputs)} output_frames={2*len(inputs)-1}')
+        outputs = [inputs[0]]
+        for in_0, in_1 in zip(inputs[:-1], inputs[1:]):
+            in_0 = in_0.to(device)
+            in_1 = in_1.to(device)
+            with torch.no_grad():
+                imgt_pred = model(in_0, in_1, embt, scale_factor=scale, eval=True)['imgt_pred']
+            outputs += [imgt_pred.cpu(), in_1.cpu()]
+        inputs = outputs
+
+    outputs = padder.unpad(*outputs)
+
+    return outputs
+
+def write(outputs, input_path, output_path, frame_rate=30):
+    '''
+        write results to the output_path.
+    '''
+
+    if osp.exists(output_path) is False:
+        os.makedirs(output_path)
+
+    
+    size = outputs[0].shape[2:][::-1]
+
+    _, file_name_with_extension = os.path.split(input_path)
+    file_name, _ = os.path.splitext(file_name_with_extension)
+
+    save_video_path = f'{output_path}/output_{file_name}.mp4'
+    writer = cv2.VideoWriter(save_video_path, cv2.VideoWriter_fourcc(*"mp4v"), 
+                        frame_rate, size)
+
+    for i, imgt_pred in enumerate(outputs):
+        imgt_pred = tensor2img(imgt_pred)
+        imgt_pred = cv2.cvtColor(imgt_pred, cv2.COLOR_RGB2BGR)
+        writer.write(imgt_pred)        
+    print(f"Demo video is saved to [{save_video_path}]")
+
+    writer.release()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ckpt', type=str, default='amt-g.pth', help="The pretrained model.") 
+    parser.add_argument('--niters', type=int, default=1, help="Iter of Interpolation. The number of frames will be double after per iter.") 
+    parser.add_argument('--input', default="test.mp4", help="Input video.") 
+    parser.add_argument('--output_path', type=str, default='results', help="Output path.") 
+    parser.add_argument('--frame_rate', type=int, default=30, help="Frames rate of the output video.")
+
+    args = parser.parse_args()
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    ckpt_path = args.ckpt
+    input_path = args.input
+    output_path = args.output_path
+    iters = int(args.niters)
+    frame_rate = int(args.frame_rate)
+
+    inputs, scale, padder = get_input_video_from_path(input_path, device)
+    model = load_model(ckpt_path, device)
+    outputs = interpolater(model, inputs, scale, padder, iters)
+    write(outputs, input_path, output_path, frame_rate)
diff --git a/opensora/models/frame_interpolation/networks/AMT-G.py b/opensora/models/frame_interpolation/networks/AMT-G.py
new file mode 100644
index 000000000..a24cb1a37
--- /dev/null
+++ b/opensora/models/frame_interpolation/networks/AMT-G.py
@@ -0,0 +1,172 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from networks.blocks.raft import (
+    coords_grid,
+    BasicUpdateBlock, BidirCorrBlock
+)
+from networks.blocks.feat_enc import (
+    LargeEncoder
+)
+from networks.blocks.ifrnet import (
+    resize,
+    Encoder,
+    InitDecoder,
+    IntermediateDecoder
+)
+from networks.blocks.multi_flow import (
+    multi_flow_combine,
+    MultiFlowDecoder
+)
+
+
+class Model(nn.Module):
+    def __init__(self, 
+                 corr_radius=3, 
+                 corr_lvls=4, 
+                 num_flows=5, 
+                 channels=[84, 96, 112, 128], 
+                 skip_channels=84):
+        super(Model, self).__init__()
+        self.radius = corr_radius
+        self.corr_levels = corr_lvls
+        self.num_flows = num_flows
+
+        self.feat_encoder = LargeEncoder(output_dim=128, norm_fn='instance', dropout=0.)
+        self.encoder = Encoder(channels, large=True)
+        self.decoder4 = InitDecoder(channels[3], channels[2], skip_channels)
+        self.decoder3 = IntermediateDecoder(channels[2], channels[1], skip_channels)
+        self.decoder2 = IntermediateDecoder(channels[1], channels[0], skip_channels)
+        self.decoder1 = MultiFlowDecoder(channels[0], skip_channels, num_flows)
+
+        self.update4 = self._get_updateblock(112, None)
+        self.update3_low = self._get_updateblock(96, 2.0)
+        self.update2_low = self._get_updateblock(84, 4.0)
+        
+        self.update3_high = self._get_updateblock(96, None)
+        self.update2_high = self._get_updateblock(84, None)
+        
+        self.comb_block = nn.Sequential(
+            nn.Conv2d(3*self.num_flows, 6*self.num_flows, 7, 1, 3),
+            nn.PReLU(6*self.num_flows),
+            nn.Conv2d(6*self.num_flows, 3, 7, 1, 3),
+        )
+
+    def _get_updateblock(self, cdim, scale_factor=None):
+        return BasicUpdateBlock(cdim=cdim, hidden_dim=192, flow_dim=64, 
+                                corr_dim=256, corr_dim2=192, fc_dim=188, 
+                                scale_factor=scale_factor, corr_levels=self.corr_levels, 
+                                radius=self.radius)
+
+    def _corr_scale_lookup(self, corr_fn, coord, flow0, flow1, embt, downsample=1):
+        # convert t -> 0 to 0 -> 1 | convert t -> 1 to 1 -> 0
+        # based on linear assumption
+        t1_scale = 1. / embt
+        t0_scale = 1. / (1. - embt)
+        if downsample != 1:
+            inv = 1 / downsample
+            flow0 = inv * resize(flow0, scale_factor=inv)
+            flow1 = inv * resize(flow1, scale_factor=inv)
+            
+        corr0, corr1 = corr_fn(coord + flow1 * t1_scale, coord + flow0 * t0_scale) 
+        corr = torch.cat([corr0, corr1], dim=1)
+        flow = torch.cat([flow0, flow1], dim=1)
+        return corr, flow
+    
+    def forward(self, img0, img1, embt, scale_factor=1.0, eval=False, **kwargs):
+        mean_ = torch.cat([img0, img1], 2).mean(1, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True)
+        img0 = img0 - mean_
+        img1 = img1 - mean_
+        img0_ = resize(img0, scale_factor) if scale_factor != 1.0 else img0
+        img1_ = resize(img1, scale_factor) if scale_factor != 1.0 else img1
+        b, _, h, w = img0_.shape
+        coord = coords_grid(b, h // 8, w // 8, img0.device)
+        
+        fmap0, fmap1 = self.feat_encoder([img0_, img1_]) # [1, 128, H//8, W//8]
+        corr_fn = BidirCorrBlock(fmap0, fmap1, radius=self.radius, num_levels=self.corr_levels)
+
+        # f0_1: [1, c0, H//2, W//2] | f0_2: [1, c1, H//4, W//4]
+        # f0_3: [1, c2, H//8, W//8] | f0_4: [1, c3, H//16, W//16]
+        f0_1, f0_2, f0_3, f0_4 = self.encoder(img0_)
+        f1_1, f1_2, f1_3, f1_4 = self.encoder(img1_)
+
+        ######################################### the 4th decoder #########################################
+        up_flow0_4, up_flow1_4, ft_3_ = self.decoder4(f0_4, f1_4, embt)
+        corr_4, flow_4 = self._corr_scale_lookup(corr_fn, coord, 
+                                                 up_flow0_4, up_flow1_4, 
+                                                 embt, downsample=1)
+
+        # residue update with lookup corr
+        delta_ft_3_, delta_flow_4 = self.update4(ft_3_, flow_4, corr_4)
+        delta_flow0_4, delta_flow1_4 = torch.chunk(delta_flow_4, 2, 1)
+        up_flow0_4 = up_flow0_4 + delta_flow0_4
+        up_flow1_4 = up_flow1_4 + delta_flow1_4
+        ft_3_ = ft_3_ + delta_ft_3_
+
+        ######################################### the 3rd decoder #########################################
+        up_flow0_3, up_flow1_3, ft_2_ = self.decoder3(ft_3_, f0_3, f1_3, up_flow0_4, up_flow1_4)
+        corr_3, flow_3 = self._corr_scale_lookup(corr_fn, 
+                                                 coord, up_flow0_3, up_flow1_3, 
+                                                 embt, downsample=2)
+
+        # residue update with lookup corr
+        delta_ft_2_, delta_flow_3 = self.update3_low(ft_2_, flow_3, corr_3)
+        delta_flow0_3, delta_flow1_3 = torch.chunk(delta_flow_3, 2, 1)
+        up_flow0_3 = up_flow0_3 + delta_flow0_3
+        up_flow1_3 = up_flow1_3 + delta_flow1_3
+        ft_2_ = ft_2_ + delta_ft_2_
+        
+        # residue update with lookup corr (hr)
+        corr_3 = resize(corr_3, scale_factor=2.0)
+        up_flow_3 = torch.cat([up_flow0_3, up_flow1_3], dim=1)
+        delta_ft_2_, delta_up_flow_3 = self.update3_high(ft_2_, up_flow_3, corr_3)
+        ft_2_ += delta_ft_2_
+        up_flow0_3 += delta_up_flow_3[:, 0:2]
+        up_flow1_3 += delta_up_flow_3[:, 2:4]
+        
+        ######################################### the 2nd decoder #########################################
+        up_flow0_2, up_flow1_2, ft_1_  = self.decoder2(ft_2_, f0_2, f1_2, up_flow0_3, up_flow1_3)
+        corr_2, flow_2 = self._corr_scale_lookup(corr_fn, 
+                                                 coord, up_flow0_2, up_flow1_2, 
+                                                 embt, downsample=4)
+        
+        # residue update with lookup corr
+        delta_ft_1_, delta_flow_2 = self.update2_low(ft_1_, flow_2, corr_2)
+        delta_flow0_2, delta_flow1_2 = torch.chunk(delta_flow_2, 2, 1)
+        up_flow0_2 = up_flow0_2 + delta_flow0_2
+        up_flow1_2 = up_flow1_2 + delta_flow1_2
+        ft_1_ = ft_1_ + delta_ft_1_
+        
+        # residue update with lookup corr (hr)
+        corr_2 = resize(corr_2, scale_factor=4.0)
+        up_flow_2 = torch.cat([up_flow0_2, up_flow1_2], dim=1)
+        delta_ft_1_, delta_up_flow_2 = self.update2_high(ft_1_, up_flow_2, corr_2)
+        ft_1_ += delta_ft_1_
+        up_flow0_2 += delta_up_flow_2[:, 0:2]
+        up_flow1_2 += delta_up_flow_2[:, 2:4]
+        
+        ######################################### the 1st decoder #########################################
+        up_flow0_1, up_flow1_1, mask, img_res = self.decoder1(ft_1_, f0_1, f1_1, up_flow0_2, up_flow1_2)
+        
+        if scale_factor != 1.0: 
+            up_flow0_1 = resize(up_flow0_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor)
+            up_flow1_1 = resize(up_flow1_1, scale_factor=(1.0/scale_factor)) * (1.0/scale_factor)
+            mask = resize(mask, scale_factor=(1.0/scale_factor))
+            img_res = resize(img_res, scale_factor=(1.0/scale_factor))
+
+        # Merge multiple predictions 
+        imgt_pred = multi_flow_combine(self.comb_block, img0, img1, up_flow0_1, up_flow1_1, 
+                                                                        mask, img_res, mean_)
+        imgt_pred = torch.clamp(imgt_pred, 0, 1)
+
+        if eval:
+            return  { 'imgt_pred': imgt_pred, }
+        else:
+            up_flow0_1 = up_flow0_1.reshape(b, self.num_flows, 2, h, w)
+            up_flow1_1 = up_flow1_1.reshape(b, self.num_flows, 2, h, w)
+            return {
+                'imgt_pred': imgt_pred,
+                'flow0_pred': [up_flow0_1, up_flow0_2, up_flow0_3, up_flow0_4],
+                'flow1_pred': [up_flow1_1, up_flow1_2, up_flow1_3, up_flow1_4],
+                'ft_pred': [ft_1_, ft_2_, ft_3_],
+            }
diff --git a/opensora/models/frame_interpolation/networks/__init__.py b/opensora/models/frame_interpolation/networks/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/models/frame_interpolation/networks/blocks/__init__.py b/opensora/models/frame_interpolation/networks/blocks/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/models/frame_interpolation/networks/blocks/feat_enc.py b/opensora/models/frame_interpolation/networks/blocks/feat_enc.py
new file mode 100644
index 000000000..3805bd315
--- /dev/null
+++ b/opensora/models/frame_interpolation/networks/blocks/feat_enc.py
@@ -0,0 +1,343 @@
+import torch
+import torch.nn as nn
+
+
+class BottleneckBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(BottleneckBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
+        self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
+        self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
+            self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes//4)
+            self.norm2 = nn.BatchNorm2d(planes//4)
+            self.norm3 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm4 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes//4)
+            self.norm2 = nn.InstanceNorm2d(planes//4)
+            self.norm3 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm4 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            self.norm3 = nn.Sequential()
+            if not stride == 1:
+                self.norm4 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
+
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+        y = self.relu(self.norm3(self.conv3(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(ResidualBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+class SmallEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
+        super(SmallEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(32)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(32)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 32
+        self.layer1 = self._make_layer(32,  stride=1)
+        self.layer2 = self._make_layer(64, stride=2)
+        self.layer3 = self._make_layer(96, stride=2)
+
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+        
+        self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+    
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x):
+
+        # if input is list, combine batch dimension
+        is_list = isinstance(x, tuple) or isinstance(x, list)
+        if is_list:
+            batch_dim = x[0].shape[0]
+            x = torch.cat(x, dim=0)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.conv2(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+
+        if is_list:
+            x = torch.split(x, [batch_dim, batch_dim], dim=0)
+
+        return x
+
+class BasicEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
+        super(BasicEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(64)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(64)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 64
+        self.layer1 = self._make_layer(64,  stride=1)
+        self.layer2 = self._make_layer(72, stride=2)
+        self.layer3 = self._make_layer(128, stride=2)
+
+        # output convolution
+        self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
+
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x):
+
+        # if input is list, combine batch dimension
+        is_list = isinstance(x, tuple) or isinstance(x, list)
+        if is_list:
+            batch_dim = x[0].shape[0]
+            x = torch.cat(x, dim=0)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        x = self.conv2(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+
+        if is_list:
+            x = torch.split(x, [batch_dim, batch_dim], dim=0)
+
+        return x
+
+class LargeEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
+        super(LargeEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(64)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(64)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 64
+        self.layer1 = self._make_layer(64, stride=1)
+        self.layer2 = self._make_layer(112, stride=2)
+        self.layer3 = self._make_layer(160, stride=2)
+        self.layer3_2 = self._make_layer(160, stride=1)
+
+        # output convolution
+        self.conv2 = nn.Conv2d(self.in_planes, output_dim, kernel_size=1)
+
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x):
+
+        # if input is list, combine batch dimension
+        is_list = isinstance(x, tuple) or isinstance(x, list)
+        if is_list:
+            batch_dim = x[0].shape[0]
+            x = torch.cat(x, dim=0)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer3_2(x)
+
+        x = self.conv2(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+
+        if is_list:
+            x = torch.split(x, [batch_dim, batch_dim], dim=0)
+
+        return x
diff --git a/opensora/models/frame_interpolation/networks/blocks/ifrnet.py b/opensora/models/frame_interpolation/networks/blocks/ifrnet.py
new file mode 100644
index 000000000..586ae6103
--- /dev/null
+++ b/opensora/models/frame_interpolation/networks/blocks/ifrnet.py
@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from utils.flow_utils import warp
+
+
+def resize(x, scale_factor):
+    return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False)
+
+def convrelu(in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias=bias), 
+        nn.PReLU(out_channels)
+    )
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channels, side_channels, bias=True):
+        super(ResBlock, self).__init__()
+        self.side_channels = side_channels
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias), 
+            nn.PReLU(in_channels)
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(side_channels, side_channels, kernel_size=3, stride=1, padding=1, bias=bias), 
+            nn.PReLU(side_channels)
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias), 
+            nn.PReLU(in_channels)
+        )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(side_channels, side_channels, kernel_size=3, stride=1, padding=1, bias=bias), 
+            nn.PReLU(side_channels)
+        )
+        self.conv5 = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=bias)
+        self.prelu = nn.PReLU(in_channels)
+
+    def forward(self, x):
+        out = self.conv1(x)
+
+        res_feat = out[:, :-self.side_channels, ...]
+        side_feat = out[:, -self.side_channels:, :, :]
+        side_feat = self.conv2(side_feat)
+        out = self.conv3(torch.cat([res_feat, side_feat], 1))
+
+        res_feat = out[:, :-self.side_channels, ...]
+        side_feat = out[:, -self.side_channels:, :, :]
+        side_feat = self.conv4(side_feat)
+        out = self.conv5(torch.cat([res_feat, side_feat], 1))
+
+        out = self.prelu(x + out)
+        return out
+    
+class Encoder(nn.Module):
+    def __init__(self, channels, large=False):
+        super(Encoder, self).__init__()
+        self.channels = channels        
+        prev_ch = 3
+        for idx, ch in enumerate(channels, 1):
+            k = 7 if large and idx == 1 else 3
+            p = 3 if k ==7 else 1
+            self.register_module(f'pyramid{idx}', 
+            nn.Sequential(
+                convrelu(prev_ch, ch, k, 2, p),
+                convrelu(ch, ch, 3, 1, 1)
+            ))
+            prev_ch = ch
+                
+    def forward(self, in_x):
+        fs = []
+        for idx in range(len(self.channels)):
+            out_x = getattr(self, f'pyramid{idx+1}')(in_x)
+            fs.append(out_x)
+            in_x = out_x
+        return fs
+    
+class InitDecoder(nn.Module):
+    def __init__(self, in_ch, out_ch, skip_ch) -> None:
+        super().__init__()
+        self.convblock = nn.Sequential(
+            convrelu(in_ch*2+1, in_ch*2), 
+            ResBlock(in_ch*2, skip_ch), 
+            nn.ConvTranspose2d(in_ch*2, out_ch+4, 4, 2, 1, bias=True)
+        )
+    def forward(self, f0, f1, embt):
+        h, w = f0.shape[2:]
+        embt = embt.repeat(1, 1, h, w)
+        out = self.convblock(torch.cat([f0, f1, embt], 1))
+        flow0, flow1 = torch.chunk(out[:, :4, ...], 2, 1)
+        ft_ = out[:, 4:, ...]
+        return flow0, flow1, ft_
+    
+class IntermediateDecoder(nn.Module):
+    def __init__(self, in_ch, out_ch, skip_ch) -> None:
+        super().__init__()
+        self.convblock = nn.Sequential(
+            convrelu(in_ch*3+4, in_ch*3), 
+            ResBlock(in_ch*3, skip_ch), 
+            nn.ConvTranspose2d(in_ch*3, out_ch+4, 4, 2, 1, bias=True)
+        )
+    def forward(self, ft_, f0, f1, flow0_in, flow1_in):
+        f0_warp = warp(f0, flow0_in)
+        f1_warp = warp(f1, flow1_in)
+        f_in = torch.cat([ft_, f0_warp, f1_warp, flow0_in, flow1_in], 1)
+        out = self.convblock(f_in)
+        flow0, flow1 = torch.chunk(out[:, :4, ...], 2, 1)
+        ft_ = out[:, 4:, ...]
+        flow0 = flow0 + 2.0 * resize(flow0_in, scale_factor=2.0)
+        flow1 = flow1 + 2.0 * resize(flow1_in, scale_factor=2.0)
+        return flow0, flow1, ft_
\ No newline at end of file
diff --git a/opensora/models/frame_interpolation/networks/blocks/multi_flow.py b/opensora/models/frame_interpolation/networks/blocks/multi_flow.py
new file mode 100644
index 000000000..4563c3262
--- /dev/null
+++ b/opensora/models/frame_interpolation/networks/blocks/multi_flow.py
@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+from utils.flow_utils import warp
+from networks.blocks.ifrnet import (
+    convrelu, resize,
+    ResBlock,
+)
+
+
+def multi_flow_combine(comb_block, img0, img1, flow0, flow1, 
+                       mask=None, img_res=None, mean=None):
+        '''
+            A parallel implementation of multiple flow field warping 
+            comb_block: An nn.Seqential object.
+            img shape: [b, c, h, w]
+            flow shape: [b, 2*num_flows, h, w]
+            mask (opt):
+                If 'mask' is None, the function conduct a simple average.
+            img_res (opt):
+                If 'img_res' is None, the function adds zero instead. 
+            mean (opt):
+                If 'mean' is None, the function adds zero instead.       
+        '''
+        b, c, h, w = flow0.shape
+        num_flows = c // 2
+        flow0   =   flow0.reshape(b, num_flows, 2, h, w).reshape(-1, 2, h, w)
+        flow1   =   flow1.reshape(b, num_flows, 2, h, w).reshape(-1, 2, h, w)
+        
+        mask    =    mask.reshape(b, num_flows, 1, h, w
+                            ).reshape(-1, 1, h, w) if mask is not None else None
+        img_res = img_res.reshape(b, num_flows, 3, h, w
+                            ).reshape(-1, 3, h, w)  if img_res is not None else 0
+        img0 = torch.stack([img0] * num_flows, 1).reshape(-1, 3, h, w)
+        img1 = torch.stack([img1] * num_flows, 1).reshape(-1, 3, h, w)
+        mean = torch.stack([mean] * num_flows, 1).reshape(-1, 1, 1, 1
+                                                    ) if mean is not None else 0
+        
+        img0_warp = warp(img0, flow0)
+        img1_warp = warp(img1, flow1)
+        img_warps = mask * img0_warp + (1 - mask) * img1_warp + mean + img_res
+        img_warps = img_warps.reshape(b, num_flows, 3, h, w)
+        imgt_pred = img_warps.mean(1) + comb_block(img_warps.view(b, -1, h, w))
+        return imgt_pred
+
+
+class MultiFlowDecoder(nn.Module):
+    def __init__(self, in_ch, skip_ch, num_flows=3):
+        super(MultiFlowDecoder, self).__init__()
+        self.num_flows = num_flows
+        self.convblock = nn.Sequential(
+            convrelu(in_ch*3+4, in_ch*3), 
+            ResBlock(in_ch*3, skip_ch), 
+            nn.ConvTranspose2d(in_ch*3, 8*num_flows, 4, 2, 1, bias=True)
+        )
+        
+    def forward(self, ft_, f0, f1, flow0, flow1):
+        n = self.num_flows
+        f0_warp = warp(f0, flow0)
+        f1_warp = warp(f1, flow1)
+        out = self.convblock(torch.cat([ft_, f0_warp, f1_warp, flow0, flow1], 1))
+        delta_flow0, delta_flow1, mask, img_res = torch.split(out, [2*n, 2*n, n, 3*n], 1)
+        mask = torch.sigmoid(mask)
+        
+        flow0 = delta_flow0 + 2.0 * resize(flow0, scale_factor=2.0
+                                           ).repeat(1, self.num_flows, 1, 1)
+        flow1 = delta_flow1 + 2.0 * resize(flow1, scale_factor=2.0
+                                           ).repeat(1, self.num_flows, 1, 1)
+        
+        return flow0, flow1, mask, img_res
\ No newline at end of file
diff --git a/opensora/models/frame_interpolation/networks/blocks/raft.py b/opensora/models/frame_interpolation/networks/blocks/raft.py
new file mode 100644
index 000000000..9fb85ad65
--- /dev/null
+++ b/opensora/models/frame_interpolation/networks/blocks/raft.py
@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def resize(x, scale_factor):
+    return F.interpolate(x, scale_factor=scale_factor, mode="bilinear", align_corners=False)
+
+
+def bilinear_sampler(img, coords, mask=False):
+    """ Wrapper for grid_sample, uses pixel coordinates """
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1,1], dim=-1)
+    xgrid = 2*xgrid/(W-1) - 1
+    ygrid = 2*ygrid/(H-1) - 1
+
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+    img = F.grid_sample(img, grid, align_corners=True)
+
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+
+    return img
+
+
+def coords_grid(batch, ht, wd, device):
+    coords = torch.meshgrid(torch.arange(ht, device=device), 
+                            torch.arange(wd, device=device), 
+                            indexing='ij')
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+
+class SmallUpdateBlock(nn.Module):
+    def __init__(self, cdim, hidden_dim, flow_dim, corr_dim, fc_dim,
+                 corr_levels=4, radius=3, scale_factor=None):
+        super(SmallUpdateBlock, self).__init__()
+        cor_planes = corr_levels * (2 * radius + 1) **2
+        self.scale_factor = scale_factor
+
+        self.convc1 = nn.Conv2d(2 * cor_planes, corr_dim, 1, padding=0)
+        self.convf1 = nn.Conv2d(4, flow_dim*2, 7, padding=3)
+        self.convf2 = nn.Conv2d(flow_dim*2, flow_dim, 3, padding=1)
+        self.conv = nn.Conv2d(corr_dim+flow_dim, fc_dim, 3, padding=1)
+
+        self.gru = nn.Sequential(
+            nn.Conv2d(fc_dim+4+cdim, hidden_dim, 3, padding=1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
+        )
+
+        self.feat_head = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(hidden_dim, cdim, 3, padding=1),
+        )
+
+        self.flow_head = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(hidden_dim, 4, 3, padding=1),
+        )
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+            
+    def forward(self, net, flow, corr):
+        net = resize(net, 1 / self.scale_factor
+                      ) if self.scale_factor is not None else net
+        cor = self.lrelu(self.convc1(corr))
+        flo = self.lrelu(self.convf1(flow))
+        flo = self.lrelu(self.convf2(flo))
+        cor_flo = torch.cat([cor, flo], dim=1)
+        inp = self.lrelu(self.conv(cor_flo))
+        inp = torch.cat([inp, flow, net], dim=1)
+
+        out = self.gru(inp)
+        delta_net = self.feat_head(out)
+        delta_flow = self.flow_head(out)
+        
+        if self.scale_factor is not None:
+            delta_net = resize(delta_net, scale_factor=self.scale_factor)
+            delta_flow = self.scale_factor * resize(delta_flow, scale_factor=self.scale_factor)
+        
+        return delta_net, delta_flow
+
+
+class BasicUpdateBlock(nn.Module):
+    def __init__(self, cdim, hidden_dim, flow_dim, corr_dim, corr_dim2, 
+                 fc_dim, corr_levels=4, radius=3, scale_factor=None, out_num=1):
+        super(BasicUpdateBlock, self).__init__()
+        cor_planes = corr_levels * (2 * radius + 1) **2
+
+        self.scale_factor = scale_factor
+        self.convc1 = nn.Conv2d(2 * cor_planes, corr_dim, 1, padding=0)
+        self.convc2 = nn.Conv2d(corr_dim, corr_dim2, 3, padding=1)
+        self.convf1 = nn.Conv2d(4, flow_dim*2, 7, padding=3)
+        self.convf2 = nn.Conv2d(flow_dim*2, flow_dim, 3, padding=1)
+        self.conv = nn.Conv2d(flow_dim+corr_dim2, fc_dim, 3, padding=1)
+
+        self.gru = nn.Sequential(
+            nn.Conv2d(fc_dim+4+cdim, hidden_dim, 3, padding=1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
+        )
+
+        self.feat_head = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(hidden_dim, cdim, 3, padding=1),
+        )
+
+        self.flow_head = nn.Sequential(
+            nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
+            nn.LeakyReLU(negative_slope=0.1, inplace=True),
+            nn.Conv2d(hidden_dim, 4*out_num, 3, padding=1),
+        )
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+            
+    def forward(self, net, flow, corr):
+        net = resize(net, 1 / self.scale_factor
+                      ) if self.scale_factor is not None else net
+        cor = self.lrelu(self.convc1(corr))
+        cor = self.lrelu(self.convc2(cor))
+        flo = self.lrelu(self.convf1(flow))
+        flo = self.lrelu(self.convf2(flo))
+        cor_flo = torch.cat([cor, flo], dim=1)
+        inp = self.lrelu(self.conv(cor_flo))
+        inp = torch.cat([inp, flow, net], dim=1)
+
+        out = self.gru(inp)
+        delta_net = self.feat_head(out)
+        delta_flow = self.flow_head(out)
+        
+        if self.scale_factor is not None:
+            delta_net = resize(delta_net, scale_factor=self.scale_factor)
+            delta_flow = self.scale_factor * resize(delta_flow, scale_factor=self.scale_factor)
+        return delta_net, delta_flow
+
+
+class BidirCorrBlock:
+    def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
+        self.num_levels = num_levels
+        self.radius = radius
+        self.corr_pyramid = []
+        self.corr_pyramid_T = []
+
+        corr = BidirCorrBlock.corr(fmap1, fmap2)
+        batch, h1, w1, dim, h2, w2 = corr.shape
+        corr_T = corr.clone().permute(0, 4, 5, 3, 1, 2)
+
+        corr = corr.reshape(batch*h1*w1, dim, h2, w2)
+        corr_T = corr_T.reshape(batch*h2*w2, dim, h1, w1)
+        
+        self.corr_pyramid.append(corr)
+        self.corr_pyramid_T.append(corr_T)
+
+        for _ in range(self.num_levels-1):
+            corr = F.avg_pool2d(corr, 2, stride=2)
+            corr_T = F.avg_pool2d(corr_T, 2, stride=2)
+            self.corr_pyramid.append(corr)
+            self.corr_pyramid_T.append(corr_T)
+
+    def __call__(self, coords0, coords1):
+        r = self.radius
+        coords0 = coords0.permute(0, 2, 3, 1)
+        coords1 = coords1.permute(0, 2, 3, 1)
+        assert coords0.shape == coords1.shape, f"coords0 shape: [{coords0.shape}] is not equal to [{coords1.shape}]"
+        batch, h1, w1, _ = coords0.shape
+
+        out_pyramid = []
+        out_pyramid_T = []
+        for i in range(self.num_levels):
+            corr = self.corr_pyramid[i]
+            corr_T = self.corr_pyramid_T[i]
+
+            dx = torch.linspace(-r, r, 2*r+1, device=coords0.device)
+            dy = torch.linspace(-r, r, 2*r+1, device=coords0.device)
+            delta = torch.stack(torch.meshgrid(dy, dx, indexing='ij'), axis=-1)
+            delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
+
+            centroid_lvl_0 = coords0.reshape(batch*h1*w1, 1, 1, 2) / 2**i
+            centroid_lvl_1 = coords1.reshape(batch*h1*w1, 1, 1, 2) / 2**i
+            coords_lvl_0 = centroid_lvl_0 + delta_lvl
+            coords_lvl_1 = centroid_lvl_1 + delta_lvl
+
+            corr = bilinear_sampler(corr, coords_lvl_0)
+            corr_T = bilinear_sampler(corr_T, coords_lvl_1)
+            corr = corr.view(batch, h1, w1, -1)
+            corr_T = corr_T.view(batch, h1, w1, -1)
+            out_pyramid.append(corr)
+            out_pyramid_T.append(corr_T)
+
+        out = torch.cat(out_pyramid, dim=-1)
+        out_T = torch.cat(out_pyramid_T, dim=-1)
+        return out.permute(0, 3, 1, 2).contiguous().float(), out_T.permute(0, 3, 1, 2).contiguous().float()
+
+    @staticmethod
+    def corr(fmap1, fmap2):
+        batch, dim, ht, wd = fmap1.shape
+        fmap1 = fmap1.view(batch, dim, ht*wd)
+        fmap2 = fmap2.view(batch, dim, ht*wd) 
+        
+        corr = torch.matmul(fmap1.transpose(1,2), fmap2)
+        corr = corr.view(batch, ht, wd, 1, ht, wd)
+        return corr  / torch.sqrt(torch.tensor(dim).float())
\ No newline at end of file
diff --git a/opensora/models/frame_interpolation/readme.md b/opensora/models/frame_interpolation/readme.md
new file mode 100644
index 000000000..54cbfa2ee
--- /dev/null
+++ b/opensora/models/frame_interpolation/readme.md
@@ -0,0 +1,17 @@
+#### Frame Interpolation
+
+We use AMT as our frame interpolation model. (Thanks [AMT](https://github.com/MCG-NKU/AMT)) After sampling, you can use frame interpolation model to interpolate your video smoothly.
+
+1. Download the pretrained weights from [AMT](https://github.com/MCG-NKU/AMT), we recommend using the largest model AMT-G to achieve the best performance. 
+2. Run the script of frame interpolation.
+```
+python opensora/models/frame_interpolation/interpolation.py --ckpt /path/to/ckpt --niters 1 --input /path/to/input/video.mp4 --output_path /path/to/output/floder --frame_rate 30
+```
+3. The output video will be stored at output_path and its duration time is equal `the total number of frames after frame interpolation / the frame rate`
+##### Frame Interpolation Specific Settings
+
+* `--ckpt`: Pretrained model of [AMT](https://github.com/MCG-NKU/AMT). We use AMT-G as our frame interpolation model. 
+* `--niter`: Iterations of interpolation. With $m$ input frames, `[N_ITER]` $=n$ corresponds to $2^n\times (m-1)+1$ output frames.
+* `--input`: Path of the input video.
+* `--output_path`: Folder Path of the output video.
+* `--frame_rate"`: Frame rate of the output video. 
diff --git a/opensora/models/frame_interpolation/utils/__init__.py b/opensora/models/frame_interpolation/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/opensora/models/frame_interpolation/utils/build_utils.py b/opensora/models/frame_interpolation/utils/build_utils.py
new file mode 100644
index 000000000..2d34c3b8d
--- /dev/null
+++ b/opensora/models/frame_interpolation/utils/build_utils.py
@@ -0,0 +1,12 @@
+import importlib
+
+
+def base_build_fn(module, cls, params):
+    return getattr(importlib.import_module(
+                    module, package=None), cls)(**params)
+
+
+def build_from_cfg(config):
+    module, cls = config['name'].rsplit(".", 1)
+    params = config.get('params', {})
+    return base_build_fn(module, cls, params)
diff --git a/opensora/models/frame_interpolation/utils/dist_utils.py b/opensora/models/frame_interpolation/utils/dist_utils.py
new file mode 100644
index 000000000..6337f9991
--- /dev/null
+++ b/opensora/models/frame_interpolation/utils/dist_utils.py
@@ -0,0 +1,48 @@
+import os
+import torch
+
+
+def get_world_size():
+    """Find OMPI world size without calling mpi functions
+    :rtype: int
+    """
+    if os.environ.get('PMI_SIZE') is not None:
+        return int(os.environ.get('PMI_SIZE') or 1)
+    elif os.environ.get('OMPI_COMM_WORLD_SIZE') is not None:
+        return int(os.environ.get('OMPI_COMM_WORLD_SIZE') or 1)
+    else:
+        return torch.cuda.device_count()
+
+
+def get_global_rank():
+    """Find OMPI world rank without calling mpi functions
+    :rtype: int
+    """
+    if os.environ.get('PMI_RANK') is not None:
+        return int(os.environ.get('PMI_RANK') or 0)
+    elif os.environ.get('OMPI_COMM_WORLD_RANK') is not None:
+        return int(os.environ.get('OMPI_COMM_WORLD_RANK') or 0)
+    else:
+        return 0
+
+
+def get_local_rank():
+    """Find OMPI local rank without calling mpi functions
+    :rtype: int
+    """
+    if os.environ.get('MPI_LOCALRANKID') is not None:
+        return int(os.environ.get('MPI_LOCALRANKID') or 0)
+    elif os.environ.get('OMPI_COMM_WORLD_LOCAL_RANK') is not None:
+        return int(os.environ.get('OMPI_COMM_WORLD_LOCAL_RANK') or 0)
+    else:
+        return 0
+
+
+def get_master_ip():
+    if os.environ.get('AZ_BATCH_MASTER_NODE') is not None:
+        return os.environ.get('AZ_BATCH_MASTER_NODE').split(':')[0]
+    elif os.environ.get('AZ_BATCHAI_MPI_MASTER_NODE') is not None:
+        return os.environ.get('AZ_BATCHAI_MPI_MASTER_NODE')
+    else:
+        return "127.0.0.1"
+
diff --git a/opensora/models/frame_interpolation/utils/flow_utils.py b/opensora/models/frame_interpolation/utils/flow_utils.py
new file mode 100644
index 000000000..84fca2049
--- /dev/null
+++ b/opensora/models/frame_interpolation/utils/flow_utils.py
@@ -0,0 +1,122 @@
+import numpy as np
+import torch
+from PIL import ImageFile
+import torch.nn.functional as F
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def warp(img, flow):
+    B, _, H, W = flow.shape
+    xx = torch.linspace(-1.0, 1.0, W).view(1, 1, 1, W).expand(B, -1, H, -1)
+    yy = torch.linspace(-1.0, 1.0, H).view(1, 1, H, 1).expand(B, -1, -1, W)
+    grid = torch.cat([xx, yy], 1).to(img)
+    flow_ = torch.cat([flow[:, 0:1, :, :] / ((W - 1.0) / 2.0), flow[:, 1:2, :, :] / ((H - 1.0) / 2.0)], 1)
+    grid_ = (grid + flow_).permute(0, 2, 3, 1)
+    output = F.grid_sample(input=img, grid=grid_, mode='bilinear', padding_mode='border', align_corners=True)
+    return output
+
+
+def make_colorwheel():
+    """
+    Generates a color wheel for optical flow visualization as presented in:
+        Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+        URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
+    Code follows the original C++ source code of Daniel Scharstein.
+    Code follows the the Matlab source code of Deqing Sun.
+    Returns:
+        np.ndarray: Color wheel
+    """
+
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros((ncols, 3))
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
+    col = col+RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
+    colorwheel[col:col+YG, 1] = 255
+    col = col+YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
+    col = col+GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
+    colorwheel[col:col+CB, 2] = 255
+    col = col+CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
+    col = col+BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+def flow_uv_to_colors(u, v, convert_to_bgr=False):
+    """
+    Applies the flow color wheel to (possibly clipped) flow components u and v.
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+    Args:
+        u (np.ndarray): Input horizontal flow of shape [H,W]
+        v (np.ndarray): Input vertical flow of shape [H,W]
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
+    colorwheel = make_colorwheel()  # shape [55x3]
+    ncols = colorwheel.shape[0]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    a = np.arctan2(-v, -u)/np.pi
+    fk = (a+1) / 2*(ncols-1)
+    k0 = np.floor(fk).astype(np.int32)
+    k1 = k0 + 1
+    k1[k1 == ncols] = 0
+    f = fk - k0
+    for i in range(colorwheel.shape[1]):
+        tmp = colorwheel[:,i]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1-f)*col0 + f*col1
+        idx = (rad <= 1)
+        col[idx]  = 1 - rad[idx] * (1-col[idx])
+        col[~idx] = col[~idx] * 0.75   # out of range
+        # Note the 2-i => BGR instead of RGB
+        ch_idx = 2-i if convert_to_bgr else i
+        flow_image[:,:,ch_idx] = np.floor(255 * col)
+    return flow_image
+
+def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
+    """
+    Expects a two dimensional flow image of shape.
+    Args:
+        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
+        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
+    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
+    if clip_flow is not None:
+        flow_uv = np.clip(flow_uv, 0, clip_flow)
+    u = flow_uv[:,:,0]
+    v = flow_uv[:,:,1]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    rad_max = np.max(rad)
+    epsilon = 1e-5
+    u = u / (rad_max + epsilon)
+    v = v / (rad_max + epsilon)
+    return flow_uv_to_colors(u, v, convert_to_bgr)
\ No newline at end of file
diff --git a/opensora/models/frame_interpolation/utils/utils.py b/opensora/models/frame_interpolation/utils/utils.py
new file mode 100644
index 000000000..0473226d4
--- /dev/null
+++ b/opensora/models/frame_interpolation/utils/utils.py
@@ -0,0 +1,297 @@
+import re
+import sys
+import torch
+import random
+import numpy as np
+from PIL import ImageFile
+import torch.nn.functional as F
+from imageio import imread, imwrite
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+class AverageMeter():
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0.
+        self.avg = 0.
+        self.sum = 0.
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+class AverageMeterGroups:
+    def __init__(self) -> None:
+        self.meter_dict = dict()
+    
+    def update(self, dict, n=1):
+        for name, val in dict.items():
+            if self.meter_dict.get(name) is None:
+                self.meter_dict[name] = AverageMeter()
+            self.meter_dict[name].update(val, n)
+    
+    def reset(self, name=None):
+        if name is None:
+            for v in self.meter_dict.values():
+                v.reset()
+        else:
+            meter = self.meter_dict.get(name)
+            if meter is not None:
+                meter.reset()
+    
+    def avg(self, name):
+        meter = self.meter_dict.get(name)
+        if meter is not None:
+            return meter.avg
+
+
+class InputPadder:
+    """ Pads images such that dimensions are divisible by divisor """
+    def __init__(self, dims, divisor=16):
+        self.ht, self.wd = dims[-2:]
+        pad_ht = (((self.ht // divisor) + 1) * divisor - self.ht) % divisor
+        pad_wd = (((self.wd // divisor) + 1) * divisor - self.wd) % divisor
+        self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
+
+    def pad(self, *inputs):
+        if len(inputs) == 1:
+            return F.pad(inputs[0], self._pad, mode='replicate')
+        else:
+            return [F.pad(x, self._pad, mode='replicate') for x in inputs]
+
+    def unpad(self, *inputs):
+        if len(inputs) == 1:
+            return self._unpad(inputs[0])
+        else:
+            return [self._unpad(x) for x in inputs]
+    
+    def _unpad(self, x):
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
+        return x[..., c[0]:c[1], c[2]:c[3]]
+
+
+def img2tensor(img):
+    if img.shape[-1] > 3:
+        img = img[:,:,:3]
+    return torch.tensor(img).permute(2, 0, 1).unsqueeze(0) / 255.0
+
+
+def tensor2img(img_t):
+    return (img_t * 255.).detach(
+                        ).squeeze(0).permute(1, 2, 0).cpu().numpy(
+                        ).clip(0, 255).astype(np.uint8)
+
+def seed_all(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def read(file):
+    if file.endswith('.float3'): return readFloat(file)
+    elif file.endswith('.flo'): return readFlow(file)
+    elif file.endswith('.ppm'): return readImage(file)
+    elif file.endswith('.pgm'): return readImage(file)
+    elif file.endswith('.png'): return readImage(file)
+    elif file.endswith('.jpg'): return readImage(file)
+    elif file.endswith('.pfm'): return readPFM(file)[0]
+    else: raise Exception('don\'t know how to read %s' % file)
+
+
+def write(file, data):
+    if file.endswith('.float3'): return writeFloat(file, data)
+    elif file.endswith('.flo'): return writeFlow(file, data)
+    elif file.endswith('.ppm'): return writeImage(file, data)
+    elif file.endswith('.pgm'): return writeImage(file, data)
+    elif file.endswith('.png'): return writeImage(file, data)
+    elif file.endswith('.jpg'): return writeImage(file, data)
+    elif file.endswith('.pfm'): return writePFM(file, data)
+    else: raise Exception('don\'t know how to write %s' % file)
+
+
+def readPFM(file):
+    file = open(file, 'rb')
+
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().rstrip()
+    if header.decode("ascii") == 'PF':
+        color = True
+    elif header.decode("ascii") == 'Pf':
+        color = False
+    else:
+        raise Exception('Not a PFM file.')
+
+    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode("ascii"))
+    if dim_match:
+        width, height = list(map(int, dim_match.groups()))
+    else:
+        raise Exception('Malformed PFM header.')
+
+    scale = float(file.readline().decode("ascii").rstrip())
+    if scale < 0:
+        endian = '<'
+        scale = -scale
+    else:
+        endian = '>'
+
+    data = np.fromfile(file, endian + 'f')
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    return data, scale
+
+
+def writePFM(file, image, scale=1):
+    file = open(file, 'wb')
+
+    color = None
+
+    if image.dtype.name != 'float32':
+        raise Exception('Image dtype must be float32.')
+
+    image = np.flipud(image)
+
+    if len(image.shape) == 3 and image.shape[2] == 3:
+        color = True
+    elif len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1:
+        color = False
+    else:
+        raise Exception('Image must have H x W x 3, H x W x 1 or H x W dimensions.')
+
+    file.write('PF\n' if color else 'Pf\n'.encode())
+    file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0]))
+
+    endian = image.dtype.byteorder
+
+    if endian == '<' or endian == '=' and sys.byteorder == 'little':
+        scale = -scale
+
+    file.write('%f\n'.encode() % scale)
+
+    image.tofile(file)
+
+
+def readFlow(name):
+    if name.endswith('.pfm') or name.endswith('.PFM'):
+        return readPFM(name)[0][:,:,0:2]
+
+    f = open(name, 'rb')
+
+    header = f.read(4)
+    if header.decode("utf-8") != 'PIEH':
+        raise Exception('Flow file header does not contain PIEH')
+
+    width = np.fromfile(f, np.int32, 1).squeeze()
+    height = np.fromfile(f, np.int32, 1).squeeze()
+
+    flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2))
+
+    return flow.astype(np.float32)
+
+
+def readImage(name):
+    if name.endswith('.pfm') or name.endswith('.PFM'):
+        data = readPFM(name)[0]
+        if len(data.shape)==3:
+            return data[:,:,0:3]
+        else:
+            return data
+    return imread(name)
+
+
+def writeImage(name, data):
+    if name.endswith('.pfm') or name.endswith('.PFM'):
+        return writePFM(name, data, 1)
+    return imwrite(name, data)
+
+
+def writeFlow(name, flow):
+    f = open(name, 'wb')
+    f.write('PIEH'.encode('utf-8'))
+    np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f)
+    flow = flow.astype(np.float32)
+    flow.tofile(f)
+
+
+def readFloat(name):
+    f = open(name, 'rb')
+
+    if(f.readline().decode("utf-8"))  != 'float\n':
+        raise Exception('float file %s did not contain <float> keyword' % name)
+
+    dim = int(f.readline())
+
+    dims = []
+    count = 1
+    for i in range(0, dim):
+        d = int(f.readline())
+        dims.append(d)
+        count *= d
+
+    dims = list(reversed(dims))
+
+    data = np.fromfile(f, np.float32, count).reshape(dims)
+    if dim > 2:
+        data = np.transpose(data, (2, 1, 0))
+        data = np.transpose(data, (1, 0, 2))
+
+    return data
+
+
+def writeFloat(name, data):
+    f = open(name, 'wb')
+
+    dim=len(data.shape)
+    if dim>3:
+        raise Exception('bad float file dimension: %d' % dim)
+
+    f.write(('float\n').encode('ascii'))
+    f.write(('%d\n' % dim).encode('ascii'))
+
+    if dim == 1:
+        f.write(('%d\n' % data.shape[0]).encode('ascii'))
+    else:
+        f.write(('%d\n' % data.shape[1]).encode('ascii'))
+        f.write(('%d\n' % data.shape[0]).encode('ascii'))
+        for i in range(2, dim):
+            f.write(('%d\n' % data.shape[i]).encode('ascii'))
+
+    data = data.astype(np.float32)
+    if dim==2:
+        data.tofile(f)
+
+    else:
+        np.transpose(data, (2, 0, 1)).tofile(f)
+
+
+def check_dim_and_resize(tensor_list):
+    shape_list = []
+    for t in tensor_list:
+        shape_list.append(t.shape[2:])
+
+    if len(set(shape_list)) > 1:
+        desired_shape = shape_list[0]
+        print(f'Inconsistent size of input video frames. All frames will be resized to {desired_shape}')
+        
+        resize_tensor_list = []
+        for t in tensor_list:
+            resize_tensor_list.append(torch.nn.functional.interpolate(t, size=tuple(desired_shape), mode='bilinear'))
+
+        tensor_list = resize_tensor_list
+
+    return tensor_list
+
diff --git a/opensora/models/super_resolution/README.md b/opensora/models/super_resolution/README.md
new file mode 100644
index 000000000..869cee78d
--- /dev/null
+++ b/opensora/models/super_resolution/README.md
@@ -0,0 +1,25 @@
+
+## Environment Preparation
+
+For video super resolution, please prepare your own python envirment from [RGT](https://github.com/zhengchen1999/RGT)  and down the [ckpt](https://drive.google.com/drive/folders/1zxrr31Kp2D_N9a-OUAPaJEn_yTaSXTfZ) into the folder like 
+```bash
+./experiments/pretrained_models/RGT_x2.pth
+```
+
+## Video Super Resolution
+The inferencing instruction is in [run.py](run.py).
+```bash
+python run.py --SR x4 --root_path /path_to_root --input_dir /path_to_input_dir --output_dir /path_to_video_output
+```
+You can configure some more detailed parameters in [run.py](run.py) such as . 
+```bash
+--mul_numwork 16 --use_chop False
+```
+We recommend using `` --use_chop = False `` when memory allows. 
+Note that in our tests.
+
+A single frame of 256x256 requires about 3G RAM-Usage, and a single 4090 card can process about one frame per second.
+
+A single frame of 512x512 takes about 19G RAM-Usage, and a single 4090 takes about 5 seconds to process a frame.
+
+
diff --git a/opensora/models/super_resolution/basicsr/__init__.py b/opensora/models/super_resolution/basicsr/__init__.py
new file mode 100644
index 000000000..3360401f7
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/__init__.py
@@ -0,0 +1,6 @@
+from .archs import *
+from .data import *
+from .metrics import *
+from .models import *
+from .test import *
+from .utils import *
diff --git a/opensora/models/super_resolution/basicsr/archs/__init__.py b/opensora/models/super_resolution/basicsr/archs/__init__.py
new file mode 100644
index 000000000..cfb1e4d7b
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/archs/__init__.py
@@ -0,0 +1,25 @@
+import importlib
+from copy import deepcopy
+from os import path as osp
+
+from basicsr.utils import get_root_logger, scandir
+from basicsr.utils.registry import ARCH_REGISTRY
+
+__all__ = ['build_network']
+
+# automatically scan and import arch modules for registry
+# scan all the files under the 'archs' folder and collect files ending with
+# '_arch.py'
+arch_folder = osp.dirname(osp.abspath(__file__))
+arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_arch.py')]
+# import all the arch modules
+_arch_modules = [importlib.import_module(f'basicsr.archs.{file_name}') for file_name in arch_filenames]
+
+
+def build_network(opt):
+    opt = deepcopy(opt)
+    network_type = opt.pop('type')
+    net = ARCH_REGISTRY.get(network_type)(**opt)
+    logger = get_root_logger()
+    logger.info(f'Network [{net.__class__.__name__}] is created.')
+    return net
diff --git a/opensora/models/super_resolution/basicsr/archs/arch_util.py b/opensora/models/super_resolution/basicsr/archs/arch_util.py
new file mode 100644
index 000000000..1719e8e9f
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/archs/arch_util.py
@@ -0,0 +1,318 @@
+import collections.abc
+import math
+import torch
+import torchvision
+import warnings
+from distutils.version import LooseVersion
+from itertools import repeat
+from torch import nn as nn
+from torch.nn import functional as F
+from torch.nn import init as init
+from torch.nn.modules.batchnorm import _BatchNorm
+
+# from basicsr.ops.dcn import ModulatedDeformConvPack, modulated_deform_conv
+from basicsr.utils import get_root_logger
+
+
+@torch.no_grad()
+def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
+    """Initialize network weights.
+
+    Args:
+        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
+        scale (float): Scale initialized weights, especially for residual
+            blocks. Default: 1.
+        bias_fill (float): The value to fill bias. Default: 0
+        kwargs (dict): Other arguments for initialization function.
+    """
+    if not isinstance(module_list, list):
+        module_list = [module_list]
+    for module in module_list:
+        for m in module.modules():
+            if isinstance(m, nn.Conv2d):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, nn.Linear):
+                init.kaiming_normal_(m.weight, **kwargs)
+                m.weight.data *= scale
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+            elif isinstance(m, _BatchNorm):
+                init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    m.bias.data.fill_(bias_fill)
+
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+
+class ResidualBlockNoBN(nn.Module):
+    """Residual block without BN.
+
+    It has a style of:
+        ---Conv-ReLU-Conv-+-
+         |________________|
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+            Default: 64.
+        res_scale (float): Residual scale. Default: 1.
+        pytorch_init (bool): If set to True, use pytorch default init,
+            otherwise, use default_init_weights. Default: False.
+    """
+
+    def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):
+        super(ResidualBlockNoBN, self).__init__()
+        self.res_scale = res_scale
+        self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
+        self.relu = nn.ReLU(inplace=True)
+
+        if not pytorch_init:
+            default_init_weights([self.conv1, self.conv2], 0.1)
+
+    def forward(self, x):
+        identity = x
+        out = self.conv2(self.relu(self.conv1(x)))
+        return identity + out * self.res_scale
+
+
+class Upsample(nn.Sequential):
+    """Upsample module.
+
+    Args:
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+    """
+
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:  # scale = 2^n
+            for _ in range(int(math.log(scale, 2))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f'scale {scale} is not supported. Supported scales: 2^n and 3.')
+        super(Upsample, self).__init__(*m)
+
+
+def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True):
+    """Warp an image or feature map with optical flow.
+
+    Args:
+        x (Tensor): Tensor with size (n, c, h, w).
+        flow (Tensor): Tensor with size (n, h, w, 2), normal value.
+        interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.
+        padding_mode (str): 'zeros' or 'border' or 'reflection'.
+            Default: 'zeros'.
+        align_corners (bool): Before pytorch 1.3, the default value is
+            align_corners=True. After pytorch 1.3, the default value is
+            align_corners=False. Here, we use the True as default.
+
+    Returns:
+        Tensor: Warped image or feature map.
+    """
+    assert x.size()[-2:] == flow.size()[1:3]
+    _, _, h, w = x.size()
+    # create mesh grid
+    grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))
+    grid = torch.stack((grid_x, grid_y), 2).float()  # W(x), H(y), 2
+    grid.requires_grad = False
+
+    vgrid = grid + flow
+    # scale grid to [-1,1]
+    vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0
+    vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0
+    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
+    output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)
+
+    # TODO, what if align_corners=False
+    return output
+
+
+def resize_flow(flow, size_type, sizes, interp_mode='bilinear', align_corners=False):
+    """Resize a flow according to ratio or shape.
+
+    Args:
+        flow (Tensor): Precomputed flow. shape [N, 2, H, W].
+        size_type (str): 'ratio' or 'shape'.
+        sizes (list[int | float]): the ratio for resizing or the final output
+            shape.
+            1) The order of ratio should be [ratio_h, ratio_w]. For
+            downsampling, the ratio should be smaller than 1.0 (i.e., ratio
+            < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,
+            ratio > 1.0).
+            2) The order of output_size should be [out_h, out_w].
+        interp_mode (str): The mode of interpolation for resizing.
+            Default: 'bilinear'.
+        align_corners (bool): Whether align corners. Default: False.
+
+    Returns:
+        Tensor: Resized flow.
+    """
+    _, _, flow_h, flow_w = flow.size()
+    if size_type == 'ratio':
+        output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])
+    elif size_type == 'shape':
+        output_h, output_w = sizes[0], sizes[1]
+    else:
+        raise ValueError(f'Size type should be ratio or shape, but got type {size_type}.')
+
+    input_flow = flow.clone()
+    ratio_h = output_h / flow_h
+    ratio_w = output_w / flow_w
+    input_flow[:, 0, :, :] *= ratio_w
+    input_flow[:, 1, :, :] *= ratio_h
+    resized_flow = F.interpolate(
+        input=input_flow, size=(output_h, output_w), mode=interp_mode, align_corners=align_corners)
+    return resized_flow
+
+
+# TODO: may write a cpp file
+def pixel_unshuffle(x, scale):
+    """ Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)
+
+
+# class DCNv2Pack(ModulatedDeformConvPack):
+#     """Modulated deformable conv for deformable alignment.
+#
+#     Different from the official DCNv2Pack, which generates offsets and masks
+#     from the preceding features, this DCNv2Pack takes another different
+#     features to generate offsets and masks.
+#
+#     Ref:
+#         Delving Deep into Deformable Alignment in Video Super-Resolution.
+#     """
+#
+#     def forward(self, x, feat):
+#         out = self.conv_offset(feat)
+#         o1, o2, mask = torch.chunk(out, 3, dim=1)
+#         offset = torch.cat((o1, o2), dim=1)
+#         mask = torch.sigmoid(mask)
+#
+#         offset_absmean = torch.mean(torch.abs(offset))
+#         if offset_absmean > 50:
+#             logger = get_root_logger()
+#             logger.warning(f'Offset abs mean is {offset_absmean}, larger than 50.')
+#
+#         if LooseVersion(torchvision.__version__) >= LooseVersion('0.9.0'):
+#             return torchvision.ops.deform_conv2d(x, offset, self.weight, self.bias, self.stride, self.padding,
+#                                                  self.dilation, mask)
+#         else:
+#             return modulated_deform_conv(x, offset, mask, self.weight, self.bias, self.stride, self.padding,
+#                                          self.dilation, self.groups, self.deformable_groups)
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn(
+            'mean is more than 2 std from [a, b] in nn.init.trunc_normal_. '
+            'The distribution of values may be incorrect.',
+            stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        low = norm_cdf((a - mean) / std)
+        up = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [low, up], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * low - 1, 2 * up - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    r"""Fills the input Tensor with values drawn from a truncated
+    normal distribution.
+
+    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/weight_init.py
+
+    The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+
+    Args:
+        tensor: an n-dimensional `torch.Tensor`
+        mean: the mean of the normal distribution
+        std: the standard deviation of the normal distribution
+        a: the minimum cutoff value
+        b: the maximum cutoff value
+
+    Examples:
+        >>> w = torch.empty(3, 5)
+        >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+# From PyTorch
+def _ntuple(n):
+
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
diff --git a/opensora/models/super_resolution/basicsr/archs/rgt_arch.py b/opensora/models/super_resolution/basicsr/archs/rgt_arch.py
new file mode 100644
index 000000000..810fd765a
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/archs/rgt_arch.py
@@ -0,0 +1,757 @@
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from torch import Tensor
+from torch.nn import functional as F
+
+from timm.models.layers import DropPath, trunc_normal_
+from einops.layers.torch import Rearrange
+from einops import rearrange, repeat
+
+import math
+import numpy as np
+
+import random
+
+from basicsr.utils.registry import ARCH_REGISTRY
+
+
+def img2windows(img, H_sp, W_sp):
+    """
+    Input: Image (B, C, H, W)
+    Output: Window Partition (B', N, C)
+    """
+    B, C, H, W = img.shape
+    img_reshape = img.view(B, C, H // H_sp, H_sp, W // W_sp, W_sp)
+    img_perm = img_reshape.permute(0, 2, 4, 3, 5, 1).contiguous().reshape(-1, H_sp* W_sp, C)
+    return img_perm
+
+
+def windows2img(img_splits_hw, H_sp, W_sp, H, W):
+    """
+    Input: Window Partition (B', N, C)
+    Output: Image (B, H, W, C)
+    """
+    B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))
+
+    img = img_splits_hw.view(B, H // H_sp, W // W_sp, H_sp, W_sp, -1)
+    img = img.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return img
+
+
+class Gate(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.conv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim) # DW Conv
+
+    def forward(self, x, H, W):
+        # Split
+        x1, x2 = x.chunk(2, dim = -1)
+        B, N, C = x.shape
+        x2 = self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C//2, H, W)).flatten(2).transpose(-1, -2).contiguous()
+
+        return x1 * x2
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.sg = Gate(hidden_features//2)
+        self.fc2 = nn.Linear(hidden_features//2, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x, H, W):
+        """
+        Input: x: (B, H*W, C), H, W
+        Output: x: (B, H*W, C)
+        """
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+
+        x = self.sg(x, H, W)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class DynamicPosBias(nn.Module):
+    # The implementation builds on Crossformer code https://github.com/cheerss/CrossFormer/blob/main/models/crossformer.py
+    """ Dynamic Relative Position Bias.
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        residual (bool):  If True, use residual strage to connect conv.
+    """
+    def __init__(self, dim, num_heads, residual):
+        super().__init__()
+        self.residual = residual
+        self.num_heads = num_heads
+        self.pos_dim = dim // 4
+        self.pos_proj = nn.Linear(2, self.pos_dim)
+        self.pos1 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos2 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim)
+        )
+        self.pos3 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.num_heads)
+        )
+    def forward(self, biases):
+        if self.residual:
+            pos = self.pos_proj(biases) # 2Gh-1 * 2Gw-1, heads
+            pos = pos + self.pos1(pos)
+            pos = pos + self.pos2(pos)
+            pos = self.pos3(pos)
+        else:
+            pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases))))
+        return pos
+
+
+class WindowAttention(nn.Module):
+    def __init__(self, dim, idx, split_size=[8,8], dim_out=None, num_heads=6, attn_drop=0., proj_drop=0., qk_scale=None, position_bias=True):
+        super().__init__()
+        self.dim = dim
+        self.dim_out = dim_out or dim
+        self.split_size = split_size
+        self.num_heads = num_heads
+        self.idx = idx
+        self.position_bias = position_bias
+
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        if idx == 0:
+            H_sp, W_sp = self.split_size[0], self.split_size[1]
+        elif idx == 1:
+            W_sp, H_sp = self.split_size[0], self.split_size[1]
+        else:
+            print ("ERROR MODE", idx)
+            exit(0)
+        self.H_sp = H_sp
+        self.W_sp = W_sp
+
+        if self.position_bias:
+            self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False)
+            # generate mother-set
+            position_bias_h = torch.arange(1 - self.H_sp, self.H_sp)
+            position_bias_w = torch.arange(1 - self.W_sp, self.W_sp)
+            biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w]))
+            biases = biases.flatten(1).transpose(0, 1).contiguous().float()
+            self.register_buffer('rpe_biases', biases)
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(self.H_sp)
+            coords_w = torch.arange(self.W_sp)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
+            coords_flatten = torch.flatten(coords, 1)
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+            relative_coords[:, :, 0] += self.H_sp - 1
+            relative_coords[:, :, 1] += self.W_sp - 1
+            relative_coords[:, :, 0] *= 2 * self.W_sp - 1
+            relative_position_index = relative_coords.sum(-1)
+            self.register_buffer('relative_position_index', relative_position_index)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+
+    def im2win(self, x, H, W):
+        B, N, C = x.shape
+        x = x.transpose(-2,-1).contiguous().view(B, C, H, W)
+        x = img2windows(x, self.H_sp, self.W_sp)
+        x = x.reshape(-1, self.H_sp* self.W_sp, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
+        return x
+
+    def forward(self, qkv, H, W, mask=None):
+        """
+        Input: qkv: (B, 3*L, C), H, W, mask: (B, N, N), N is the window size
+        Output: x (B, H, W, C)
+        """
+        q,k,v = qkv[0], qkv[1], qkv[2]
+
+        B, L, C = q.shape
+        assert L == H * W, "flatten img_tokens has wrong size"
+
+        # partition the q,k,v, image to window
+        q = self.im2win(q, H, W)
+        k = self.im2win(k, H, W)
+        v = self.im2win(v, H, W)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))  # B head N C @ B head C N --> B head N N
+
+        # calculate drpe
+        if self.position_bias:
+            pos = self.pos(self.rpe_biases)
+            # select position bias
+            relative_position_bias = pos[self.relative_position_index.view(-1)].view(
+                self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1)
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        N = attn.shape[3]
+
+        # use mask for shift window
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+
+        attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v)
+        x = x.transpose(1, 2).reshape(-1, self.H_sp* self.W_sp, C)  # B head N N @ B head N C
+
+        # merge the window, window to image
+        x = windows2img(x, self.H_sp, self.W_sp, H, W)  # B H' W' C
+
+        return x
+
+
+class L_SA(nn.Module):
+    # The implementation builds on CAT code https://github.com/zhengchen1999/CAT/blob/main/basicsr/archs/cat_arch.py
+    def __init__(self, dim, num_heads,
+                 split_size=[2,4], shift_size=[1,2], qkv_bias=False, qk_scale=None,
+                 drop=0., attn_drop=0., idx=0, reso=64, rs_id=0):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.split_size = split_size
+        self.shift_size = shift_size
+        self.idx = idx
+        self.rs_id = rs_id
+        self.patches_resolution = reso
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+
+        assert 0 <= self.shift_size[0] < self.split_size[0], "shift_size must in 0-split_size0"
+        assert 0 <= self.shift_size[1] < self.split_size[1], "shift_size must in 0-split_size1"
+
+        self.branch_num = 2
+
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(drop)
+
+        self.attns = nn.ModuleList([
+                WindowAttention(
+                    dim//2, idx = i,
+                    split_size=split_size, num_heads=num_heads//2, dim_out=dim//2,
+                    qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop, position_bias=True)
+                for i in range(self.branch_num)])
+
+        if (self.rs_id % 2 == 0 and self.idx > 0 and (self.idx - 2) % 4 == 0) or (self.rs_id % 2 != 0 and self.idx % 4 == 0):
+            attn_mask = self.calculate_mask(self.patches_resolution, self.patches_resolution)
+
+            self.register_buffer("attn_mask_0", attn_mask[0])
+            self.register_buffer("attn_mask_1", attn_mask[1])
+        else:
+            attn_mask = None
+
+            self.register_buffer("attn_mask_0", None)
+            self.register_buffer("attn_mask_1", None)
+
+        self.get_v = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1,groups=dim) # DW Conv
+
+    def calculate_mask(self, H, W):
+        # The implementation builds on Swin Transformer code https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
+        # calculate attention mask for Rwin
+        img_mask_0 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=0
+        img_mask_1 = torch.zeros((1, H, W, 1))  # 1 H W 1 idx=1
+        h_slices_0 = (slice(0, -self.split_size[0]),
+                    slice(-self.split_size[0], -self.shift_size[0]),
+                    slice(-self.shift_size[0], None))
+        w_slices_0 = (slice(0, -self.split_size[1]),
+                    slice(-self.split_size[1], -self.shift_size[1]),
+                    slice(-self.shift_size[1], None))
+
+        h_slices_1 = (slice(0, -self.split_size[1]),
+                    slice(-self.split_size[1], -self.shift_size[1]),
+                    slice(-self.shift_size[1], None))
+        w_slices_1 = (slice(0, -self.split_size[0]),
+                    slice(-self.split_size[0], -self.shift_size[0]),
+                    slice(-self.shift_size[0], None))
+        cnt = 0
+        for h in h_slices_0:
+            for w in w_slices_0:
+                img_mask_0[:, h, w, :] = cnt
+                cnt += 1
+        cnt = 0
+        for h in h_slices_1:
+            for w in w_slices_1:
+                img_mask_1[:, h, w, :] = cnt
+                cnt += 1
+
+        # calculate mask for H-Shift
+        img_mask_0 = img_mask_0.view(1, H // self.split_size[0], self.split_size[0], W // self.split_size[1], self.split_size[1], 1)
+        img_mask_0 = img_mask_0.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.split_size[0], self.split_size[1], 1) # nW, sw[0], sw[1], 1
+        mask_windows_0 = img_mask_0.view(-1, self.split_size[0] * self.split_size[1])
+        attn_mask_0 = mask_windows_0.unsqueeze(1) - mask_windows_0.unsqueeze(2)
+        attn_mask_0 = attn_mask_0.masked_fill(attn_mask_0 != 0, float(-100.0)).masked_fill(attn_mask_0 == 0, float(0.0))
+
+        # calculate mask for V-Shift
+        img_mask_1 = img_mask_1.view(1, H // self.split_size[1], self.split_size[1], W // self.split_size[0], self.split_size[0], 1)
+        img_mask_1 = img_mask_1.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, self.split_size[1], self.split_size[0], 1) # nW, sw[1], sw[0], 1
+        mask_windows_1 = img_mask_1.view(-1, self.split_size[1] * self.split_size[0])
+        attn_mask_1 = mask_windows_1.unsqueeze(1) - mask_windows_1.unsqueeze(2)
+        attn_mask_1 = attn_mask_1.masked_fill(attn_mask_1 != 0, float(-100.0)).masked_fill(attn_mask_1 == 0, float(0.0))
+
+        return attn_mask_0, attn_mask_1
+
+    def forward(self, x, H, W):
+        """
+        Input: x: (B, H*W, C), x_size: (H, W)
+        Output: x: (B, H*W, C)
+        """
+
+        B, L, C = x.shape
+        assert L == H * W, "flatten img_tokens has wrong size"
+
+        qkv = self.qkv(x).reshape(B, -1, 3, C).permute(2, 0, 1, 3) # 3, B, HW, C
+        # v without partition
+        v = qkv[2].transpose(-2,-1).contiguous().view(B, C, H, W)
+
+
+        max_split_size = max(self.split_size[0], self.split_size[1])
+        pad_l = pad_t = 0
+        pad_r = (max_split_size - W % max_split_size) % max_split_size
+        pad_b = (max_split_size - H % max_split_size) % max_split_size
+
+        qkv = qkv.reshape(3*B, H, W, C).permute(0, 3, 1, 2) # 3B C H W
+        qkv = F.pad(qkv, (pad_l, pad_r, pad_t, pad_b)).reshape(3, B, C, -1).transpose(-2, -1) # l r t b
+        _H = pad_b + H
+        _W = pad_r + W
+        _L = _H * _W
+
+        if (self.rs_id % 2 == 0 and self.idx > 0 and (self.idx - 2) % 4 == 0) or (self.rs_id % 2 != 0 and self.idx % 4 == 0):
+            qkv = qkv.view(3, B, _H, _W, C)
+            # H-Shift
+            qkv_0 = torch.roll(qkv[:,:,:,:,:C//2], shifts=(-self.shift_size[0], -self.shift_size[1]), dims=(2, 3))
+            qkv_0 = qkv_0.view(3, B, _L, C//2)
+            # V-Shift
+            qkv_1 = torch.roll(qkv[:,:,:,:,C//2:], shifts=(-self.shift_size[1], -self.shift_size[0]), dims=(2, 3))
+            qkv_1 = qkv_1.view(3, B, _L, C//2)
+
+            if self.patches_resolution != _H or self.patches_resolution != _W:
+                mask_tmp = self.calculate_mask(_H, _W)
+                # H-Rwin
+                x1_shift = self.attns[0](qkv_0, _H, _W, mask=mask_tmp[0].to(x.device))
+                # V-Rwin
+                x2_shift = self.attns[1](qkv_1, _H, _W, mask=mask_tmp[1].to(x.device))
+
+            else:
+                # H-Rwin
+                x1_shift = self.attns[0](qkv_0, _H, _W, mask=self.attn_mask_0)
+                # V-Rwin
+                x2_shift = self.attns[1](qkv_1, _H, _W, mask=self.attn_mask_1)
+
+            x1 = torch.roll(x1_shift, shifts=(self.shift_size[0], self.shift_size[1]), dims=(1, 2))
+            x2 = torch.roll(x2_shift, shifts=(self.shift_size[1], self.shift_size[0]), dims=(1, 2))
+            x1 = x1[:, :H, :W, :].reshape(B, L, C//2)
+            x2 = x2[:, :H, :W, :].reshape(B, L, C//2)
+            # Concat
+            attened_x = torch.cat([x1,x2], dim=2)
+        else:
+            # V-Rwin
+            x1 = self.attns[0](qkv[:,:,:,:C//2], _H, _W)[:, :H, :W, :].reshape(B, L, C//2)
+            # H-Rwin
+            x2 = self.attns[1](qkv[:,:,:,C//2:], _H, _W)[:, :H, :W, :].reshape(B, L, C//2)
+            # Concat
+            attened_x = torch.cat([x1,x2], dim=2)
+
+        # mix
+        lcm = self.get_v(v)
+        lcm = lcm.permute(0, 2, 3, 1).contiguous().view(B, L, C)
+
+        x = attened_x + lcm
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class RG_SA(nn.Module):
+    """
+    Recursive-Generalization Self-Attention (RG-SA).
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+        c_ratio (float): channel adjustment factor.
+    """
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., c_ratio=0.5):
+        super(RG_SA, self).__init__()
+        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.cr = int(dim * c_ratio) # scaled channel dimension
+
+        # self.scale = qk_scale or head_dim ** -0.5
+        self.scale = qk_scale or (head_dim * c_ratio) ** -0.5
+
+        # RGM
+        self.reduction1 = nn.Conv2d(dim, dim, kernel_size=4, stride=4, groups=dim)
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1,groups=dim)
+        self.conv = nn.Conv2d(dim, self.cr, kernel_size=1, stride=1)
+        self.norm_act = nn.Sequential(
+            nn.LayerNorm(self.cr),
+            nn.GELU())
+        # CA
+        self.q = nn.Linear(dim, self.cr, bias=qkv_bias)
+        self.k = nn.Linear(self.cr, self.cr, bias=qkv_bias)
+        self.v = nn.Linear(self.cr, dim, bias=qkv_bias)
+
+        # CPE
+        self.cpe = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim)
+
+        self.proj = nn.Linear(dim, dim)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, H, W):
+        B, N, C = x.shape
+
+        _scale = 1
+
+        # reduction
+        _x = x.permute(0, 2, 1).reshape(B, C, H, W).contiguous()
+
+        if self.training:
+            _time = max(int(math.log(H//4, 4)), int(math.log(W//4, 4)))
+        else:
+            _time = max(int(math.log(H//16, 4)), int(math.log(W//16, 4)))
+            if _time < 2: _time = 2 # testing _time must equal or larger than training _time (2)
+
+        _scale = 4 ** _time
+
+        # Recursion xT
+        for _ in range(_time):
+            _x = self.reduction1(_x)
+
+        _x = self.conv(self.dwconv(_x)).reshape(B, self.cr, -1).permute(0, 2, 1).contiguous()  # shape=(B, N', C')
+        _x = self.norm_act(_x)
+
+        # q, k, v, where q_shape=(B, N, C'), k_shape=(B, N', C'), v_shape=(B, N', C)
+        q = self.q(x).reshape(B, N, self.num_heads, int(self.cr / self.num_heads)).permute(0, 2, 1, 3)
+        k = self.k(_x).reshape(B, -1, self.num_heads, int(self.cr / self.num_heads)).permute(0, 2, 1, 3)
+        v = self.v(_x).reshape(B, -1, self.num_heads, int(C / self.num_heads)).permute(0, 2, 1, 3)
+        
+        # corss-attention
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        # CPE
+        # v_shape=(B, H, N', C//H)
+        v = v + self.cpe(v.transpose(1, 2).reshape(B, -1, C).transpose(1, 2).contiguous().view(B, C, H // _scale, W // _scale)).view(B, C, -1).view(B, self.num_heads, int(C / self.num_heads), -1).transpose(-1, -2)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0.,
+                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, idx=0, 
+                 rs_id=0, split_size=[2,4], shift_size=[1,2], reso=64, c_ratio=0.5, layerscale_value=1e-4):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        if idx % 2 == 0:
+            self.attn = L_SA(
+                dim, split_size=split_size, shift_size=shift_size, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
+                drop=drop, idx=idx, reso=reso, rs_id=rs_id
+            )
+        else:
+            self.attn = RG_SA(
+                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
+                proj_drop=drop, c_ratio=c_ratio
+            )
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, out_features=dim, act_layer=act_layer)
+        self.norm2 = norm_layer(dim)
+
+        # HAI
+        self.gamma = nn.Parameter(layerscale_value * torch.ones((dim)), requires_grad=True)
+
+    def forward(self, x, x_size):
+        H , W = x_size
+
+        res = x
+
+        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
+        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
+
+        # HAI
+        x = x + (res * self.gamma)
+
+        return x
+
+
+class ResidualGroup(nn.Module):
+
+    def __init__(   self,
+                    dim,
+                    reso,
+                    num_heads,
+                    mlp_ratio=4.,
+                    qkv_bias=False,
+                    qk_scale=None,
+                    drop=0.,
+                    attn_drop=0.,
+                    drop_paths=None,
+                    act_layer=nn.GELU,
+                    norm_layer=nn.LayerNorm,
+                    depth=2,
+                    use_chk=False,
+                    resi_connection='1conv',
+                    rs_id=0,
+                    split_size=[8,8],
+                    c_ratio = 0.5):
+        super().__init__()
+        self.use_chk = use_chk
+        self.reso = reso
+
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=drop_paths[i],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                idx = i,
+                rs_id = rs_id,
+                split_size = split_size,
+                shift_size = [split_size[0]//2, split_size[1]//2],
+                c_ratio = c_ratio,
+                )for i in range(depth)])
+
+
+        if resi_connection == '1conv':
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == '3conv':
+            self.conv = nn.Sequential(
+                nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(dim // 4, dim, 3, 1, 1))
+
+    def forward(self, x, x_size):
+        """
+        Input: x: (B, H*W, C), x_size: (H, W)
+        Output: x: (B, H*W, C)
+        """
+        H, W = x_size
+        res = x
+        for blk in self.blocks:
+            if self.use_chk:
+                x = checkpoint.checkpoint(blk, x, x_size)
+            else:
+                x = blk(x, x_size)
+        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
+        x = self.conv(x)
+        x = rearrange(x, "b c h w -> b (h w) c")
+        x = res + x
+
+        return x
+
+
+class Upsample(nn.Sequential):
+    """Upsample module.
+    Args:
+        scale (int): Scale factor. Supported scales: 2^n and 3.
+        num_feat (int): Channel number of intermediate features.
+    """
+    def __init__(self, scale, num_feat):
+        m = []
+        if (scale & (scale - 1)) == 0:  # scale = 2^n
+            for _ in range(int(math.log(scale, 2))):
+                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
+                m.append(nn.PixelShuffle(2))
+        elif scale == 3:
+            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
+            m.append(nn.PixelShuffle(3))
+        else:
+            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
+        super(Upsample, self).__init__(*m)
+
+
+@ARCH_REGISTRY.register()
+class RGT(nn.Module):
+
+    def __init__(self,
+                img_size=64,
+                in_chans=3,
+                embed_dim=180,
+                depth=[2,2,2,2],
+                num_heads=[2,2,2,2],
+                mlp_ratio=4.,
+                qkv_bias=True,
+                qk_scale=None,
+                drop_rate=0.,
+                attn_drop_rate=0.,
+                drop_path_rate=0.1,
+                act_layer=nn.GELU,
+                norm_layer=nn.LayerNorm,
+                use_chk=False,
+                upscale=2,
+                img_range=1.,
+                resi_connection='1conv',
+                split_size=[8,8],
+                c_ratio=0.5,
+                **kwargs):
+        super().__init__()
+
+        num_in_ch = in_chans
+        num_out_ch = in_chans
+        num_feat = 64
+        self.img_range = img_range
+        if in_chans == 3:
+            rgb_mean = (0.4488, 0.4371, 0.4040)
+            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
+        else:
+            self.mean = torch.zeros(1, 1, 1, 1)
+        self.upscale = upscale
+
+        # ------------------------- 1, Shallow Feature Extraction ------------------------- #
+        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        # ------------------------- 2, Deep Feature Extraction ------------------------- #
+        self.num_layers = len(depth)
+        self.use_chk = use_chk
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        heads=num_heads
+
+        self.before_RG = nn.Sequential(
+            Rearrange('b c h w -> b (h w) c'),
+            nn.LayerNorm(embed_dim)
+        )
+
+        curr_dim = embed_dim
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, np.sum(depth))]  # stochastic depth decay rule
+
+        self.layers = nn.ModuleList()
+        for i in range(self.num_layers):
+            layer = ResidualGroup(
+                dim=embed_dim,
+                num_heads=heads[i],
+                reso=img_size,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_paths=dpr[sum(depth[:i]):sum(depth[:i + 1])],
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                depth=depth[i],
+                use_chk=use_chk,
+                resi_connection=resi_connection,
+                rs_id=i,
+                split_size = split_size,
+                c_ratio = c_ratio
+                )
+            self.layers.append(layer)
+
+        self.norm = norm_layer(curr_dim)
+        # build the last conv layer in deep feature extraction
+        if resi_connection == '1conv':
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == '3conv':
+            # to save parameters and memory
+            self.conv_after_body = nn.Sequential(
+                nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), nn.LeakyReLU(negative_slope=0.2, inplace=True),
+                nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1))
+
+        # ------------------------- 3, Reconstruction ------------------------- #
+        self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True))
+        self.upsample = Upsample(upscale, num_feat)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward_features(self, x):
+        _, _, H, W = x.shape
+        x_size = [H, W]
+        x = self.before_RG(x)
+        for layer in self.layers:
+            x = layer(x, x_size)
+        x = self.norm(x)
+        x = rearrange(x, "b (h w) c -> b c h w", h=H, w=W)
+
+        return x
+
+    def forward(self, x):
+        """
+        Input: x: (B, C, H, W)
+        """
+        self.mean = self.mean.type_as(x)
+        x = (x - self.mean) * self.img_range
+
+        x = self.conv_first(x)
+        x = self.conv_after_body(self.forward_features(x)) + x
+        x = self.conv_before_upsample(x)
+        x = self.conv_last(self.upsample(x))
+
+        x = x / self.img_range + self.mean
+        return x
+
+
+if __name__ == '__main__':
+    upscale = 1
+    height = 62
+    width = 66
+    model = RGT(
+        upscale=2,
+        in_chans=3,
+        img_size=64,
+        img_range=1.,
+        depth=[6,6,6,6,6,6],
+        embed_dim=180,
+        num_heads=[6,6,6,6,6,6],
+        mlp_ratio=2,
+        resi_connection='1conv',
+        split_size=[8, 8],
+        upsampler='pixelshuffle').cuda()
+    # print(model)
+    print(height, width)
+
+    x = torch.randn((1, 3, height, width)).cuda()
+    x = model(x)
+    print(x.shape)
\ No newline at end of file
diff --git a/opensora/models/super_resolution/basicsr/archs/vgg_arch.py b/opensora/models/super_resolution/basicsr/archs/vgg_arch.py
new file mode 100644
index 000000000..1e44c0a11
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/archs/vgg_arch.py
@@ -0,0 +1,161 @@
+import os
+import torch
+from collections import OrderedDict
+from torch import nn as nn
+from torchvision.models import vgg as vgg
+
+from basicsr.utils.registry import ARCH_REGISTRY
+
+VGG_PRETRAIN_PATH = 'experiments/pretrained_models/vgg19-dcbb9e9d.pth'
+NAMES = {
+    'vgg11': [
+        'conv1_1', 'relu1_1', 'pool1', 'conv2_1', 'relu2_1', 'pool2', 'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2',
+        'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2',
+        'pool5'
+    ],
+    'vgg13': [
+        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'pool4',
+        'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'pool5'
+    ],
+    'vgg16': [
+        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'pool3', 'conv4_1', 'relu4_1', 'conv4_2',
+        'relu4_2', 'conv4_3', 'relu4_3', 'pool4', 'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3',
+        'pool5'
+    ],
+    'vgg19': [
+        'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1', 'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',
+        'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3', 'conv4_1',
+        'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4', 'conv5_1', 'relu5_1',
+        'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5'
+    ]
+}
+
+
+def insert_bn(names):
+    """Insert bn layer after each conv.
+
+    Args:
+        names (list): The list of layer names.
+
+    Returns:
+        list: The list of layer names with bn layers.
+    """
+    names_bn = []
+    for name in names:
+        names_bn.append(name)
+        if 'conv' in name:
+            position = name.replace('conv', '')
+            names_bn.append('bn' + position)
+    return names_bn
+
+
+@ARCH_REGISTRY.register()
+class VGGFeatureExtractor(nn.Module):
+    """VGG network for feature extraction.
+
+    In this implementation, we allow users to choose whether use normalization
+    in the input feature and the type of vgg network. Note that the pretrained
+    path must fit the vgg type.
+
+    Args:
+        layer_name_list (list[str]): Forward function returns the corresponding
+            features according to the layer_name_list.
+            Example: {'relu1_1', 'relu2_1', 'relu3_1'}.
+        vgg_type (str): Set the type of vgg network. Default: 'vgg19'.
+        use_input_norm (bool): If True, normalize the input image. Importantly,
+            the input feature must in the range [0, 1]. Default: True.
+        range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
+            Default: False.
+        requires_grad (bool): If true, the parameters of VGG network will be
+            optimized. Default: False.
+        remove_pooling (bool): If true, the max pooling operations in VGG net
+            will be removed. Default: False.
+        pooling_stride (int): The stride of max pooling operation. Default: 2.
+    """
+
+    def __init__(self,
+                 layer_name_list,
+                 vgg_type='vgg19',
+                 use_input_norm=True,
+                 range_norm=False,
+                 requires_grad=False,
+                 remove_pooling=False,
+                 pooling_stride=2):
+        super(VGGFeatureExtractor, self).__init__()
+
+        self.layer_name_list = layer_name_list
+        self.use_input_norm = use_input_norm
+        self.range_norm = range_norm
+
+        self.names = NAMES[vgg_type.replace('_bn', '')]
+        if 'bn' in vgg_type:
+            self.names = insert_bn(self.names)
+
+        # only borrow layers that will be used to avoid unused params
+        max_idx = 0
+        for v in layer_name_list:
+            idx = self.names.index(v)
+            if idx > max_idx:
+                max_idx = idx
+
+        if os.path.exists(VGG_PRETRAIN_PATH):
+            vgg_net = getattr(vgg, vgg_type)(pretrained=False)
+            state_dict = torch.load(VGG_PRETRAIN_PATH, map_location=lambda storage, loc: storage)
+            vgg_net.load_state_dict(state_dict)
+        else:
+            vgg_net = getattr(vgg, vgg_type)(pretrained=True)
+
+        features = vgg_net.features[:max_idx + 1]
+
+        modified_net = OrderedDict()
+        for k, v in zip(self.names, features):
+            if 'pool' in k:
+                # if remove_pooling is true, pooling operation will be removed
+                if remove_pooling:
+                    continue
+                else:
+                    # in some cases, we may want to change the default stride
+                    modified_net[k] = nn.MaxPool2d(kernel_size=2, stride=pooling_stride)
+            else:
+                modified_net[k] = v
+
+        self.vgg_net = nn.Sequential(modified_net)
+
+        if not requires_grad:
+            self.vgg_net.eval()
+            for param in self.parameters():
+                param.requires_grad = False
+        else:
+            self.vgg_net.train()
+            for param in self.parameters():
+                param.requires_grad = True
+
+        if self.use_input_norm:
+            # the mean is for image with range [0, 1]
+            self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+            # the std is for image with range [0, 1]
+            self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+    def forward(self, x):
+        """Forward function.
+
+        Args:
+            x (Tensor): Input tensor with shape (n, c, h, w).
+
+        Returns:
+            Tensor: Forward results.
+        """
+        if self.range_norm:
+            x = (x + 1) / 2
+        if self.use_input_norm:
+            x = (x - self.mean) / self.std
+
+        output = {}
+        for key, layer in self.vgg_net._modules.items():
+            x = layer(x)
+            if key in self.layer_name_list:
+                output[key] = x.clone()
+
+        return output
\ No newline at end of file
diff --git a/opensora/models/super_resolution/basicsr/data/__init__.py b/opensora/models/super_resolution/basicsr/data/__init__.py
new file mode 100644
index 000000000..12a7ea053
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/__init__.py
@@ -0,0 +1,101 @@
+import importlib
+import numpy as np
+import random
+import torch
+import torch.utils.data
+from copy import deepcopy
+from functools import partial
+from os import path as osp
+
+from basicsr.data.prefetch_dataloader import PrefetchDataLoader
+from basicsr.utils import get_root_logger, scandir
+from basicsr.utils.dist_util import get_dist_info
+from basicsr.utils.registry import DATASET_REGISTRY
+
+__all__ = ['build_dataset', 'build_dataloader']
+
+# automatically scan and import dataset modules for registry
+# scan all the files under the data folder with '_dataset' in file names
+data_folder = osp.dirname(osp.abspath(__file__))
+dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
+# import all the dataset modules
+_dataset_modules = [importlib.import_module(f'basicsr.data.{file_name}') for file_name in dataset_filenames]
+
+
+def build_dataset(dataset_opt):
+    """Build dataset from options.
+
+    Args:
+        dataset_opt (dict): Configuration for dataset. It must contain:
+            name (str): Dataset name.
+            type (str): Dataset type.
+    """
+    dataset_opt = deepcopy(dataset_opt)
+    dataset = DATASET_REGISTRY.get(dataset_opt['type'])(dataset_opt)
+    # logger = get_root_logger()
+    # logger.info(f'Dataset [{dataset.__class__.__name__}] - {dataset_opt["name"]} is built.')
+    return dataset
+
+
+def build_dataloader(dataset, dataset_opt, num_gpu=1, dist=False, sampler=None, seed=None):
+    """Build dataloader.
+
+    Args:
+        dataset (torch.utils.data.Dataset): Dataset.
+        dataset_opt (dict): Dataset options. It contains the following keys:
+            phase (str): 'train' or 'val'.
+            num_worker_per_gpu (int): Number of workers for each GPU.
+            batch_size_per_gpu (int): Training batch size for each GPU.
+        num_gpu (int): Number of GPUs. Used only in the train phase.
+            Default: 1.
+        dist (bool): Whether in distributed training. Used only in the train
+            phase. Default: False.
+        sampler (torch.utils.data.sampler): Data sampler. Default: None.
+        seed (int | None): Seed. Default: None
+    """
+    phase = dataset_opt['phase']
+    rank, _ = get_dist_info()
+    if phase == 'train':
+        if dist:  # distributed training
+            batch_size = dataset_opt['batch_size_per_gpu']
+            num_workers = dataset_opt['num_worker_per_gpu']
+        else:  # non-distributed training
+            multiplier = 1 if num_gpu == 0 else num_gpu
+            batch_size = dataset_opt['batch_size_per_gpu'] * multiplier
+            num_workers = dataset_opt['num_worker_per_gpu'] * multiplier
+        dataloader_args = dict(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            sampler=sampler,
+            drop_last=True)
+        if sampler is None:
+            dataloader_args['shuffle'] = True
+        dataloader_args['worker_init_fn'] = partial(
+            worker_init_fn, num_workers=num_workers, rank=rank, seed=seed) if seed is not None else None
+    elif phase in ['val', 'test']:  # validation
+        dataloader_args = dict(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
+    else:
+        raise ValueError(f"Wrong dataset phase: {phase}. Supported ones are 'train', 'val' and 'test'.")
+
+    dataloader_args['pin_memory'] = dataset_opt.get('pin_memory', False)
+    dataloader_args['persistent_workers'] = dataset_opt.get('persistent_workers', False)
+
+    prefetch_mode = dataset_opt.get('prefetch_mode')
+    if prefetch_mode == 'cpu':  # CPUPrefetcher
+        num_prefetch_queue = dataset_opt.get('num_prefetch_queue', 1)
+        logger = get_root_logger()
+        logger.info(f'Use {prefetch_mode} prefetch dataloader: num_prefetch_queue = {num_prefetch_queue}')
+        return PrefetchDataLoader(num_prefetch_queue=num_prefetch_queue, **dataloader_args)
+    else:
+        # prefetch_mode=None: Normal dataloader
+        # prefetch_mode='cuda': dataloader for CUDAPrefetcher
+        return torch.utils.data.DataLoader(**dataloader_args)
+
+
+def worker_init_fn(worker_id, num_workers, rank, seed):
+    # Set the worker seed to num_workers * rank + worker_id + seed
+    worker_seed = num_workers * rank + worker_id + seed
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
diff --git a/opensora/models/super_resolution/basicsr/data/data_sampler.py b/opensora/models/super_resolution/basicsr/data/data_sampler.py
new file mode 100644
index 000000000..575452d9f
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/data_sampler.py
@@ -0,0 +1,48 @@
+import math
+import torch
+from torch.utils.data.sampler import Sampler
+
+
+class EnlargedSampler(Sampler):
+    """Sampler that restricts data loading to a subset of the dataset.
+
+    Modified from torch.utils.data.distributed.DistributedSampler
+    Support enlarging the dataset for iteration-based training, for saving
+    time when restart the dataloader after each epoch
+
+    Args:
+        dataset (torch.utils.data.Dataset): Dataset used for sampling.
+        num_replicas (int | None): Number of processes participating in
+            the training. It is usually the world_size.
+        rank (int | None): Rank of the current process within num_replicas.
+        ratio (int): Enlarging ratio. Default: 1.
+    """
+
+    def __init__(self, dataset, num_replicas, rank, ratio=1):
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.num_samples = math.ceil(len(self.dataset) * ratio / self.num_replicas)
+        self.total_size = self.num_samples * self.num_replicas
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+        indices = torch.randperm(self.total_size, generator=g).tolist()
+
+        dataset_size = len(self.dataset)
+        indices = [v % dataset_size for v in indices]
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+    def set_epoch(self, epoch):
+        self.epoch = epoch
diff --git a/opensora/models/super_resolution/basicsr/data/data_util.py b/opensora/models/super_resolution/basicsr/data/data_util.py
new file mode 100644
index 000000000..de4811dad
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/data_util.py
@@ -0,0 +1,280 @@
+import cv2
+import numpy as np
+import torch
+from os import path as osp
+from torch.nn import functional as F
+
+from basicsr.utils import img2tensor, scandir
+
+
+def generate_frame_indices(crt_idx, max_frame_num, num_frames, padding='reflection'):
+    """Generate an index list for reading `num_frames` frames from a sequence
+    of images.
+
+    Args:
+        crt_idx (int): Current center index.
+        max_frame_num (int): Max number of the sequence of images (from 1).
+        num_frames (int): Reading num_frames frames.
+        padding (str): Padding mode, one of
+            'replicate' | 'reflection' | 'reflection_circle' | 'circle'
+            Examples: current_idx = 0, num_frames = 5
+            The generated frame indices under different padding mode:
+            replicate: [0, 0, 0, 1, 2]
+            reflection: [2, 1, 0, 1, 2]
+            reflection_circle: [4, 3, 0, 1, 2]
+            circle: [3, 4, 0, 1, 2]
+
+    Returns:
+        list[int]: A list of indices.
+    """
+    assert num_frames % 2 == 1, 'num_frames should be an odd number.'
+    assert padding in ('replicate', 'reflection', 'reflection_circle', 'circle'), f'Wrong padding mode: {padding}.'
+
+    max_frame_num = max_frame_num - 1  # start from 0
+    num_pad = num_frames // 2
+
+    indices = []
+    for i in range(crt_idx - num_pad, crt_idx + num_pad + 1):
+        if i < 0:
+            if padding == 'replicate':
+                pad_idx = 0
+            elif padding == 'reflection':
+                pad_idx = -i
+            elif padding == 'reflection_circle':
+                pad_idx = crt_idx + num_pad - i
+            else:
+                pad_idx = num_frames + i
+        elif i > max_frame_num:
+            if padding == 'replicate':
+                pad_idx = max_frame_num
+            elif padding == 'reflection':
+                pad_idx = max_frame_num * 2 - i
+            elif padding == 'reflection_circle':
+                pad_idx = (crt_idx - num_pad) - (i - max_frame_num)
+            else:
+                pad_idx = i - num_frames
+        else:
+            pad_idx = i
+        indices.append(pad_idx)
+    return indices
+
+
+def paired_paths_from_lmdb(folders, keys):
+    """Generate paired paths from lmdb files.
+
+    Contents of lmdb. Taking the `lq.lmdb` for example, the file structure is:
+
+    lq.lmdb
+    ├── data.mdb
+    ├── lock.mdb
+    ├── meta_info.txt
+
+    The data.mdb and lock.mdb are standard lmdb files and you can refer to
+    https://lmdb.readthedocs.io/en/release/ for more details.
+
+    The meta_info.txt is a specified txt file to record the meta information
+    of our datasets. It will be automatically created when preparing
+    datasets by our provided dataset tools.
+    Each line in the txt file records
+    1)image name (with extension),
+    2)image shape,
+    3)compression level, separated by a white space.
+    Example: `baboon.png (120,125,3) 1`
+
+    We use the image name without extension as the lmdb key.
+    Note that we use the same key for the corresponding lq and gt images.
+
+    Args:
+        folders (list[str]): A list of folder path. The order of list should
+            be [input_folder, gt_folder].
+        keys (list[str]): A list of keys identifying folders. The order should
+            be in consistent with folders, e.g., ['lq', 'gt'].
+            Note that this key is different from lmdb keys.
+
+    Returns:
+        list[str]: Returned path list.
+    """
+    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
+                               f'But got {len(folders)}')
+    assert len(keys) == 2, f'The len of keys should be 2 with [input_key, gt_key]. But got {len(keys)}'
+    input_folder, gt_folder = folders
+    input_key, gt_key = keys
+
+    if not (input_folder.endswith('.lmdb') and gt_folder.endswith('.lmdb')):
+        raise ValueError(f'{input_key} folder and {gt_key} folder should both in lmdb '
+                         f'formats. But received {input_key}: {input_folder}; '
+                         f'{gt_key}: {gt_folder}')
+    # ensure that the two meta_info files are the same
+    with open(osp.join(input_folder, 'meta_info.txt')) as fin:
+        input_lmdb_keys = [line.split('.')[0] for line in fin]
+    with open(osp.join(gt_folder, 'meta_info.txt')) as fin:
+        gt_lmdb_keys = [line.split('.')[0] for line in fin]
+    if set(input_lmdb_keys) != set(gt_lmdb_keys):
+        raise ValueError(f'Keys in {input_key}_folder and {gt_key}_folder are different.')
+    else:
+        paths = []
+        for lmdb_key in sorted(input_lmdb_keys):
+            paths.append(dict([(f'{input_key}_path', lmdb_key), (f'{gt_key}_path', lmdb_key)]))
+        return paths
+
+
+def paired_paths_from_meta_info_file(folders, keys, meta_info_file, filename_tmpl):
+    """Generate paired paths from an meta information file.
+
+    Each line in the meta information file contains the image names and
+    image shape (usually for gt), separated by a white space.
+
+    Example of an meta information file:
+    ```
+    0001_s001.png (480,480,3)
+    0001_s002.png (480,480,3)
+    ```
+
+    Args:
+        folders (list[str]): A list of folder path. The order of list should
+            be [input_folder, gt_folder].
+        keys (list[str]): A list of keys identifying folders. The order should
+            be in consistent with folders, e.g., ['lq', 'gt'].
+        meta_info_file (str): Path to the meta information file.
+        filename_tmpl (str): Template for each filename. Note that the
+            template excludes the file extension. Usually the filename_tmpl is
+            for files in the input folder.
+
+    Returns:
+        list[str]: Returned path list.
+    """
+    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
+                               f'But got {len(folders)}')
+    assert len(keys) == 2, f'The len of keys should be 2 with [input_key, gt_key]. But got {len(keys)}'
+    input_folder, gt_folder = folders
+    input_key, gt_key = keys
+
+    with open(meta_info_file, 'r') as fin:
+        gt_names = [line.strip().split(' ')[0] for line in fin]
+
+    paths = []
+    for gt_name in gt_names:
+        basename, ext = osp.splitext(osp.basename(gt_name))
+        input_name = f'{filename_tmpl.format(basename)}{ext}'
+        input_path = osp.join(input_folder, input_name)
+        gt_path = osp.join(gt_folder, gt_name)
+        paths.append(dict([(f'{input_key}_path', input_path), (f'{gt_key}_path', gt_path)]))
+    return paths
+
+
+def paired_paths_from_folder(folders, keys, filename_tmpl):
+    """Generate paired paths from folders.
+
+    Args:
+        folders (list[str]): A list of folder path. The order of list should
+            be [input_folder, gt_folder].
+        keys (list[str]): A list of keys identifying folders. The order should
+            be in consistent with folders, e.g., ['lq', 'gt'].
+        filename_tmpl (str): Template for each filename. Note that the
+            template excludes the file extension. Usually the filename_tmpl is
+            for files in the input folder.
+
+    Returns:
+        list[str]: Returned path list.
+    """
+    assert len(folders) == 2, ('The len of folders should be 2 with [input_folder, gt_folder]. '
+                               f'But got {len(folders)}')
+    assert len(keys) == 2, f'The len of keys should be 2 with [input_key, gt_key]. But got {len(keys)}'
+    input_folder, gt_folder = folders
+    input_key, gt_key = keys
+
+    input_paths = list(scandir(input_folder))
+    gt_paths = list(scandir(gt_folder))
+    assert len(input_paths) == len(gt_paths), (f'{input_key} and {gt_key} datasets have different number of images: '
+                                               f'{len(input_paths)}, {len(gt_paths)}.')
+    paths = []
+    for gt_path in gt_paths:
+        basename, ext = osp.splitext(osp.basename(gt_path))
+        input_name = f'{filename_tmpl.format(basename)}{ext}'
+        input_path = osp.join(input_folder, input_name)
+        assert input_name in input_paths, f'{input_name} is not in {input_key}_paths.'
+        gt_path = osp.join(gt_folder, gt_path)
+        paths.append(dict([(f'{input_key}_path', input_path), (f'{gt_key}_path', gt_path)]))
+    return paths
+
+
+def paths_from_folder(folder):
+    """Generate paths from folder.
+
+    Args:
+        folder (str): Folder path.
+
+    Returns:
+        list[str]: Returned path list.
+    """
+
+    paths = list(scandir(folder))
+    paths = [osp.join(folder, path) for path in paths]
+    return paths
+
+
+def paths_from_lmdb(folder):
+    """Generate paths from lmdb.
+
+    Args:
+        folder (str): Folder path.
+
+    Returns:
+        list[str]: Returned path list.
+    """
+    if not folder.endswith('.lmdb'):
+        raise ValueError(f'Folder {folder}folder should in lmdb format.')
+    with open(osp.join(folder, 'meta_info.txt')) as fin:
+        paths = [line.split('.')[0] for line in fin]
+    return paths
+
+
+def generate_gaussian_kernel(kernel_size=13, sigma=1.6):
+    """Generate Gaussian kernel used in `duf_downsample`.
+
+    Args:
+        kernel_size (int): Kernel size. Default: 13.
+        sigma (float): Sigma of the Gaussian kernel. Default: 1.6.
+
+    Returns:
+        np.array: The Gaussian kernel.
+    """
+    from scipy.ndimage import filters as filters
+    kernel = np.zeros((kernel_size, kernel_size))
+    # set element at the middle to one, a dirac delta
+    kernel[kernel_size // 2, kernel_size // 2] = 1
+    # gaussian-smooth the dirac, resulting in a gaussian filter
+    return filters.gaussian_filter(kernel, sigma)
+
+
+def duf_downsample(x, kernel_size=13, scale=4):
+    """Downsamping with Gaussian kernel used in the DUF official code.
+
+    Args:
+        x (Tensor): Frames to be downsampled, with shape (b, t, c, h, w).
+        kernel_size (int): Kernel size. Default: 13.
+        scale (int): Downsampling factor. Supported scale: (2, 3, 4).
+            Default: 4.
+
+    Returns:
+        Tensor: DUF downsampled frames.
+    """
+    assert scale in (2, 3, 4), f'Only support scale (2, 3, 4), but got {scale}.'
+
+    squeeze_flag = False
+    if x.ndim == 4:
+        squeeze_flag = True
+        x = x.unsqueeze(0)
+    b, t, c, h, w = x.size()
+    x = x.view(-1, 1, h, w)
+    pad_w, pad_h = kernel_size // 2 + scale * 2, kernel_size // 2 + scale * 2
+    x = F.pad(x, (pad_w, pad_w, pad_h, pad_h), 'reflect')
+
+    gaussian_filter = generate_gaussian_kernel(kernel_size, 0.4 * scale)
+    gaussian_filter = torch.from_numpy(gaussian_filter).type_as(x).unsqueeze(0).unsqueeze(0)
+    x = F.conv2d(x, gaussian_filter, stride=scale)
+    x = x[:, :, 2:-2, 2:-2]
+    x = x.view(b, t, c, x.size(2), x.size(3))
+    if squeeze_flag:
+        x = x.squeeze(0)
+    return x
diff --git a/opensora/models/super_resolution/basicsr/data/paired_image_dataset.py b/opensora/models/super_resolution/basicsr/data/paired_image_dataset.py
new file mode 100644
index 000000000..23749f8d5
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/paired_image_dataset.py
@@ -0,0 +1,113 @@
+from torch.utils import data as data
+from torchvision.transforms.functional import normalize
+
+from basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb, paired_paths_from_meta_info_file
+from basicsr.data.transforms import augment, paired_random_crop
+from basicsr.utils import FileClient, imfrombytes, img2tensor
+from basicsr.utils.matlab_functions import bgr2ycbcr
+from basicsr.utils.registry import DATASET_REGISTRY
+
+import numpy as np
+
+@DATASET_REGISTRY.register()
+class PairedImageDataset(data.Dataset):
+    """Paired image dataset for image restoration.
+
+    Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
+
+    There are three modes:
+    1. 'lmdb': Use lmdb files.
+        If opt['io_backend'] == lmdb.
+    2. 'meta_info_file': Use meta information file to generate paths.
+        If opt['io_backend'] != lmdb and opt['meta_info_file'] is not None.
+    3. 'folder': Scan folders to generate paths.
+        The rest.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_gt (str): Data root path for gt.
+            dataroot_lq (str): Data root path for lq.
+            meta_info_file (str): Path for meta information file.
+            io_backend (dict): IO backend type and other kwarg.
+            filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
+                Default: '{}'.
+            gt_size (int): Cropped patched size for gt patches.
+            use_hflip (bool): Use horizontal flips.
+            use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
+
+            scale (bool): Scale, which will be added automatically.
+            phase (str): 'train' or 'val'.
+    """
+
+    def __init__(self, opt):
+        super(PairedImageDataset, self).__init__()
+        self.opt = opt
+        # file client (io backend)
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+        self.mean = opt['mean'] if 'mean' in opt else None
+        self.std = opt['std'] if 'std' in opt else None
+
+        self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
+        if 'filename_tmpl' in opt:
+            self.filename_tmpl = opt['filename_tmpl']
+        else:
+            self.filename_tmpl = '{}'
+
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
+            self.io_backend_opt['client_keys'] = ['lq', 'gt']
+            self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
+        elif 'meta_info_file' in self.opt and self.opt['meta_info_file'] is not None:
+            self.paths = paired_paths_from_meta_info_file([self.lq_folder, self.gt_folder], ['lq', 'gt'],
+                                                          self.opt['meta_info_file'], self.filename_tmpl)
+        else:
+            self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        scale = self.opt['scale']
+
+        # Load gt and lq images. Dimension order: HWC; channel order: BGR;
+        
+        # image range: [0, 1], float32., H W 3
+        gt_path = self.paths[index]['gt_path']
+        img_bytes = self.file_client.get(gt_path, 'gt')
+        img_gt = imfrombytes(img_bytes, float32=True)
+        lq_path = self.paths[index]['lq_path']
+        img_bytes = self.file_client.get(lq_path, 'lq')
+        img_lq = imfrombytes(img_bytes, float32=True)
+
+        # augmentation for training
+        if self.opt['phase'] == 'train':
+            gt_size = self.opt['gt_size']
+            # random crop
+            img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
+            # flip, rotation
+            img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
+
+        # color space transform
+        if 'color' in self.opt and self.opt['color'] == 'y':
+            img_gt = bgr2ycbcr(img_gt, y_only=True)[..., None]
+            img_lq = bgr2ycbcr(img_lq, y_only=True)[..., None]
+
+        # crop the unmatched GT images during validation or testing, especially for SR benchmark datasets
+        # TODO: It is better to update the datasets, rather than force to crop
+        if self.opt['phase'] != 'train':
+            img_gt = img_gt[0:img_lq.shape[0] * scale, 0:img_lq.shape[1] * scale, :]
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
+        # normalize
+        if self.mean is not None or self.std is not None:
+            normalize(img_lq, self.mean, self.std, inplace=True)
+            normalize(img_gt, self.mean, self.std, inplace=True)
+
+        # print(img_lq.shape,img_gt.shape,img_lq.min(),img_gt.min(),img_lq.max(),img_gt.max(),lq_path,gt_path)
+
+        return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
+
+    def __len__(self):
+        return len(self.paths)
diff --git a/opensora/models/super_resolution/basicsr/data/prefetch_dataloader.py b/opensora/models/super_resolution/basicsr/data/prefetch_dataloader.py
new file mode 100644
index 000000000..508842505
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/prefetch_dataloader.py
@@ -0,0 +1,125 @@
+import queue as Queue
+import threading
+import torch
+from torch.utils.data import DataLoader
+
+
+class PrefetchGenerator(threading.Thread):
+    """A general prefetch generator.
+
+    Ref:
+    https://stackoverflow.com/questions/7323664/python-generator-pre-fetch
+
+    Args:
+        generator: Python generator.
+        num_prefetch_queue (int): Number of prefetch queue.
+    """
+
+    def __init__(self, generator, num_prefetch_queue):
+        threading.Thread.__init__(self)
+        self.queue = Queue.Queue(num_prefetch_queue)
+        self.generator = generator
+        self.daemon = True
+        self.start()
+
+    def run(self):
+        for item in self.generator:
+            self.queue.put(item)
+        self.queue.put(None)
+
+    def __next__(self):
+        next_item = self.queue.get()
+        if next_item is None:
+            raise StopIteration
+        return next_item
+
+    def __iter__(self):
+        return self
+
+
+class PrefetchDataLoader(DataLoader):
+    """Prefetch version of dataloader.
+
+    Ref:
+    https://github.com/IgorSusmelj/pytorch-styleguide/issues/5#
+
+    TODO:
+    Need to test on single gpu and ddp (multi-gpu). There is a known issue in
+    ddp.
+
+    Args:
+        num_prefetch_queue (int): Number of prefetch queue.
+        kwargs (dict): Other arguments for dataloader.
+    """
+
+    def __init__(self, num_prefetch_queue, **kwargs):
+        self.num_prefetch_queue = num_prefetch_queue
+        super(PrefetchDataLoader, self).__init__(**kwargs)
+
+    def __iter__(self):
+        return PrefetchGenerator(super().__iter__(), self.num_prefetch_queue)
+
+
+class CPUPrefetcher():
+    """CPU prefetcher.
+
+    Args:
+        loader: Dataloader.
+    """
+
+    def __init__(self, loader):
+        self.ori_loader = loader
+        self.loader = iter(loader)
+
+    def next(self):
+        try:
+            return next(self.loader)
+        except StopIteration:
+            return None
+
+    def reset(self):
+        self.loader = iter(self.ori_loader)
+
+
+class CUDAPrefetcher():
+    """CUDA prefetcher.
+
+    Ref:
+    https://github.com/NVIDIA/apex/issues/304#
+
+    It may consums more GPU memory.
+
+    Args:
+        loader: Dataloader.
+        opt (dict): Options.
+    """
+
+    def __init__(self, loader, opt):
+        self.ori_loader = loader
+        self.loader = iter(loader)
+        self.opt = opt
+        self.stream = torch.cuda.Stream()
+        self.device = torch.device('cuda' if opt['num_gpu'] != 0 else 'cpu')
+        self.preload()
+
+    def preload(self):
+        try:
+            self.batch = next(self.loader)  # self.batch is a dict
+        except StopIteration:
+            self.batch = None
+            return None
+        # put tensors to gpu
+        with torch.cuda.stream(self.stream):
+            for k, v in self.batch.items():
+                if torch.is_tensor(v):
+                    self.batch[k] = self.batch[k].to(device=self.device, non_blocking=True)
+
+    def next(self):
+        torch.cuda.current_stream().wait_stream(self.stream)
+        batch = self.batch
+        self.preload()
+        return batch
+
+    def reset(self):
+        self.loader = iter(self.ori_loader)
+        self.preload()
diff --git a/opensora/models/super_resolution/basicsr/data/single_image_dataset.py b/opensora/models/super_resolution/basicsr/data/single_image_dataset.py
new file mode 100644
index 000000000..795803a10
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/single_image_dataset.py
@@ -0,0 +1,69 @@
+from os import path as osp
+from torch.utils import data as data
+from torchvision.transforms.functional import normalize
+
+from basicsr.data.data_util import paths_from_lmdb
+from basicsr.utils import FileClient, imfrombytes, img2tensor, scandir
+from basicsr.utils.matlab_functions import rgb2ycbcr
+from basicsr.utils.registry import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+class SingleImageDataset(data.Dataset):
+    """Read only lq images in the test phase.
+
+    Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc).
+
+    There are two modes:
+    1. 'meta_info_file': Use meta information file to generate paths.
+    2. 'folder': Scan folders to generate paths.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_lq (str): Data root path for lq.
+            meta_info_file (str): Path for meta information file.
+            io_backend (dict): IO backend type and other kwarg.
+    """
+
+    def __init__(self, opt):
+        super(SingleImageDataset, self).__init__()
+        self.opt = opt
+        # file client (io backend)
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+        self.mean = opt['mean'] if 'mean' in opt else None
+        self.std = opt['std'] if 'std' in opt else None
+        self.lq_folder = opt['dataroot_lq']
+
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = [self.lq_folder]
+            self.io_backend_opt['client_keys'] = ['lq']
+            self.paths = paths_from_lmdb(self.lq_folder)
+        elif 'meta_info_file' in self.opt:
+            with open(self.opt['meta_info_file'], 'r') as fin:
+                self.paths = [osp.join(self.lq_folder, line.rstrip().split(' ')[0]) for line in fin]
+        else:
+            self.paths = sorted(list(scandir(self.lq_folder, full_path=True)))
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        # load lq image
+        lq_path = self.paths[index]
+        img_bytes = self.file_client.get(lq_path, 'lq')
+        img_lq = imfrombytes(img_bytes, float32=True)
+
+        # color space transform
+        if 'color' in self.opt and self.opt['color'] == 'y':
+            img_lq = rgb2ycbcr(img_lq, y_only=True)[..., None]
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_lq = img2tensor(img_lq, bgr2rgb=True, float32=True)
+        # normalize
+        if self.mean is not None or self.std is not None:
+            normalize(img_lq, self.mean, self.std, inplace=True)
+        return {'lq': img_lq, 'lq_path': lq_path}
+
+    def __len__(self):
+        return len(self.paths)
diff --git a/opensora/models/super_resolution/basicsr/data/transforms.py b/opensora/models/super_resolution/basicsr/data/transforms.py
new file mode 100644
index 000000000..d9bbb5fb7
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/data/transforms.py
@@ -0,0 +1,179 @@
+import cv2
+import random
+import torch
+
+
+def mod_crop(img, scale):
+    """Mod crop images, used during testing.
+
+    Args:
+        img (ndarray): Input image.
+        scale (int): Scale factor.
+
+    Returns:
+        ndarray: Result image.
+    """
+    img = img.copy()
+    if img.ndim in (2, 3):
+        h, w = img.shape[0], img.shape[1]
+        h_remainder, w_remainder = h % scale, w % scale
+        img = img[:h - h_remainder, :w - w_remainder, ...]
+    else:
+        raise ValueError(f'Wrong img ndim: {img.ndim}.')
+    return img
+
+
+def paired_random_crop(img_gts, img_lqs, gt_patch_size, scale, gt_path=None):
+    """Paired random crop. Support Numpy array and Tensor inputs.
+
+    It crops lists of lq and gt images with corresponding locations.
+
+    Args:
+        img_gts (list[ndarray] | ndarray | list[Tensor] | Tensor): GT images. Note that all images
+            should have the same shape. If the input is an ndarray, it will
+            be transformed to a list containing itself.
+        img_lqs (list[ndarray] | ndarray): LQ images. Note that all images
+            should have the same shape. If the input is an ndarray, it will
+            be transformed to a list containing itself.
+        gt_patch_size (int): GT patch size.
+        scale (int): Scale factor.
+        gt_path (str): Path to ground-truth. Default: None.
+
+    Returns:
+        list[ndarray] | ndarray: GT images and LQ images. If returned results
+            only have one element, just return ndarray.
+    """
+
+    if not isinstance(img_gts, list):
+        img_gts = [img_gts]
+    if not isinstance(img_lqs, list):
+        img_lqs = [img_lqs]
+
+    # determine input type: Numpy array or Tensor
+    input_type = 'Tensor' if torch.is_tensor(img_gts[0]) else 'Numpy'
+
+    if input_type == 'Tensor':
+        h_lq, w_lq = img_lqs[0].size()[-2:]
+        h_gt, w_gt = img_gts[0].size()[-2:]
+    else:
+        h_lq, w_lq = img_lqs[0].shape[0:2]
+        h_gt, w_gt = img_gts[0].shape[0:2]
+    lq_patch_size = gt_patch_size // scale
+
+    if h_gt != h_lq * scale or w_gt != w_lq * scale:
+        raise ValueError(f'Scale mismatches. GT ({h_gt}, {w_gt}) is not {scale}x ',
+                         f'multiplication of LQ ({h_lq}, {w_lq}).')
+    if h_lq < lq_patch_size or w_lq < lq_patch_size:
+        raise ValueError(f'LQ ({h_lq}, {w_lq}) is smaller than patch size '
+                         f'({lq_patch_size}, {lq_patch_size}). '
+                         f'Please remove {gt_path}.')
+
+    # randomly choose top and left coordinates for lq patch
+    top = random.randint(0, h_lq - lq_patch_size)
+    left = random.randint(0, w_lq - lq_patch_size)
+
+    # crop lq patch
+    if input_type == 'Tensor':
+        img_lqs = [v[:, :, top:top + lq_patch_size, left:left + lq_patch_size] for v in img_lqs]
+    else:
+        img_lqs = [v[top:top + lq_patch_size, left:left + lq_patch_size, ...] for v in img_lqs]
+
+    # crop corresponding gt patch
+    top_gt, left_gt = int(top * scale), int(left * scale)
+    if input_type == 'Tensor':
+        img_gts = [v[:, :, top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size] for v in img_gts]
+    else:
+        img_gts = [v[top_gt:top_gt + gt_patch_size, left_gt:left_gt + gt_patch_size, ...] for v in img_gts]
+    if len(img_gts) == 1:
+        img_gts = img_gts[0]
+    if len(img_lqs) == 1:
+        img_lqs = img_lqs[0]
+    return img_gts, img_lqs
+
+
+def augment(imgs, hflip=True, rotation=True, flows=None, return_status=False):
+    """Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees).
+
+    We use vertical flip and transpose for rotation implementation.
+    All the images in the list use the same augmentation.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Images to be augmented. If the input
+            is an ndarray, it will be transformed to a list.
+        hflip (bool): Horizontal flip. Default: True.
+        rotation (bool): Ratotation. Default: True.
+        flows (list[ndarray]: Flows to be augmented. If the input is an
+            ndarray, it will be transformed to a list.
+            Dimension is (h, w, 2). Default: None.
+        return_status (bool): Return the status of flip and rotation.
+            Default: False.
+
+    Returns:
+        list[ndarray] | ndarray: Augmented images and flows. If returned
+            results only have one element, just return ndarray.
+
+    """
+    hflip = hflip and random.random() < 0.5
+    vflip = rotation and random.random() < 0.5
+    rot90 = rotation and random.random() < 0.5
+
+    def _augment(img):
+        if hflip:  # horizontal
+            cv2.flip(img, 1, img)
+        if vflip:  # vertical
+            cv2.flip(img, 0, img)
+        if rot90:
+            img = img.transpose(1, 0, 2)
+        return img
+
+    def _augment_flow(flow):
+        if hflip:  # horizontal
+            cv2.flip(flow, 1, flow)
+            flow[:, :, 0] *= -1
+        if vflip:  # vertical
+            cv2.flip(flow, 0, flow)
+            flow[:, :, 1] *= -1
+        if rot90:
+            flow = flow.transpose(1, 0, 2)
+            flow = flow[:, :, [1, 0]]
+        return flow
+
+    if not isinstance(imgs, list):
+        imgs = [imgs]
+    imgs = [_augment(img) for img in imgs]
+    if len(imgs) == 1:
+        imgs = imgs[0]
+
+    if flows is not None:
+        if not isinstance(flows, list):
+            flows = [flows]
+        flows = [_augment_flow(flow) for flow in flows]
+        if len(flows) == 1:
+            flows = flows[0]
+        return imgs, flows
+    else:
+        if return_status:
+            return imgs, (hflip, vflip, rot90)
+        else:
+            return imgs
+
+
+def img_rotate(img, angle, center=None, scale=1.0):
+    """Rotate image.
+
+    Args:
+        img (ndarray): Image to be rotated.
+        angle (float): Rotation angle in degrees. Positive values mean
+            counter-clockwise rotation.
+        center (tuple[int]): Rotation center. If the center is None,
+            initialize it as the center of the image. Default: None.
+        scale (float): Isotropic scale factor. Default: 1.0.
+    """
+    (h, w) = img.shape[:2]
+
+    if center is None:
+        center = (w // 2, h // 2)
+
+    matrix = cv2.getRotationMatrix2D(center, angle, scale)
+    rotated_img = cv2.warpAffine(img, matrix, (w, h))
+    return rotated_img
diff --git a/opensora/models/super_resolution/basicsr/losses/__init__.py b/opensora/models/super_resolution/basicsr/losses/__init__.py
new file mode 100644
index 000000000..14942900c
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/losses/__init__.py
@@ -0,0 +1,26 @@
+from copy import deepcopy
+
+from basicsr.utils import get_root_logger
+from basicsr.utils.registry import LOSS_REGISTRY
+from .losses import (CharbonnierLoss, GANLoss, L1Loss, MSELoss, WeightedTVLoss, g_path_regularize,
+                     gradient_penalty_loss, r1_penalty)
+
+__all__ = [
+    'L1Loss', 'MSELoss', 'CharbonnierLoss', 'WeightedTVLoss', 'GANLoss', 'gradient_penalty_loss',
+    'r1_penalty', 'g_path_regularize'
+]
+
+
+def build_loss(opt):
+    """Build loss from options.
+
+    Args:
+        opt (dict): Configuration. It must contain:
+            type (str): Model type.
+    """
+    opt = deepcopy(opt)
+    loss_type = opt.pop('type')
+    loss = LOSS_REGISTRY.get(loss_type)(**opt)
+    logger = get_root_logger()
+    logger.info(f'Loss [{loss.__class__.__name__}] is created.')
+    return loss
diff --git a/opensora/models/super_resolution/basicsr/losses/loss_util.py b/opensora/models/super_resolution/basicsr/losses/loss_util.py
new file mode 100644
index 000000000..744eeb46d
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/losses/loss_util.py
@@ -0,0 +1,95 @@
+import functools
+from torch.nn import functional as F
+
+
+def reduce_loss(loss, reduction):
+    """Reduce loss as specified.
+
+    Args:
+        loss (Tensor): Elementwise loss tensor.
+        reduction (str): Options are 'none', 'mean' and 'sum'.
+
+    Returns:
+        Tensor: Reduced loss tensor.
+    """
+    reduction_enum = F._Reduction.get_enum(reduction)
+    # none: 0, elementwise_mean:1, sum: 2
+    if reduction_enum == 0:
+        return loss
+    elif reduction_enum == 1:
+        return loss.mean()
+    else:
+        return loss.sum()
+
+
+def weight_reduce_loss(loss, weight=None, reduction='mean'):
+    """Apply element-wise weight and reduce loss.
+
+    Args:
+        loss (Tensor): Element-wise loss.
+        weight (Tensor): Element-wise weights. Default: None.
+        reduction (str): Same as built-in losses of PyTorch. Options are
+            'none', 'mean' and 'sum'. Default: 'mean'.
+
+    Returns:
+        Tensor: Loss values.
+    """
+    # if weight is specified, apply element-wise weight
+    if weight is not None:
+        assert weight.dim() == loss.dim()
+        assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
+        loss = loss * weight
+
+    # if weight is not specified or reduction is sum, just reduce the loss
+    if weight is None or reduction == 'sum':
+        loss = reduce_loss(loss, reduction)
+    # if reduction is mean, then compute mean over weight region
+    elif reduction == 'mean':
+        if weight.size(1) > 1:
+            weight = weight.sum()
+        else:
+            weight = weight.sum() * loss.size(1)
+        loss = loss.sum() / weight
+
+    return loss
+
+
+def weighted_loss(loss_func):
+    """Create a weighted version of a given loss function.
+
+    To use this decorator, the loss function must have the signature like
+    `loss_func(pred, target, **kwargs)`. The function only needs to compute
+    element-wise loss without any reduction. This decorator will add weight
+    and reduction arguments to the function. The decorated function will have
+    the signature like `loss_func(pred, target, weight=None, reduction='mean',
+    **kwargs)`.
+
+    :Example:
+
+    >>> import torch
+    >>> @weighted_loss
+    >>> def l1_loss(pred, target):
+    >>>     return (pred - target).abs()
+
+    >>> pred = torch.Tensor([0, 2, 3])
+    >>> target = torch.Tensor([1, 1, 1])
+    >>> weight = torch.Tensor([1, 0, 1])
+
+    >>> l1_loss(pred, target)
+    tensor(1.3333)
+    >>> l1_loss(pred, target, weight)
+    tensor(1.5000)
+    >>> l1_loss(pred, target, reduction='none')
+    tensor([1., 1., 2.])
+    >>> l1_loss(pred, target, weight, reduction='sum')
+    tensor(3.)
+    """
+
+    @functools.wraps(loss_func)
+    def wrapper(pred, target, weight=None, reduction='mean', **kwargs):
+        # get element-wise loss
+        loss = loss_func(pred, target, **kwargs)
+        loss = weight_reduce_loss(loss, weight, reduction)
+        return loss
+
+    return wrapper
diff --git a/opensora/models/super_resolution/basicsr/losses/losses.py b/opensora/models/super_resolution/basicsr/losses/losses.py
new file mode 100644
index 000000000..55436902e
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/losses/losses.py
@@ -0,0 +1,492 @@
+import math
+import torch
+from torch import autograd as autograd
+from torch import nn as nn
+from torch.nn import functional as F
+
+from basicsr.archs.vgg_arch import VGGFeatureExtractor
+from basicsr.utils.registry import LOSS_REGISTRY
+from .loss_util import weighted_loss
+
+_reduction_modes = ['none', 'mean', 'sum']
+
+
+@weighted_loss
+def l1_loss(pred, target):
+    return F.l1_loss(pred, target, reduction='none')
+
+
+@weighted_loss
+def mse_loss(pred, target):
+    return F.mse_loss(pred, target, reduction='none')
+
+
+@weighted_loss
+def charbonnier_loss(pred, target, eps=1e-12):
+    return torch.sqrt((pred - target)**2 + eps)
+
+
+@LOSS_REGISTRY.register()
+class L1Loss(nn.Module):
+    """L1 (mean absolute error, MAE) loss.
+
+    Args:
+        loss_weight (float): Loss weight for L1 loss. Default: 1.0.
+        reduction (str): Specifies the reduction to apply to the output.
+            Supported choices are 'none' | 'mean' | 'sum'. Default: 'mean'.
+    """
+
+    def __init__(self, loss_weight=1.0, reduction='mean'):
+        super(L1Loss, self).__init__()
+        if reduction not in ['none', 'mean', 'sum']:
+            raise ValueError(f'Unsupported reduction mode: {reduction}. Supported ones are: {_reduction_modes}')
+
+        self.loss_weight = loss_weight
+        self.reduction = reduction
+
+    def forward(self, pred, target, weight=None, **kwargs):
+        """
+        Args:
+            pred (Tensor): of shape (N, C, H, W). Predicted tensor.
+            target (Tensor): of shape (N, C, H, W). Ground truth tensor.
+            weight (Tensor, optional): of shape (N, C, H, W). Element-wise weights. Default: None.
+        """
+        return self.loss_weight * l1_loss(pred, target, weight, reduction=self.reduction)
+
+
+@LOSS_REGISTRY.register()
+class MSELoss(nn.Module):
+    """MSE (L2) loss.
+
+    Args:
+        loss_weight (float): Loss weight for MSE loss. Default: 1.0.
+        reduction (str): Specifies the reduction to apply to the output.
+            Supported choices are 'none' | 'mean' | 'sum'. Default: 'mean'.
+    """
+
+    def __init__(self, loss_weight=1.0, reduction='mean'):
+        super(MSELoss, self).__init__()
+        if reduction not in ['none', 'mean', 'sum']:
+            raise ValueError(f'Unsupported reduction mode: {reduction}. Supported ones are: {_reduction_modes}')
+
+        self.loss_weight = loss_weight
+        self.reduction = reduction
+
+    def forward(self, pred, target, weight=None, **kwargs):
+        """
+        Args:
+            pred (Tensor): of shape (N, C, H, W). Predicted tensor.
+            target (Tensor): of shape (N, C, H, W). Ground truth tensor.
+            weight (Tensor, optional): of shape (N, C, H, W). Element-wise weights. Default: None.
+        """
+        return self.loss_weight * mse_loss(pred, target, weight, reduction=self.reduction)
+
+
+@LOSS_REGISTRY.register()
+class CharbonnierLoss(nn.Module):
+    """Charbonnier loss (one variant of Robust L1Loss, a differentiable
+    variant of L1Loss).
+
+    Described in "Deep Laplacian Pyramid Networks for Fast and Accurate
+        Super-Resolution".
+
+    Args:
+        loss_weight (float): Loss weight for L1 loss. Default: 1.0.
+        reduction (str): Specifies the reduction to apply to the output.
+            Supported choices are 'none' | 'mean' | 'sum'. Default: 'mean'.
+        eps (float): A value used to control the curvature near zero. Default: 1e-12.
+    """
+
+    def __init__(self, loss_weight=1.0, reduction='mean', eps=1e-12):
+        super(CharbonnierLoss, self).__init__()
+        if reduction not in ['none', 'mean', 'sum']:
+            raise ValueError(f'Unsupported reduction mode: {reduction}. Supported ones are: {_reduction_modes}')
+
+        self.loss_weight = loss_weight
+        self.reduction = reduction
+        self.eps = eps
+
+    def forward(self, pred, target, weight=None, **kwargs):
+        """
+        Args:
+            pred (Tensor): of shape (N, C, H, W). Predicted tensor.
+            target (Tensor): of shape (N, C, H, W). Ground truth tensor.
+            weight (Tensor, optional): of shape (N, C, H, W). Element-wise weights. Default: None.
+        """
+        return self.loss_weight * charbonnier_loss(pred, target, weight, eps=self.eps, reduction=self.reduction)
+
+
+@LOSS_REGISTRY.register()
+class WeightedTVLoss(L1Loss):
+    """Weighted TV loss.
+
+    Args:
+        loss_weight (float): Loss weight. Default: 1.0.
+    """
+
+    def __init__(self, loss_weight=1.0, reduction='mean'):
+        if reduction not in ['mean', 'sum']:
+            raise ValueError(f'Unsupported reduction mode: {reduction}. Supported ones are: mean | sum')
+        super(WeightedTVLoss, self).__init__(loss_weight=loss_weight, reduction=reduction)
+
+    def forward(self, pred, weight=None):
+        if weight is None:
+            y_weight = None
+            x_weight = None
+        else:
+            y_weight = weight[:, :, :-1, :]
+            x_weight = weight[:, :, :, :-1]
+
+        y_diff = super().forward(pred[:, :, :-1, :], pred[:, :, 1:, :], weight=y_weight)
+        x_diff = super().forward(pred[:, :, :, :-1], pred[:, :, :, 1:], weight=x_weight)
+
+        loss = x_diff + y_diff
+
+        return loss
+
+
+@LOSS_REGISTRY.register()
+class PerceptualLoss(nn.Module):
+    """Perceptual loss with commonly used style loss.
+
+    Args:
+        layer_weights (dict): The weight for each layer of vgg feature.
+            Here is an example: {'conv5_4': 1.}, which means the conv5_4
+            feature layer (before relu5_4) will be extracted with weight
+            1.0 in calculating losses.
+        vgg_type (str): The type of vgg network used as feature extractor.
+            Default: 'vgg19'.
+        use_input_norm (bool):  If True, normalize the input image in vgg.
+            Default: True.
+        range_norm (bool): If True, norm images with range [-1, 1] to [0, 1].
+            Default: False.
+        perceptual_weight (float): If `perceptual_weight > 0`, the perceptual
+            loss will be calculated and the loss will multiplied by the
+            weight. Default: 1.0.
+        style_weight (float): If `style_weight > 0`, the style loss will be
+            calculated and the loss will multiplied by the weight.
+            Default: 0.
+        criterion (str): Criterion used for perceptual loss. Default: 'l1'.
+    """
+
+    def __init__(self,
+                 layer_weights,
+                 vgg_type='vgg19',
+                 use_input_norm=True,
+                 range_norm=False,
+                 perceptual_weight=1.0,
+                 style_weight=0.,
+                 criterion='l1'):
+        super(PerceptualLoss, self).__init__()
+        self.perceptual_weight = perceptual_weight
+        self.style_weight = style_weight
+        self.layer_weights = layer_weights
+        self.vgg = VGGFeatureExtractor(
+            layer_name_list=list(layer_weights.keys()),
+            vgg_type=vgg_type,
+            use_input_norm=use_input_norm,
+            range_norm=range_norm)
+
+        self.criterion_type = criterion
+        if self.criterion_type == 'l1':
+            self.criterion = torch.nn.L1Loss()
+        elif self.criterion_type == 'l2':
+            self.criterion = torch.nn.L2loss()
+        elif self.criterion_type == 'fro':
+            self.criterion = None
+        else:
+            raise NotImplementedError(f'{criterion} criterion has not been supported.')
+
+    def forward(self, x, gt):
+        """Forward function.
+
+        Args:
+            x (Tensor): Input tensor with shape (n, c, h, w).
+            gt (Tensor): Ground-truth tensor with shape (n, c, h, w).
+
+        Returns:
+            Tensor: Forward results.
+        """
+        # extract vgg features
+        x_features = self.vgg(x)
+        gt_features = self.vgg(gt.detach())
+
+        # calculate perceptual loss
+        if self.perceptual_weight > 0:
+            percep_loss = 0
+            for k in x_features.keys():
+                if self.criterion_type == 'fro':
+                    percep_loss += torch.norm(x_features[k] - gt_features[k], p='fro') * self.layer_weights[k]
+                else:
+                    percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
+            percep_loss *= self.perceptual_weight
+        else:
+            percep_loss = None
+
+        # calculate style loss
+        if self.style_weight > 0:
+            style_loss = 0
+            for k in x_features.keys():
+                if self.criterion_type == 'fro':
+                    style_loss += torch.norm(
+                        self._gram_mat(x_features[k]) - self._gram_mat(gt_features[k]), p='fro') * self.layer_weights[k]
+                else:
+                    style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(
+                        gt_features[k])) * self.layer_weights[k]
+            style_loss *= self.style_weight
+        else:
+            style_loss = None
+
+        return percep_loss, style_loss
+
+    def _gram_mat(self, x):
+        """Calculate Gram matrix.
+
+        Args:
+            x (torch.Tensor): Tensor with shape of (n, c, h, w).
+
+        Returns:
+            torch.Tensor: Gram matrix.
+        """
+        n, c, h, w = x.size()
+        features = x.view(n, c, w * h)
+        features_t = features.transpose(1, 2)
+        gram = features.bmm(features_t) / (c * h * w)
+        return gram
+
+
+@LOSS_REGISTRY.register()
+class GANLoss(nn.Module):
+    """Define GAN loss.
+
+    Args:
+        gan_type (str): Support 'vanilla', 'lsgan', 'wgan', 'hinge'.
+        real_label_val (float): The value for real label. Default: 1.0.
+        fake_label_val (float): The value for fake label. Default: 0.0.
+        loss_weight (float): Loss weight. Default: 1.0.
+            Note that loss_weight is only for generators; and it is always 1.0
+            for discriminators.
+    """
+
+    def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0, loss_weight=1.0):
+        super(GANLoss, self).__init__()
+        self.gan_type = gan_type
+        self.loss_weight = loss_weight
+        self.real_label_val = real_label_val
+        self.fake_label_val = fake_label_val
+
+        if self.gan_type == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif self.gan_type == 'lsgan':
+            self.loss = nn.MSELoss()
+        elif self.gan_type == 'wgan':
+            self.loss = self._wgan_loss
+        elif self.gan_type == 'wgan_softplus':
+            self.loss = self._wgan_softplus_loss
+        elif self.gan_type == 'hinge':
+            self.loss = nn.ReLU()
+        else:
+            raise NotImplementedError(f'GAN type {self.gan_type} is not implemented.')
+
+    def _wgan_loss(self, input, target):
+        """wgan loss.
+
+        Args:
+            input (Tensor): Input tensor.
+            target (bool): Target label.
+
+        Returns:
+            Tensor: wgan loss.
+        """
+        return -input.mean() if target else input.mean()
+
+    def _wgan_softplus_loss(self, input, target):
+        """wgan loss with soft plus. softplus is a smooth approximation to the
+        ReLU function.
+
+        In StyleGAN2, it is called:
+            Logistic loss for discriminator;
+            Non-saturating loss for generator.
+
+        Args:
+            input (Tensor): Input tensor.
+            target (bool): Target label.
+
+        Returns:
+            Tensor: wgan loss.
+        """
+        return F.softplus(-input).mean() if target else F.softplus(input).mean()
+
+    def get_target_label(self, input, target_is_real):
+        """Get target label.
+
+        Args:
+            input (Tensor): Input tensor.
+            target_is_real (bool): Whether the target is real or fake.
+
+        Returns:
+            (bool | Tensor): Target tensor. Return bool for wgan, otherwise,
+                return Tensor.
+        """
+
+        if self.gan_type in ['wgan', 'wgan_softplus']:
+            return target_is_real
+        target_val = (self.real_label_val if target_is_real else self.fake_label_val)
+        return input.new_ones(input.size()) * target_val
+
+    def forward(self, input, target_is_real, is_disc=False):
+        """
+        Args:
+            input (Tensor): The input for the loss module, i.e., the network
+                prediction.
+            target_is_real (bool): Whether the targe is real or fake.
+            is_disc (bool): Whether the loss for discriminators or not.
+                Default: False.
+
+        Returns:
+            Tensor: GAN loss value.
+        """
+        target_label = self.get_target_label(input, target_is_real)
+        if self.gan_type == 'hinge':
+            if is_disc:  # for discriminators in hinge-gan
+                input = -input if target_is_real else input
+                loss = self.loss(1 + input).mean()
+            else:  # for generators in hinge-gan
+                loss = -input.mean()
+        else:  # other gan types
+            loss = self.loss(input, target_label)
+
+        # loss_weight is always 1.0 for discriminators
+        return loss if is_disc else loss * self.loss_weight
+
+
+@LOSS_REGISTRY.register()
+class MultiScaleGANLoss(GANLoss):
+    """
+    MultiScaleGANLoss accepts a list of predictions
+    """
+
+    def __init__(self, gan_type, real_label_val=1.0, fake_label_val=0.0, loss_weight=1.0):
+        super(MultiScaleGANLoss, self).__init__(gan_type, real_label_val, fake_label_val, loss_weight)
+
+    def forward(self, input, target_is_real, is_disc=False):
+        """
+        The input is a list of tensors, or a list of (a list of tensors)
+        """
+        if isinstance(input, list):
+            loss = 0
+            for pred_i in input:
+                if isinstance(pred_i, list):
+                    # Only compute GAN loss for the last layer
+                    # in case of multiscale feature matching
+                    pred_i = pred_i[-1]
+                # Safe operation: 0-dim tensor calling self.mean() does nothing
+                loss_tensor = super().forward(pred_i, target_is_real, is_disc).mean()
+                loss += loss_tensor
+            return loss / len(input)
+        else:
+            return super().forward(input, target_is_real, is_disc)
+
+
+def r1_penalty(real_pred, real_img):
+    """R1 regularization for discriminator. The core idea is to
+        penalize the gradient on real data alone: when the
+        generator distribution produces the true data distribution
+        and the discriminator is equal to 0 on the data manifold, the
+        gradient penalty ensures that the discriminator cannot create
+        a non-zero gradient orthogonal to the data manifold without
+        suffering a loss in the GAN game.
+
+        Ref:
+        Eq. 9 in Which training methods for GANs do actually converge.
+        """
+    grad_real = autograd.grad(outputs=real_pred.sum(), inputs=real_img, create_graph=True)[0]
+    grad_penalty = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()
+    return grad_penalty
+
+
+def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
+    noise = torch.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])
+    grad = autograd.grad(outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True)[0]
+    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))
+
+    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)
+
+    path_penalty = (path_lengths - path_mean).pow(2).mean()
+
+    return path_penalty, path_lengths.detach().mean(), path_mean.detach()
+
+
+def gradient_penalty_loss(discriminator, real_data, fake_data, weight=None):
+    """Calculate gradient penalty for wgan-gp.
+
+    Args:
+        discriminator (nn.Module): Network for the discriminator.
+        real_data (Tensor): Real input data.
+        fake_data (Tensor): Fake input data.
+        weight (Tensor): Weight tensor. Default: None.
+
+    Returns:
+        Tensor: A tensor for gradient penalty.
+    """
+
+    batch_size = real_data.size(0)
+    alpha = real_data.new_tensor(torch.rand(batch_size, 1, 1, 1))
+
+    # interpolate between real_data and fake_data
+    interpolates = alpha * real_data + (1. - alpha) * fake_data
+    interpolates = autograd.Variable(interpolates, requires_grad=True)
+
+    disc_interpolates = discriminator(interpolates)
+    gradients = autograd.grad(
+        outputs=disc_interpolates,
+        inputs=interpolates,
+        grad_outputs=torch.ones_like(disc_interpolates),
+        create_graph=True,
+        retain_graph=True,
+        only_inputs=True)[0]
+
+    if weight is not None:
+        gradients = gradients * weight
+
+    gradients_penalty = ((gradients.norm(2, dim=1) - 1)**2).mean()
+    if weight is not None:
+        gradients_penalty /= torch.mean(weight)
+
+    return gradients_penalty
+
+
+@LOSS_REGISTRY.register()
+class GANFeatLoss(nn.Module):
+    """Define feature matching loss for gans
+
+    Args:
+        criterion (str): Support 'l1', 'l2', 'charbonnier'.
+        loss_weight (float): Loss weight. Default: 1.0.
+        reduction (str): Specifies the reduction to apply to the output.
+            Supported choices are 'none' | 'mean' | 'sum'. Default: 'mean'.
+    """
+
+    def __init__(self, criterion='l1', loss_weight=1.0, reduction='mean'):
+        super(GANFeatLoss, self).__init__()
+        if criterion == 'l1':
+            self.loss_op = L1Loss(loss_weight, reduction)
+        elif criterion == 'l2':
+            self.loss_op = MSELoss(loss_weight, reduction)
+        elif criterion == 'charbonnier':
+            self.loss_op = CharbonnierLoss(loss_weight, reduction)
+        else:
+            raise ValueError(f'Unsupported loss mode: {criterion}. Supported ones are: l1|l2|charbonnier')
+
+        self.loss_weight = loss_weight
+
+    def forward(self, pred_fake, pred_real):
+        num_d = len(pred_fake)
+        loss = 0
+        for i in range(num_d):  # for each discriminator
+            # last output is the final prediction, exclude it
+            num_intermediate_outputs = len(pred_fake[i]) - 1
+            for j in range(num_intermediate_outputs):  # for each layer output
+                unweighted_loss = self.loss_op(pred_fake[i][j], pred_real[i][j].detach())
+                loss += unweighted_loss / num_d
+        return loss * self.loss_weight
diff --git a/opensora/models/super_resolution/basicsr/metrics/__init__.py b/opensora/models/super_resolution/basicsr/metrics/__init__.py
new file mode 100644
index 000000000..c65580aec
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/metrics/__init__.py
@@ -0,0 +1,19 @@
+from copy import deepcopy
+
+from basicsr.utils.registry import METRIC_REGISTRY
+from .psnr_ssim import calculate_psnr, calculate_ssim
+
+__all__ = ['calculate_psnr', 'calculate_ssim']
+
+
+def calculate_metric(data, opt):
+    """Calculate metric from data and options.
+
+    Args:
+        opt (dict): Configuration. It must contain:
+            type (str): Model type.
+    """
+    opt = deepcopy(opt)
+    metric_type = opt.pop('type')
+    metric = METRIC_REGISTRY.get(metric_type)(**data, **opt)
+    return metric
diff --git a/opensora/models/super_resolution/basicsr/metrics/metric_util.py b/opensora/models/super_resolution/basicsr/metrics/metric_util.py
new file mode 100644
index 000000000..0b2177787
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/metrics/metric_util.py
@@ -0,0 +1,45 @@
+import numpy as np
+
+from basicsr.utils.matlab_functions import bgr2ycbcr
+
+
+def reorder_image(img, input_order='HWC'):
+    """Reorder images to 'HWC' order.
+
+    If the input_order is (h, w), return (h, w, 1);
+    If the input_order is (c, h, w), return (h, w, c);
+    If the input_order is (h, w, c), return as it is.
+
+    Args:
+        img (ndarray): Input image.
+        input_order (str): Whether the input order is 'HWC' or 'CHW'.
+            If the input image shape is (h, w), input_order will not have
+            effects. Default: 'HWC'.
+
+    Returns:
+        ndarray: reordered image.
+    """
+
+    if input_order not in ['HWC', 'CHW']:
+        raise ValueError(f"Wrong input_order {input_order}. Supported input_orders are 'HWC' and 'CHW'")
+    if len(img.shape) == 2:
+        img = img[..., None]
+    if input_order == 'CHW':
+        img = img.transpose(1, 2, 0)
+    return img
+
+
+def to_y_channel(img):
+    """Change to Y channel of YCbCr.
+
+    Args:
+        img (ndarray): Images with range [0, 255].
+
+    Returns:
+        (ndarray): Images with range [0, 255] (float type) without round.
+    """
+    img = img.astype(np.float32) / 255.
+    if img.ndim == 3 and img.shape[2] == 3:
+        img = bgr2ycbcr(img, y_only=True)
+        img = img[..., None]
+    return img * 255.
diff --git a/opensora/models/super_resolution/basicsr/metrics/psnr_ssim.py b/opensora/models/super_resolution/basicsr/metrics/psnr_ssim.py
new file mode 100644
index 000000000..cb00426b9
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/metrics/psnr_ssim.py
@@ -0,0 +1,128 @@
+import cv2
+import numpy as np
+
+from basicsr.metrics.metric_util import reorder_image, to_y_channel
+from basicsr.utils.registry import METRIC_REGISTRY
+
+
+@METRIC_REGISTRY.register()
+def calculate_psnr(img, img2, crop_border, input_order='HWC', test_y_channel=False, **kwargs):
+    """Calculate PSNR (Peak Signal-to-Noise Ratio).
+
+    Ref: https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio
+
+    Args:
+        img (ndarray): Images with range [0, 255].
+        img2 (ndarray): Images with range [0, 255].
+        crop_border (int): Cropped pixels in each edge of an image. These
+            pixels are not involved in the PSNR calculation.
+        input_order (str): Whether the input order is 'HWC' or 'CHW'.
+            Default: 'HWC'.
+        test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
+
+    Returns:
+        float: psnr result.
+    """
+
+    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
+    if input_order not in ['HWC', 'CHW']:
+        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are "HWC" and "CHW"')
+    img = reorder_image(img, input_order=input_order)
+    img2 = reorder_image(img2, input_order=input_order)
+    img = img.astype(np.float64)
+    img2 = img2.astype(np.float64)
+
+    if crop_border != 0:
+        img = img[crop_border:-crop_border, crop_border:-crop_border, ...]
+        img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
+
+    if test_y_channel:
+        img = to_y_channel(img)
+        img2 = to_y_channel(img2)
+
+    mse = np.mean((img - img2)**2)
+    if mse == 0:
+        return float('inf')
+    return 20. * np.log10(255. / np.sqrt(mse))
+
+
+def _ssim(img, img2):
+    """Calculate SSIM (structural similarity) for one channel images.
+
+    It is called by func:`calculate_ssim`.
+
+    Args:
+        img (ndarray): Images with range [0, 255] with order 'HWC'.
+        img2 (ndarray): Images with range [0, 255] with order 'HWC'.
+
+    Returns:
+        float: ssim result.
+    """
+
+    c1 = (0.01 * 255)**2
+    c2 = (0.03 * 255)**2
+
+    img = img.astype(np.float64)
+    img2 = img2.astype(np.float64)
+    kernel = cv2.getGaussianKernel(11, 1.5)
+    window = np.outer(kernel, kernel.transpose())
+
+    mu1 = cv2.filter2D(img, -1, window)[5:-5, 5:-5]
+    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+    mu1_sq = mu1**2
+    mu2_sq = mu2**2
+    mu1_mu2 = mu1 * mu2
+    sigma1_sq = cv2.filter2D(img**2, -1, window)[5:-5, 5:-5] - mu1_sq
+    sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+    sigma12 = cv2.filter2D(img * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+    ssim_map = ((2 * mu1_mu2 + c1) * (2 * sigma12 + c2)) / ((mu1_sq + mu2_sq + c1) * (sigma1_sq + sigma2_sq + c2))
+    return ssim_map.mean()
+
+
+@METRIC_REGISTRY.register()
+def calculate_ssim(img, img2, crop_border, input_order='HWC', test_y_channel=False, **kwargs):
+    """Calculate SSIM (structural similarity).
+
+    Ref:
+    Image quality assessment: From error visibility to structural similarity
+
+    The results are the same as that of the official released MATLAB code in
+    https://ece.uwaterloo.ca/~z70wang/research/ssim/.
+
+    For three-channel images, SSIM is calculated for each channel and then
+    averaged.
+
+    Args:
+        img (ndarray): Images with range [0, 255].
+        img2 (ndarray): Images with range [0, 255].
+        crop_border (int): Cropped pixels in each edge of an image. These
+            pixels are not involved in the SSIM calculation.
+        input_order (str): Whether the input order is 'HWC' or 'CHW'.
+            Default: 'HWC'.
+        test_y_channel (bool): Test on Y channel of YCbCr. Default: False.
+
+    Returns:
+        float: ssim result.
+    """
+
+    assert img.shape == img2.shape, (f'Image shapes are different: {img.shape}, {img2.shape}.')
+    if input_order not in ['HWC', 'CHW']:
+        raise ValueError(f'Wrong input_order {input_order}. Supported input_orders are "HWC" and "CHW"')
+    img = reorder_image(img, input_order=input_order)
+    img2 = reorder_image(img2, input_order=input_order)
+    img = img.astype(np.float64)
+    img2 = img2.astype(np.float64)
+
+    if crop_border != 0:
+        img = img[crop_border:-crop_border, crop_border:-crop_border, ...]
+        img2 = img2[crop_border:-crop_border, crop_border:-crop_border, ...]
+
+    if test_y_channel:
+        img = to_y_channel(img)
+        img2 = to_y_channel(img2)
+
+    ssims = []
+    for i in range(img.shape[2]):
+        ssims.append(_ssim(img[..., i], img2[..., i]))
+    return np.array(ssims).mean()
diff --git a/opensora/models/super_resolution/basicsr/models/__init__.py b/opensora/models/super_resolution/basicsr/models/__init__.py
new file mode 100644
index 000000000..285ce3ef9
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/models/__init__.py
@@ -0,0 +1,30 @@
+import importlib
+from copy import deepcopy
+from os import path as osp
+
+from basicsr.utils import get_root_logger, scandir
+from basicsr.utils.registry import MODEL_REGISTRY
+
+__all__ = ['build_model']
+
+# automatically scan and import model modules for registry
+# scan all the files under the 'models' folder and collect files ending with
+# '_model.py'
+model_folder = osp.dirname(osp.abspath(__file__))
+model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
+# import all the model modules
+_model_modules = [importlib.import_module(f'basicsr.models.{file_name}') for file_name in model_filenames]
+
+
+def build_model(opt):
+    """Build model from options.
+
+    Args:
+        opt (dict): Configuration. It must contain:
+            model_type (str): Model type.
+    """
+    opt = deepcopy(opt)
+    model = MODEL_REGISTRY.get(opt['model_type'])(opt)
+    logger = get_root_logger()
+    logger.info(f'Model [{model.__class__.__name__}] is created.')
+    return model
diff --git a/opensora/models/super_resolution/basicsr/models/base_model.py b/opensora/models/super_resolution/basicsr/models/base_model.py
new file mode 100644
index 000000000..f06f9ca2c
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/models/base_model.py
@@ -0,0 +1,380 @@
+import os
+import time
+import torch
+from collections import OrderedDict
+from copy import deepcopy
+from torch.nn.parallel import DataParallel, DistributedDataParallel
+
+from basicsr.models import lr_scheduler as lr_scheduler
+from basicsr.utils import get_root_logger
+from basicsr.utils.dist_util import master_only
+
+
+class BaseModel():
+    """Base model."""
+
+    def __init__(self, opt):
+        self.opt = opt
+        self.device = torch.device('cuda' if opt['num_gpu'] != 0 else 'cpu')
+        self.is_train = opt['is_train']
+        self.schedulers = []
+        self.optimizers = []
+
+    def feed_data(self, data):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        pass
+
+    def save(self, epoch, current_iter):
+        """Save networks and training state."""
+        pass
+
+    def validation(self, dataloader, current_iter, tb_logger, save_img=False):
+        """Validation function.
+
+        Args:
+            dataloader (torch.utils.data.DataLoader): Validation dataloader.
+            current_iter (int): Current iteration.
+            tb_logger (tensorboard logger): Tensorboard logger.
+            save_img (bool): Whether to save images. Default: False.
+        """
+        if self.opt['dist']:
+            self.dist_validation(dataloader, current_iter, tb_logger, save_img)
+        else:
+            self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
+
+    def _initialize_best_metric_results(self, dataset_name):
+        """Initialize the best metric results dict for recording the best metric value and iteration."""
+        if hasattr(self, 'best_metric_results') and dataset_name in self.best_metric_results:
+            return
+        elif not hasattr(self, 'best_metric_results'):
+            self.best_metric_results = dict()
+
+        # add a dataset record
+        record = dict()
+        for metric, content in self.opt['val']['metrics'].items():
+            better = content.get('better', 'higher')
+            init_val = float('-inf') if better == 'higher' else float('inf')
+            record[metric] = dict(better=better, val=init_val, iter=-1)
+        self.best_metric_results[dataset_name] = record
+
+    def _update_best_metric_result(self, dataset_name, metric, val, current_iter):
+        if self.best_metric_results[dataset_name][metric]['better'] == 'higher':
+            if val >= self.best_metric_results[dataset_name][metric]['val']:
+                self.best_metric_results[dataset_name][metric]['val'] = val
+                self.best_metric_results[dataset_name][metric]['iter'] = current_iter
+        else:
+            if val <= self.best_metric_results[dataset_name][metric]['val']:
+                self.best_metric_results[dataset_name][metric]['val'] = val
+                self.best_metric_results[dataset_name][metric]['iter'] = current_iter
+
+    def model_ema(self, decay=0.999):
+        net_g = self.get_bare_model(self.net_g)
+
+        net_g_params = dict(net_g.named_parameters())
+        net_g_ema_params = dict(self.net_g_ema.named_parameters())
+
+        for k in net_g_ema_params.keys():
+            net_g_ema_params[k].data.mul_(decay).add_(net_g_params[k].data, alpha=1 - decay)
+
+    def get_current_log(self):
+        return self.log_dict
+
+    def model_to_device(self, net):
+        """Model to device. It also warps models with DistributedDataParallel
+        or DataParallel.
+
+        Args:
+            net (nn.Module)
+        """
+        net = net.to(self.device)
+        if self.opt['dist']:
+            find_unused_parameters = self.opt.get('find_unused_parameters', False)
+            net = DistributedDataParallel(
+                net, device_ids=[torch.cuda.current_device()], find_unused_parameters=find_unused_parameters)
+        elif self.opt['num_gpu'] > 1:
+            net = DataParallel(net)
+        return net
+
+    def get_optimizer(self, optim_type, params, lr, **kwargs):
+        if optim_type == 'Adam':
+            optimizer = torch.optim.Adam(params, lr, **kwargs)
+        else:
+            raise NotImplementedError(f'optimizer {optim_type} is not supperted yet.')
+        return optimizer
+
+    def setup_schedulers(self):
+        """Set up schedulers."""
+        train_opt = self.opt['train']
+        scheduler_type = train_opt['scheduler'].pop('type')
+        if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']:
+            for optimizer in self.optimizers:
+                self.schedulers.append(lr_scheduler.MultiStepRestartLR(optimizer, **train_opt['scheduler']))
+        elif scheduler_type == 'CosineAnnealingRestartLR':
+            for optimizer in self.optimizers:
+                self.schedulers.append(lr_scheduler.CosineAnnealingRestartLR(optimizer, **train_opt['scheduler']))
+        else:
+            raise NotImplementedError(f'Scheduler {scheduler_type} is not implemented yet.')
+
+    def get_bare_model(self, net):
+        """Get bare model, especially under wrapping with
+        DistributedDataParallel or DataParallel.
+        """
+        if isinstance(net, (DataParallel, DistributedDataParallel)):
+            net = net.module
+        return net
+
+    @master_only
+    def print_network(self, net):
+        """Print the str and parameter number of a network.
+
+        Args:
+            net (nn.Module)
+        """
+        if isinstance(net, (DataParallel, DistributedDataParallel)):
+            net_cls_str = f'{net.__class__.__name__} - {net.module.__class__.__name__}'
+        else:
+            net_cls_str = f'{net.__class__.__name__}'
+
+        net = self.get_bare_model(net)
+        net_str = str(net)
+        net_params = sum(map(lambda x: x.numel(), net.parameters()))
+
+        logger = get_root_logger()
+        logger.info(f'Network: {net_cls_str}, with parameters: {net_params:,d}')
+        logger.info(net_str)
+
+    def _set_lr(self, lr_groups_l):
+        """Set learning rate for warmup.
+
+        Args:
+            lr_groups_l (list): List for lr_groups, each for an optimizer.
+        """
+        for optimizer, lr_groups in zip(self.optimizers, lr_groups_l):
+            for param_group, lr in zip(optimizer.param_groups, lr_groups):
+                param_group['lr'] = lr
+
+    def _get_init_lr(self):
+        """Get the initial lr, which is set by the scheduler.
+        """
+        init_lr_groups_l = []
+        for optimizer in self.optimizers:
+            init_lr_groups_l.append([v['initial_lr'] for v in optimizer.param_groups])
+        return init_lr_groups_l
+
+    def update_learning_rate(self, current_iter, warmup_iter=-1):
+        """Update learning rate.
+
+        Args:
+            current_iter (int): Current iteration.
+            warmup_iter (int)： Warmup iter numbers. -1 for no warmup.
+                Default： -1.
+        """
+        if current_iter > 1:
+            for scheduler in self.schedulers:
+                scheduler.step()
+        # set up warm-up learning rate
+        if current_iter < warmup_iter:
+            # get initial lr for each group
+            init_lr_g_l = self._get_init_lr()
+            # modify warming-up learning rates
+            # currently only support linearly warm up
+            warm_up_lr_l = []
+            for init_lr_g in init_lr_g_l:
+                warm_up_lr_l.append([v / warmup_iter * current_iter for v in init_lr_g])
+            # set learning rate
+            self._set_lr(warm_up_lr_l)
+
+    def get_current_learning_rate(self):
+        return [param_group['lr'] for param_group in self.optimizers[0].param_groups]
+
+    @master_only
+    def save_network(self, net, net_label, current_iter, param_key='params'):
+        """Save networks.
+
+        Args:
+            net (nn.Module | list[nn.Module]): Network(s) to be saved.
+            net_label (str): Network label.
+            current_iter (int): Current iter number.
+            param_key (str | list[str]): The parameter key(s) to save network.
+                Default: 'params'.
+        """
+        if current_iter == -1:
+            current_iter = 'latest'
+        save_filename = f'{net_label}_{current_iter}.pth'
+        save_path = os.path.join(self.opt['path']['models'], save_filename)
+
+        net = net if isinstance(net, list) else [net]
+        param_key = param_key if isinstance(param_key, list) else [param_key]
+        assert len(net) == len(param_key), 'The lengths of net and param_key should be the same.'
+
+        save_dict = {}
+        for net_, param_key_ in zip(net, param_key):
+            net_ = self.get_bare_model(net_)
+            state_dict = net_.state_dict()
+            for key, param in state_dict.items():
+                if key.startswith('module.'):  # remove unnecessary 'module.'
+                    key = key[7:]
+                state_dict[key] = param.cpu()
+            save_dict[param_key_] = state_dict
+
+        # avoid occasional writing errors
+        retry = 3
+        while retry > 0:
+            try:
+                torch.save(save_dict, save_path)
+            except Exception as e:
+                logger = get_root_logger()
+                logger.warning(f'Save model error: {e}, remaining retry times: {retry - 1}')
+                time.sleep(1)
+            else:
+                break
+            finally:
+                retry -= 1
+        if retry == 0:
+            logger.warning(f'Still cannot save {save_path}. Just ignore it.')
+            # raise IOError(f'Cannot save {save_path}.')
+
+    def _print_different_keys_loading(self, crt_net, load_net, strict=True):
+        """Print keys with different name or different size when loading models.
+
+        1. Print keys with different names.
+        2. If strict=False, print the same key but with different tensor size.
+            It also ignore these keys with different sizes (not load).
+
+        Args:
+            crt_net (torch model): Current network.
+            load_net (dict): Loaded network.
+            strict (bool): Whether strictly loaded. Default: True.
+        """
+        crt_net = self.get_bare_model(crt_net)
+        crt_net = crt_net.state_dict()
+        crt_net_keys = set(crt_net.keys())
+        load_net_keys = set(load_net.keys())
+
+        logger = get_root_logger()
+        if crt_net_keys != load_net_keys:
+            logger.warning('Current net - loaded net:')
+            for v in sorted(list(crt_net_keys - load_net_keys)):
+                logger.warning(f'  {v}')
+            logger.warning('Loaded net - current net:')
+            for v in sorted(list(load_net_keys - crt_net_keys)):
+                logger.warning(f'  {v}')
+
+        # check the size for the same keys
+        if not strict:
+            common_keys = crt_net_keys & load_net_keys
+            for k in common_keys:
+                if crt_net[k].size() != load_net[k].size():
+                    logger.warning(f'Size different, ignore [{k}]: crt_net: '
+                                   f'{crt_net[k].shape}; load_net: {load_net[k].shape}')
+                    load_net[k + '.ignore'] = load_net.pop(k)
+
+    def load_network(self, net, load_path, strict=True, param_key='params'):
+        """Load network.
+
+        Args:
+            load_path (str): The path of networks to be loaded.
+            net (nn.Module): Network.
+            strict (bool): Whether strictly loaded.
+            param_key (str): The parameter key of loaded network. If set to
+                None, use the root 'path'.
+                Default: 'params'.
+        """
+        logger = get_root_logger()
+        net = self.get_bare_model(net)
+        load_net = torch.load(load_path, map_location=lambda storage, loc: storage)
+        if param_key is not None:
+            if param_key not in load_net and 'params' in load_net:
+                param_key = 'params'
+                logger.info('Loading: params_ema does not exist, use params.')
+            load_net = load_net[param_key]
+        logger.info(f'Loading {net.__class__.__name__} model from {load_path}, with param key: [{param_key}].')
+        # remove unnecessary 'module.'
+        for k, v in deepcopy(load_net).items():
+            if k.startswith('module.'):
+                load_net[k[7:]] = v
+                load_net.pop(k)
+        self._print_different_keys_loading(net, load_net, strict)
+        net.load_state_dict(load_net, strict=strict)
+
+    @master_only
+    def save_training_state(self, epoch, current_iter):
+        """Save training states during training, which will be used for
+        resuming.
+
+        Args:
+            epoch (int): Current epoch.
+            current_iter (int): Current iteration.
+        """
+        if current_iter != -1:
+            state = {'epoch': epoch, 'iter': current_iter, 'optimizers': [], 'schedulers': []}
+            for o in self.optimizers:
+                state['optimizers'].append(o.state_dict())
+            for s in self.schedulers:
+                state['schedulers'].append(s.state_dict())
+            save_filename = f'{current_iter}.state'
+            save_path = os.path.join(self.opt['path']['training_states'], save_filename)
+
+            # avoid occasional writing errors
+            retry = 3
+            while retry > 0:
+                try:
+                    torch.save(state, save_path)
+                except Exception as e:
+                    logger = get_root_logger()
+                    logger.warning(f'Save training state error: {e}, remaining retry times: {retry - 1}')
+                    time.sleep(1)
+                else:
+                    break
+                finally:
+                    retry -= 1
+            if retry == 0:
+                logger.warning(f'Still cannot save {save_path}. Just ignore it.')
+                # raise IOError(f'Cannot save {save_path}.')
+
+    def resume_training(self, resume_state):
+        """Reload the optimizers and schedulers for resumed training.
+
+        Args:
+            resume_state (dict): Resume state.
+        """
+        resume_optimizers = resume_state['optimizers']
+        resume_schedulers = resume_state['schedulers']
+        assert len(resume_optimizers) == len(self.optimizers), 'Wrong lengths of optimizers'
+        assert len(resume_schedulers) == len(self.schedulers), 'Wrong lengths of schedulers'
+        for i, o in enumerate(resume_optimizers):
+            self.optimizers[i].load_state_dict(o)
+        for i, s in enumerate(resume_schedulers):
+            self.schedulers[i].load_state_dict(s)
+
+    def reduce_loss_dict(self, loss_dict):
+        """reduce loss dict.
+
+        In distributed training, it averages the losses among different GPUs .
+
+        Args:
+            loss_dict (OrderedDict): Loss dict.
+        """
+        with torch.no_grad():
+            if self.opt['dist']:
+                keys = []
+                losses = []
+                for name, value in loss_dict.items():
+                    keys.append(name)
+                    losses.append(value)
+                losses = torch.stack(losses, 0)
+                torch.distributed.reduce(losses, dst=0)
+                if self.opt['rank'] == 0:
+                    losses /= self.opt['world_size']
+                loss_dict = {key: loss for key, loss in zip(keys, losses)}
+
+            log_dict = OrderedDict()
+            for name, value in loss_dict.items():
+                log_dict[name] = value.mean().item()
+
+            return log_dict
diff --git a/opensora/models/super_resolution/basicsr/models/lr_scheduler.py b/opensora/models/super_resolution/basicsr/models/lr_scheduler.py
new file mode 100644
index 000000000..11e1c6c7a
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/models/lr_scheduler.py
@@ -0,0 +1,96 @@
+import math
+from collections import Counter
+from torch.optim.lr_scheduler import _LRScheduler
+
+
+class MultiStepRestartLR(_LRScheduler):
+    """ MultiStep with restarts learning rate scheme.
+
+    Args:
+        optimizer (torch.nn.optimizer): Torch optimizer.
+        milestones (list): Iterations that will decrease learning rate.
+        gamma (float): Decrease ratio. Default: 0.1.
+        restarts (list): Restart iterations. Default: [0].
+        restart_weights (list): Restart weights at each restart iteration.
+            Default: [1].
+        last_epoch (int): Used in _LRScheduler. Default: -1.
+    """
+
+    def __init__(self, optimizer, milestones, gamma=0.1, restarts=(0, ), restart_weights=(1, ), last_epoch=-1):
+        self.milestones = Counter(milestones)
+        self.gamma = gamma
+        self.restarts = restarts
+        self.restart_weights = restart_weights
+        assert len(self.restarts) == len(self.restart_weights), 'restarts and their weights do not match.'
+        super(MultiStepRestartLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch in self.restarts:
+            weight = self.restart_weights[self.restarts.index(self.last_epoch)]
+            return [group['initial_lr'] * weight for group in self.optimizer.param_groups]
+        if self.last_epoch not in self.milestones:
+            return [group['lr'] for group in self.optimizer.param_groups]
+        return [group['lr'] * self.gamma**self.milestones[self.last_epoch] for group in self.optimizer.param_groups]
+
+
+def get_position_from_periods(iteration, cumulative_period):
+    """Get the position from a period list.
+
+    It will return the index of the right-closest number in the period list.
+    For example, the cumulative_period = [100, 200, 300, 400],
+    if iteration == 50, return 0;
+    if iteration == 210, return 2;
+    if iteration == 300, return 2.
+
+    Args:
+        iteration (int): Current iteration.
+        cumulative_period (list[int]): Cumulative period list.
+
+    Returns:
+        int: The position of the right-closest number in the period list.
+    """
+    for i, period in enumerate(cumulative_period):
+        if iteration <= period:
+            return i
+
+
+class CosineAnnealingRestartLR(_LRScheduler):
+    """ Cosine annealing with restarts learning rate scheme.
+
+    An example of config:
+    periods = [10, 10, 10, 10]
+    restart_weights = [1, 0.5, 0.5, 0.5]
+    eta_min=1e-7
+
+    It has four cycles, each has 10 iterations. At 10th, 20th, 30th, the
+    scheduler will restart with the weights in restart_weights.
+
+    Args:
+        optimizer (torch.nn.optimizer): Torch optimizer.
+        periods (list): Period for each cosine anneling cycle.
+        restart_weights (list): Restart weights at each restart iteration.
+            Default: [1].
+        eta_min (float): The minimum lr. Default: 0.
+        last_epoch (int): Used in _LRScheduler. Default: -1.
+    """
+
+    def __init__(self, optimizer, periods, restart_weights=(1, ), eta_min=0, last_epoch=-1):
+        self.periods = periods
+        self.restart_weights = restart_weights
+        self.eta_min = eta_min
+        assert (len(self.periods) == len(
+            self.restart_weights)), 'periods and restart_weights should have the same length.'
+        self.cumulative_period = [sum(self.periods[0:i + 1]) for i in range(0, len(self.periods))]
+        super(CosineAnnealingRestartLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        idx = get_position_from_periods(self.last_epoch, self.cumulative_period)
+        current_weight = self.restart_weights[idx]
+        nearest_restart = 0 if idx == 0 else self.cumulative_period[idx - 1]
+        current_period = self.periods[idx]
+
+        return [
+            self.eta_min + current_weight * 0.5 * (base_lr - self.eta_min) *
+            (1 + math.cos(math.pi * ((self.last_epoch - nearest_restart) / current_period)))
+            for base_lr in self.base_lrs
+        ]
diff --git a/opensora/models/super_resolution/basicsr/models/rgt_model.py b/opensora/models/super_resolution/basicsr/models/rgt_model.py
new file mode 100644
index 000000000..f63b23513
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/models/rgt_model.py
@@ -0,0 +1,127 @@
+import torch
+from torch.nn import functional as F
+
+from basicsr.utils.registry import MODEL_REGISTRY
+from basicsr.models.sr_model import SRModel
+
+
+@MODEL_REGISTRY.register()
+class RGTModel(SRModel):
+
+    def test(self):
+        self.use_chop = self.opt['val']['use_chop'] if 'use_chop' in self.opt['val'] else False
+        if not self.use_chop:
+            if hasattr(self, 'net_g_ema'):
+                self.net_g_ema.eval()
+                with torch.no_grad():
+                    self.output = self.net_g_ema(self.lq)
+            else:
+                self.net_g.eval()
+                with torch.no_grad():
+                    self.output = self.net_g(self.lq)
+                self.net_g.train()
+
+        # test by partitioning
+        else:
+            _, C, h, w = self.lq.size()
+            split_token_h = h // 200 + 1  # number of horizontal cut sections
+            split_token_w = w // 200 + 1  # number of vertical cut sections
+
+            patch_size_tmp_h = split_token_h
+            patch_size_tmp_w = split_token_w
+            
+            # padding
+            mod_pad_h, mod_pad_w = 0, 0
+            if h % patch_size_tmp_h != 0:
+                mod_pad_h = patch_size_tmp_h - h % patch_size_tmp_h
+            if w % patch_size_tmp_w != 0:
+                mod_pad_w = patch_size_tmp_w - w % patch_size_tmp_w
+                
+            img = self.lq
+            img = torch.cat([img, torch.flip(img, [2])], 2)[:, :, :h+mod_pad_h, :]
+            img = torch.cat([img, torch.flip(img, [3])], 3)[:, :, :, :w+mod_pad_w]
+
+            _, _, H, W = img.size()
+            split_h = H // split_token_h  # height of each partition
+            split_w = W // split_token_w  # width of each partition
+
+            # overlapping
+            shave_h = 16
+            shave_w = 16
+            scale = self.opt.get('scale', 1)
+            ral = H // split_h
+            row = W // split_w
+            slices = []  # list of partition borders
+            for i in range(ral):
+                for j in range(row):
+                    if i == 0 and i == ral - 1:
+                        top = slice(i * split_h, (i + 1) * split_h)
+                    elif i == 0:
+                        top = slice(i*split_h, (i+1)*split_h+shave_h)
+                    elif i == ral - 1:
+                        top = slice(i*split_h-shave_h, (i+1)*split_h)
+                    else:
+                        top = slice(i*split_h-shave_h, (i+1)*split_h+shave_h)
+                    if j == 0 and j == row - 1:
+                        left = slice(j*split_w, (j+1)*split_w)
+                    elif j == 0:
+                        left = slice(j*split_w, (j+1)*split_w+shave_w)
+                    elif j == row - 1:
+                        left = slice(j*split_w-shave_w, (j+1)*split_w)
+                    else:
+                        left = slice(j*split_w-shave_w, (j+1)*split_w+shave_w)
+                    temp = (top, left)
+                    slices.append(temp)
+            img_chops = []  # list of partitions
+            for temp in slices:
+                top, left = temp
+                img_chops.append(img[..., top, left])
+            if hasattr(self, 'net_g_ema'):
+                self.net_g_ema.eval()
+                with torch.no_grad():
+                    outputs = []
+                    for chop in img_chops:
+                        out = self.net_g_ema(chop)  # image processing of each partition
+                        outputs.append(out)
+                    _img = torch.zeros(1, C, H * scale, W * scale)
+                    # merge
+                    for i in range(ral):
+                        for j in range(row):
+                            top = slice(i * split_h * scale, (i + 1) * split_h * scale)
+                            left = slice(j * split_w * scale, (j + 1) * split_w * scale)
+                            if i == 0:
+                                _top = slice(0, split_h * scale)
+                            else:
+                                _top = slice(shave_h*scale, (shave_h+split_h)*scale)
+                            if j == 0:
+                                _left = slice(0, split_w*scale)
+                            else:
+                                _left = slice(shave_w*scale, (shave_w+split_w)*scale)
+                            _img[..., top, left] = outputs[i * row + j][..., _top, _left]
+                    self.output = _img
+            else:
+                self.net_g.eval()
+                with torch.no_grad():
+                    outputs = []
+                    for chop in img_chops:
+                        out = self.net_g(chop)  # image processing of each partition
+                        outputs.append(out)
+                    _img = torch.zeros(1, C, H * scale, W * scale)
+                    # merge
+                    for i in range(ral):
+                        for j in range(row):
+                            top = slice(i * split_h * scale, (i + 1) * split_h * scale)
+                            left = slice(j * split_w * scale, (j + 1) * split_w * scale)
+                            if i == 0:
+                                _top = slice(0, split_h * scale)
+                            else:
+                                _top = slice(shave_h * scale, (shave_h + split_h) * scale)
+                            if j == 0:
+                                _left = slice(0, split_w * scale)
+                            else:
+                                _left = slice(shave_w * scale, (shave_w + split_w) * scale)
+                            _img[..., top, left] = outputs[i * row + j][..., _top, _left]
+                    self.output = _img
+                self.net_g.train()
+            _, _, h, w = self.output.size()
+            self.output = self.output[:, :, 0:h - mod_pad_h * scale, 0:w - mod_pad_w * scale]
diff --git a/opensora/models/super_resolution/basicsr/models/sr_model.py b/opensora/models/super_resolution/basicsr/models/sr_model.py
new file mode 100644
index 000000000..e076553c2
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/models/sr_model.py
@@ -0,0 +1,235 @@
+import torch
+from collections import OrderedDict
+from os import path as osp
+from tqdm import tqdm
+
+from basicsr.archs import build_network
+from basicsr.losses import build_loss
+from basicsr.metrics import calculate_metric
+from basicsr.utils import get_root_logger, imwrite, tensor2img
+from basicsr.utils.registry import MODEL_REGISTRY
+from .base_model import BaseModel
+
+
+@MODEL_REGISTRY.register()
+class SRModel(BaseModel):
+    """Base SR model for single image super-resolution."""
+
+    def __init__(self, opt):
+        super(SRModel, self).__init__(opt)
+
+        # define network
+        self.net_g = build_network(opt['network_g'])
+        self.net_g = self.model_to_device(self.net_g)
+        self.print_network(self.net_g)
+
+        # load pretrained models
+        load_path = self.opt['path'].get('pretrain_network_g', None)
+        if load_path is not None:
+            param_key = self.opt['path'].get('param_key_g', 'params')
+            self.load_network(self.net_g, load_path, self.opt['path'].get('strict_load_g', True), param_key)
+
+        if self.is_train:
+            self.init_training_settings()
+
+    def init_training_settings(self):
+        self.net_g.train()
+        train_opt = self.opt['train']
+
+        self.ema_decay = train_opt.get('ema_decay', 0)
+        if self.ema_decay > 0:
+            logger = get_root_logger()
+            logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')
+            # define network net_g with Exponential Moving Average (EMA)
+            # net_g_ema is used only for testing on one GPU and saving
+            # There is no need to wrap with DistributedDataParallel
+            self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
+            # load pretrained model
+            load_path = self.opt['path'].get('pretrain_network_g', None)
+            if load_path is not None:
+                self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
+            else:
+                self.model_ema(0)  # copy net_g weight
+            self.net_g_ema.eval()
+
+        # define losses
+        if train_opt.get('pixel_opt'):
+            self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
+        else:
+            self.cri_pix = None
+
+        if train_opt.get('perceptual_opt'):
+            self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
+        else:
+            self.cri_perceptual = None
+
+        if self.cri_pix is None and self.cri_perceptual is None:
+            raise ValueError('Both pixel and perceptual losses are None.')
+
+        # set up optimizers and schedulers
+        self.setup_optimizers()
+        self.setup_schedulers()
+
+    def setup_optimizers(self):
+        train_opt = self.opt['train']
+        optim_params = []
+        for k, v in self.net_g.named_parameters():
+            if v.requires_grad:
+                optim_params.append(v)
+            else:
+                logger = get_root_logger()
+                logger.warning(f'Params {k} will not be optimized.')
+
+        optim_type = train_opt['optim_g'].pop('type')
+        self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g'])
+        self.optimizers.append(self.optimizer_g)
+
+    def feed_data(self, data):
+        self.lq = data['lq'].to(self.device)
+        if 'gt' in data:
+            self.gt = data['gt'].to(self.device)
+
+    def optimize_parameters(self, current_iter):
+        self.optimizer_g.zero_grad()
+        self.output = self.net_g(self.lq)
+
+        l_total = 0
+        loss_dict = OrderedDict()
+        # pixel loss
+        if self.cri_pix:
+            l_pix = self.cri_pix(self.output, self.gt)
+            l_total += l_pix
+            loss_dict['l_pix'] = l_pix
+        # perceptual loss
+        if self.cri_perceptual:
+            l_percep, l_style = self.cri_perceptual(self.output, self.gt)
+            if l_percep is not None:
+                l_total += l_percep
+                loss_dict['l_percep'] = l_percep
+            if l_style is not None:
+                l_total += l_style
+                loss_dict['l_style'] = l_style
+
+        l_total.backward()
+        self.optimizer_g.step()
+
+        self.log_dict = self.reduce_loss_dict(loss_dict)
+
+        if self.ema_decay > 0:
+            self.model_ema(decay=self.ema_decay)
+
+    def test(self):
+        if hasattr(self, 'net_g_ema'):
+            self.net_g_ema.eval()
+            with torch.no_grad():
+                self.output = self.net_g_ema(self.lq)
+        else:
+            self.net_g.eval()
+            with torch.no_grad():
+                self.output = self.net_g(self.lq)
+            self.net_g.train()
+
+    def dist_validation(self, dataloader, current_iter, tb_logger, save_img):
+        if self.opt['rank'] == 0:
+            self.nondist_validation(dataloader, current_iter, tb_logger, save_img)
+
+    def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
+        dataset_name = dataloader.dataset.opt['name']
+        with_metrics = self.opt['val'].get('metrics') is not None
+        use_pbar = self.opt['val'].get('pbar', False)
+        # print(with_metrics,use_pbar)
+
+        if with_metrics:
+            if not hasattr(self, 'metric_results'):  # only execute in the first run
+                self.metric_results = {metric: 0 for metric in self.opt['val']['metrics'].keys()}
+            # initialize the best metric results for each dataset_name (supporting multiple validation datasets)
+            self._initialize_best_metric_results(dataset_name)
+        # zero self.metric_results
+        if with_metrics:
+            self.metric_results = {metric: 0 for metric in self.metric_results}
+
+        metric_data = dict()
+        if use_pbar:
+            pbar = tqdm(total=len(dataloader), unit='image')
+
+        for idx, val_data in enumerate(dataloader):
+            img_name = osp.splitext(osp.basename(val_data['lq_path'][0]))[0]
+            self.feed_data(val_data)
+            self.test()
+            
+            # this is img data 
+            visuals = self.get_current_visuals()
+            sr_img = tensor2img([visuals['result']])
+
+            metric_data['img'] = sr_img
+            if 'gt' in visuals:
+                gt_img = tensor2img([visuals['gt']])
+                metric_data['img2'] = gt_img
+                del self.gt
+
+            # tentative for out of GPU memory
+            del self.lq
+            del self.output
+            torch.cuda.empty_cache()
+
+            if save_img:
+                if self.opt['is_train']:
+                    save_img_path = osp.join(self.opt['path']['visualization'], img_name,
+                                             f'{img_name}_{current_iter}.png')
+                else:
+                    if self.opt['val']['suffix']:
+                        save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
+                                                 f'{img_name}_{self.opt["val"]["suffix"]}.png')
+                    else:
+                        save_img_path = osp.join(self.opt['path']['visualization'], dataset_name,
+                                                 f'{img_name}_{self.opt["name"]}.png')
+                # save img
+                imwrite(sr_img, save_img_path)
+
+            if with_metrics:
+                # calculate metrics
+                for name, opt_ in self.opt['val']['metrics'].items():
+                    self.metric_results[name] += calculate_metric(metric_data, opt_)
+            if use_pbar:
+                pbar.update(1)
+                pbar.set_description(f'Test {img_name}')
+        if use_pbar:
+            pbar.close()
+
+        if with_metrics:
+            for metric in self.metric_results.keys():
+                self.metric_results[metric] /= (idx + 1)
+                # update the best metric result
+                self._update_best_metric_result(dataset_name, metric, self.metric_results[metric], current_iter)
+
+            self._log_validation_metric_values(current_iter, dataset_name, tb_logger)
+
+    def _log_validation_metric_values(self, current_iter, dataset_name, tb_logger):
+        log_str = f'Validation {dataset_name}\n'
+        for metric, value in self.metric_results.items():
+            log_str += f'\t # {metric}: {value:.4f}'
+            if hasattr(self, 'best_metric_results'):
+                log_str += (f'\tBest: {self.best_metric_results[dataset_name][metric]["val"]:.4f} @ '
+                            f'{self.best_metric_results[dataset_name][metric]["iter"]} iter')
+            log_str += '\n'
+
+        logger = get_root_logger()
+        logger.info(log_str)
+        if tb_logger:
+            for metric, value in self.metric_results.items():
+                tb_logger.add_scalar(f'metrics/{dataset_name}/{metric}', value, current_iter)
+
+    def get_current_visuals(self):
+        out_dict = OrderedDict()
+        out_dict['lq'] = self.lq.detach().cpu()
+        out_dict['result'] = self.output.detach().cpu()
+        if hasattr(self, 'gt'):
+            out_dict['gt'] = self.gt.detach().cpu()
+        return out_dict
+
+    def save(self, epoch, current_iter):
+        if hasattr(self, 'net_g_ema'):
+            self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema'])
+        else:
+            self.save_network(self.net_g, 'net_g', current_iter)
+        self.save_training_state(epoch, current_iter)
diff --git a/opensora/models/super_resolution/basicsr/test_img.py b/opensora/models/super_resolution/basicsr/test_img.py
new file mode 100644
index 000000000..4ac293017
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/test_img.py
@@ -0,0 +1,48 @@
+import logging
+import torch
+from os import path as osp
+from basicsr.data import build_dataloader, build_dataset
+from basicsr.models import build_model
+from basicsr.utils import get_root_logger, get_time_str, make_exp_dirs
+from basicsr.utils.options import dict2str, parse_options
+
+
+def image_sr(args):
+    # parse options, set distributed setting, set ramdom seed
+    opt, _ = parse_options(args.root_path, is_train=False)
+    torch.backends.cudnn.benchmark = True
+    # torch.backends.cudnn.deterministic = True
+
+    # create test dataset and dataloader
+    test_loaders = []
+    for _, dataset_opt in sorted(opt['datasets'].items()):
+        dataset_opt['dataroot_lq'] = osp.join(args.output_dir, f'temp_LR')
+        if args.SR == 'x4': 
+            opt['upscale'] = opt['network_g']['upscale'] = 4
+            opt['val']['suffix'] = 'x4'
+            opt['path']['pretrain_network_g'] = osp.join(args.root_path, f'experiments/pretrained_models/RGT_x4.pth')
+        if args.SR == 'x2': 
+            opt['upscale'] = opt['network_g']['upscale'] = 2
+            opt['val']['suffix'] = 'x2'
+            
+        test_set = build_dataset(dataset_opt)
+        test_loader = build_dataloader(
+            test_set, dataset_opt, num_gpu=opt['num_gpu'], dist=opt['dist'], sampler=None, seed=opt['manual_seed'])
+        test_loaders.append(test_loader)
+
+    opt['path']['pretrain_network_g'] = args.ckpt_path
+    opt['val']['use_chop'] = args.use_chop
+    opt['path']['visualization'] = osp.join(args.output_dir, f'temp_results')
+    opt['path']['results_root'] = osp.join(args.output_dir, f'temp_results')
+
+    # create model
+    model = build_model(opt)
+    for test_loader in test_loaders:
+        test_set_name = test_loader.dataset.opt['name']
+        model.validation(test_loader, current_iter=opt['name'], tb_logger=None, save_img=opt['val']['save_img'])
+
+
+if __name__ == '__main__':
+    root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
+    # print(root_path)
+    # image_sr(root_path)
diff --git a/opensora/models/super_resolution/basicsr/utils/__init__.py b/opensora/models/super_resolution/basicsr/utils/__init__.py
new file mode 100644
index 000000000..4860d5261
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/__init__.py
@@ -0,0 +1,30 @@
+from .file_client import FileClient
+from .img_util import crop_border, imfrombytes, img2tensor, imwrite, tensor2img
+from .logger import AvgTimer, MessageLogger, get_env_info, get_root_logger, init_tb_logger, init_wandb_logger
+from .misc import check_resume, get_time_str, make_exp_dirs, mkdir_and_rename, scandir, set_random_seed, sizeof_fmt
+
+__all__ = [
+    # file_client.py
+    'FileClient',
+    # img_util.py
+    'img2tensor',
+    'tensor2img',
+    'imfrombytes',
+    'imwrite',
+    'crop_border',
+    # logger.py
+    'MessageLogger',
+    'AvgTimer',
+    'init_tb_logger',
+    'init_wandb_logger',
+    'get_root_logger',
+    'get_env_info',
+    # misc.py
+    'set_random_seed',
+    'get_time_str',
+    'mkdir_and_rename',
+    'make_exp_dirs',
+    'scandir',
+    'check_resume',
+    'sizeof_fmt',
+]
diff --git a/opensora/models/super_resolution/basicsr/utils/dist_util.py b/opensora/models/super_resolution/basicsr/utils/dist_util.py
new file mode 100644
index 000000000..0fab887b2
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/dist_util.py
@@ -0,0 +1,82 @@
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E501
+import functools
+import os
+import subprocess
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+
+def init_dist(launcher, backend='nccl', **kwargs):
+    if mp.get_start_method(allow_none=True) is None:
+        mp.set_start_method('spawn')
+    if launcher == 'pytorch':
+        _init_dist_pytorch(backend, **kwargs)
+    elif launcher == 'slurm':
+        _init_dist_slurm(backend, **kwargs)
+    else:
+        raise ValueError(f'Invalid launcher type: {launcher}')
+
+
+def _init_dist_pytorch(backend, **kwargs):
+    rank = int(os.environ['RANK'])
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(rank % num_gpus)
+    dist.init_process_group(backend=backend, **kwargs)
+
+
+def _init_dist_slurm(backend, port=None):
+    """Initialize slurm distributed training environment.
+
+    If argument ``port`` is not specified, then the master port will be system
+    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system
+    environment variable, then a default port ``29500`` will be used.
+
+    Args:
+        backend (str): Backend of torch.distributed.
+        port (int, optional): Master port. Defaults to None.
+    """
+    proc_id = int(os.environ['SLURM_PROCID'])
+    ntasks = int(os.environ['SLURM_NTASKS'])
+    node_list = os.environ['SLURM_NODELIST']
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(proc_id % num_gpus)
+    addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')
+    # specify master port
+    if port is not None:
+        os.environ['MASTER_PORT'] = str(port)
+    elif 'MASTER_PORT' in os.environ:
+        pass  # use MASTER_PORT in the environment variable
+    else:
+        # 29500 is torch.distributed default port
+        os.environ['MASTER_PORT'] = '29500'
+    os.environ['MASTER_ADDR'] = addr
+    os.environ['WORLD_SIZE'] = str(ntasks)
+    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)
+    os.environ['RANK'] = str(proc_id)
+    dist.init_process_group(backend=backend)
+
+
+def get_dist_info():
+    if dist.is_available():
+        initialized = dist.is_initialized()
+    else:
+        initialized = False
+    if initialized:
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    else:
+        rank = 0
+        world_size = 1
+    return rank, world_size
+
+
+def master_only(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        rank, _ = get_dist_info()
+        if rank == 0:
+            return func(*args, **kwargs)
+
+    return wrapper
diff --git a/opensora/models/super_resolution/basicsr/utils/file_client.py b/opensora/models/super_resolution/basicsr/utils/file_client.py
new file mode 100644
index 000000000..89d83ab9e
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/file_client.py
@@ -0,0 +1,167 @@
+# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py  # noqa: E501
+from abc import ABCMeta, abstractmethod
+
+
+class BaseStorageBackend(metaclass=ABCMeta):
+    """Abstract class of storage backends.
+
+    All backends need to implement two apis: ``get()`` and ``get_text()``.
+    ``get()`` reads the file as a byte stream and ``get_text()`` reads the file
+    as texts.
+    """
+
+    @abstractmethod
+    def get(self, filepath):
+        pass
+
+    @abstractmethod
+    def get_text(self, filepath):
+        pass
+
+
+class MemcachedBackend(BaseStorageBackend):
+    """Memcached storage backend.
+
+    Attributes:
+        server_list_cfg (str): Config file for memcached server list.
+        client_cfg (str): Config file for memcached client.
+        sys_path (str | None): Additional path to be appended to `sys.path`.
+            Default: None.
+    """
+
+    def __init__(self, server_list_cfg, client_cfg, sys_path=None):
+        if sys_path is not None:
+            import sys
+            sys.path.append(sys_path)
+        try:
+            import mc
+        except ImportError:
+            raise ImportError('Please install memcached to enable MemcachedBackend.')
+
+        self.server_list_cfg = server_list_cfg
+        self.client_cfg = client_cfg
+        self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg)
+        # mc.pyvector servers as a point which points to a memory cache
+        self._mc_buffer = mc.pyvector()
+
+    def get(self, filepath):
+        filepath = str(filepath)
+        import mc
+        self._client.Get(filepath, self._mc_buffer)
+        value_buf = mc.ConvertBuffer(self._mc_buffer)
+        return value_buf
+
+    def get_text(self, filepath):
+        raise NotImplementedError
+
+
+class HardDiskBackend(BaseStorageBackend):
+    """Raw hard disks storage backend."""
+
+    def get(self, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'rb') as f:
+            value_buf = f.read()
+        return value_buf
+
+    def get_text(self, filepath):
+        filepath = str(filepath)
+        with open(filepath, 'r') as f:
+            value_buf = f.read()
+        return value_buf
+
+
+class LmdbBackend(BaseStorageBackend):
+    """Lmdb storage backend.
+
+    Args:
+        db_paths (str | list[str]): Lmdb database paths.
+        client_keys (str | list[str]): Lmdb client keys. Default: 'default'.
+        readonly (bool, optional): Lmdb environment parameter. If True,
+            disallow any write operations. Default: True.
+        lock (bool, optional): Lmdb environment parameter. If False, when
+            concurrent access occurs, do not lock the database. Default: False.
+        readahead (bool, optional): Lmdb environment parameter. If False,
+            disable the OS filesystem readahead mechanism, which may improve
+            random read performance when a database is larger than RAM.
+            Default: False.
+
+    Attributes:
+        db_paths (list): Lmdb database path.
+        _client (list): A list of several lmdb envs.
+    """
+
+    def __init__(self, db_paths, client_keys='default', readonly=True, lock=False, readahead=False, **kwargs):
+        try:
+            import lmdb
+        except ImportError:
+            raise ImportError('Please install lmdb to enable LmdbBackend.')
+
+        if isinstance(client_keys, str):
+            client_keys = [client_keys]
+
+        if isinstance(db_paths, list):
+            self.db_paths = [str(v) for v in db_paths]
+        elif isinstance(db_paths, str):
+            self.db_paths = [str(db_paths)]
+        assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '
+                                                        f'but received {len(client_keys)} and {len(self.db_paths)}.')
+
+        self._client = {}
+        for client, path in zip(client_keys, self.db_paths):
+            self._client[client] = lmdb.open(path, readonly=readonly, lock=lock, readahead=readahead, **kwargs)
+
+    def get(self, filepath, client_key):
+        """Get values according to the filepath from one lmdb named client_key.
+
+        Args:
+            filepath (str | obj:`Path`): Here, filepath is the lmdb key.
+            client_key (str): Used for distinguishing different lmdb envs.
+        """
+        filepath = str(filepath)
+        assert client_key in self._client, (f'client_key {client_key} is not in lmdb clients.')
+        client = self._client[client_key]
+        with client.begin(write=False) as txn:
+            value_buf = txn.get(filepath.encode('ascii'))
+        return value_buf
+
+    def get_text(self, filepath):
+        raise NotImplementedError
+
+
+class FileClient(object):
+    """A general file client to access files in different backend.
+
+    The client loads a file or text in a specified backend from its path
+    and return it as a binary file. it can also register other backend
+    accessor with a given name and backend class.
+
+    Attributes:
+        backend (str): The storage backend type. Options are "disk",
+            "memcached" and "lmdb".
+        client (:obj:`BaseStorageBackend`): The backend object.
+    """
+
+    _backends = {
+        'disk': HardDiskBackend,
+        'memcached': MemcachedBackend,
+        'lmdb': LmdbBackend,
+    }
+
+    def __init__(self, backend='disk', **kwargs):
+        if backend not in self._backends:
+            raise ValueError(f'Backend {backend} is not supported. Currently supported ones'
+                             f' are {list(self._backends.keys())}')
+        self.backend = backend
+        self.client = self._backends[backend](**kwargs)
+
+    def get(self, filepath, client_key='default'):
+        # client_key is used only for lmdb, where different fileclients have
+        # different lmdb environments.
+        if self.backend == 'lmdb':
+            return self.client.get(filepath, client_key)
+        else:
+            return self.client.get(filepath)
+
+    def get_text(self, filepath):
+        return self.client.get_text(filepath)
diff --git a/opensora/models/super_resolution/basicsr/utils/img_util.py b/opensora/models/super_resolution/basicsr/utils/img_util.py
new file mode 100644
index 000000000..3a5f1da09
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/img_util.py
@@ -0,0 +1,172 @@
+import cv2
+import math
+import numpy as np
+import os
+import torch
+from torchvision.utils import make_grid
+
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)
+
+
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    """Convert torch Tensors into image numpy arrays.
+
+    After clamping to [min, max], values will be normalized to [0, 1].
+
+    Args:
+        tensor (Tensor or list[Tensor]): Accept shapes:
+            1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
+            2) 3D Tensor of shape (3/1 x H x W);
+            3) 2D Tensor of shape (H x W).
+            Tensor channel should be in RGB order.
+        rgb2bgr (bool): Whether to change rgb to bgr.
+        out_type (numpy type): output types. If ``np.uint8``, transform outputs
+            to uint8 type with range [0, 255]; otherwise, float type with
+            range [0, 1]. Default: ``np.uint8``.
+        min_max (tuple[int]): min and max values for clamp.
+
+    Returns:
+        (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
+        shape (H x W). The channel order is BGR.
+    """
+    if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
+
+    if torch.is_tensor(tensor):
+        tensor = [tensor]
+    result = []
+    for _tensor in tensor:
+        _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
+        _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
+
+        n_dim = _tensor.dim()
+        if n_dim == 4:
+            img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if rgb2bgr:
+                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 3:
+            img_np = _tensor.numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if img_np.shape[2] == 1:  # gray image
+                img_np = np.squeeze(img_np, axis=2)
+            else:
+                if rgb2bgr:
+                    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 2:
+            img_np = _tensor.numpy()
+        else:
+            raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
+        if out_type == np.uint8:
+            # Unlike MATLAB, numpy.unit8() WILL NOT round by default.
+            img_np = (img_np * 255.0).round()
+        img_np = img_np.astype(out_type)
+        result.append(img_np)
+    if len(result) == 1:
+        result = result[0]
+    return result
+
+
+def tensor2img_fast(tensor, rgb2bgr=True, min_max=(0, 1)):
+    """This implementation is slightly faster than tensor2img.
+    It now only supports torch tensor with shape (1, c, h, w).
+
+    Args:
+        tensor (Tensor): Now only support torch tensor with (1, c, h, w).
+        rgb2bgr (bool): Whether to change rgb to bgr. Default: True.
+        min_max (tuple[int]): min and max values for clamp.
+    """
+    output = tensor.squeeze(0).detach().clamp_(*min_max).permute(1, 2, 0)
+    output = (output - min_max[0]) / (min_max[1] - min_max[0]) * 255
+    output = output.type(torch.uint8).cpu().numpy()
+    if rgb2bgr:
+        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+    return output
+
+
+def imfrombytes(content, flag='color', float32=False):
+    """Read an image from bytes.
+
+    Args:
+        content (bytes): Image bytes got from files or other streams.
+        flag (str): Flags specifying the color type of a loaded image,
+            candidates are `color`, `grayscale` and `unchanged`.
+        float32 (bool): Whether to change to float32., If True, will also norm
+            to [0, 1]. Default: False.
+
+    Returns:
+        ndarray: Loaded image array.
+    """
+    img_np = np.frombuffer(content, np.uint8)
+    imread_flags = {'color': cv2.IMREAD_COLOR, 'grayscale': cv2.IMREAD_GRAYSCALE, 'unchanged': cv2.IMREAD_UNCHANGED}
+    img = cv2.imdecode(img_np, imread_flags[flag])
+    if float32:
+        img = img.astype(np.float32) / 255.
+    return img
+
+
+def imwrite(img, file_path, params=None, auto_mkdir=True):
+    """Write image to file.
+
+    Args:
+        img (ndarray): Image array to be written.
+        file_path (str): Image file path.
+        params (None or list): Same as opencv's :func:`imwrite` interface.
+        auto_mkdir (bool): If the parent folder of `file_path` does not exist,
+            whether to create it automatically.
+
+    Returns:
+        bool: Successful or not.
+    """
+    if auto_mkdir:
+        dir_name = os.path.abspath(os.path.dirname(file_path))
+        os.makedirs(dir_name, exist_ok=True)
+    ok = cv2.imwrite(file_path, img, params)
+    if not ok:
+        raise IOError('Failed in writing images.')
+
+
+def crop_border(imgs, crop_border):
+    """Crop borders of images.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Images with shape (h, w, c).
+        crop_border (int): Crop border for each end of height and weight.
+
+    Returns:
+        list[ndarray]: Cropped images.
+    """
+    if crop_border == 0:
+        return imgs
+    else:
+        if isinstance(imgs, list):
+            return [v[crop_border:-crop_border, crop_border:-crop_border, ...] for v in imgs]
+        else:
+            return imgs[crop_border:-crop_border, crop_border:-crop_border, ...]
diff --git a/opensora/models/super_resolution/basicsr/utils/logger.py b/opensora/models/super_resolution/basicsr/utils/logger.py
new file mode 100644
index 000000000..73553dc66
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/logger.py
@@ -0,0 +1,213 @@
+import datetime
+import logging
+import time
+
+from .dist_util import get_dist_info, master_only
+
+initialized_logger = {}
+
+
+class AvgTimer():
+
+    def __init__(self, window=200):
+        self.window = window  # average window
+        self.current_time = 0
+        self.total_time = 0
+        self.count = 0
+        self.avg_time = 0
+        self.start()
+
+    def start(self):
+        self.start_time = self.tic = time.time()
+
+    def record(self):
+        self.count += 1
+        self.toc = time.time()
+        self.current_time = self.toc - self.tic
+        self.total_time += self.current_time
+        # calculate average time
+        self.avg_time = self.total_time / self.count
+
+        # reset
+        if self.count > self.window:
+            self.count = 0
+            self.total_time = 0
+
+        self.tic = time.time()
+
+    def get_current_time(self):
+        return self.current_time
+
+    def get_avg_time(self):
+        return self.avg_time
+
+
+class MessageLogger():
+    """Message logger for printing.
+
+    Args:
+        opt (dict): Config. It contains the following keys:
+            name (str): Exp name.
+            logger (dict): Contains 'print_freq' (str) for logger interval.
+            train (dict): Contains 'total_iter' (int) for total iters.
+            use_tb_logger (bool): Use tensorboard logger.
+        start_iter (int): Start iter. Default: 1.
+        tb_logger (obj:`tb_logger`): Tensorboard logger. Default： None.
+    """
+
+    def __init__(self, opt, start_iter=1, tb_logger=None):
+        self.exp_name = opt['name']
+        self.interval = opt['logger']['print_freq']
+        self.start_iter = start_iter
+        self.max_iters = opt['train']['total_iter']
+        self.use_tb_logger = opt['logger']['use_tb_logger']
+        self.tb_logger = tb_logger
+        self.start_time = time.time()
+        self.logger = get_root_logger()
+
+    def reset_start_time(self):
+        self.start_time = time.time()
+
+    @master_only
+    def __call__(self, log_vars):
+        """Format logging message.
+
+        Args:
+            log_vars (dict): It contains the following keys:
+                epoch (int): Epoch number.
+                iter (int): Current iter.
+                lrs (list): List for learning rates.
+
+                time (float): Iter time.
+                data_time (float): Data time for each iter.
+        """
+        # epoch, iter, learning rates
+        epoch = log_vars.pop('epoch')
+        current_iter = log_vars.pop('iter')
+        lrs = log_vars.pop('lrs')
+
+        message = (f'[{self.exp_name[:5]}..][epoch:{epoch:3d}, iter:{current_iter:8,d}, lr:(')
+        for v in lrs:
+            message += f'{v:.3e},'
+        message += ')] '
+
+        # time and estimated time
+        if 'time' in log_vars.keys():
+            iter_time = log_vars.pop('time')
+            data_time = log_vars.pop('data_time')
+
+            total_time = time.time() - self.start_time
+            time_sec_avg = total_time / (current_iter - self.start_iter + 1)
+            eta_sec = time_sec_avg * (self.max_iters - current_iter - 1)
+            eta_str = str(datetime.timedelta(seconds=int(eta_sec)))
+            message += f'[eta: {eta_str}, '
+            message += f'time (data): {iter_time:.3f} ({data_time:.3f})] '
+
+        # other items, especially losses
+        for k, v in log_vars.items():
+            message += f'{k}: {v:.4e} '
+            # tensorboard logger
+            if self.use_tb_logger and 'debug' not in self.exp_name:
+                if k.startswith('l_'):
+                    self.tb_logger.add_scalar(f'losses/{k}', v, current_iter)
+                else:
+                    self.tb_logger.add_scalar(k, v, current_iter)
+        self.logger.info(message)
+
+
+@master_only
+def init_tb_logger(log_dir):
+    from torch.utils.tensorboard import SummaryWriter
+    tb_logger = SummaryWriter(log_dir=log_dir)
+    return tb_logger
+
+
+@master_only
+def init_wandb_logger(opt):
+    """We now only use wandb to sync tensorboard log."""
+    import wandb
+    logger = get_root_logger()
+
+    project = opt['logger']['wandb']['project']
+    resume_id = opt['logger']['wandb'].get('resume_id')
+    if resume_id:
+        wandb_id = resume_id
+        resume = 'allow'
+        logger.warning(f'Resume wandb logger with id={wandb_id}.')
+    else:
+        wandb_id = wandb.util.generate_id()
+        resume = 'never'
+
+    wandb.init(id=wandb_id, resume=resume, name=opt['name'], config=opt, project=project, sync_tensorboard=True)
+
+    logger.info(f'Use wandb logger with id={wandb_id}; project={project}.')
+
+
+def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
+    """Get the root logger.
+
+    The logger will be initialized if it has not been initialized. By default a
+    StreamHandler will be added. If `log_file` is specified, a FileHandler will
+    also be added.
+
+    Args:
+        logger_name (str): root logger name. Default: 'basicsr'.
+        log_file (str | None): The log filename. If specified, a FileHandler
+            will be added to the root logger.
+        log_level (int): The root logger level. Note that only the process of
+            rank 0 is affected, while other processes will set the level to
+            "Error" and be silent most of the time.
+
+    Returns:
+        logging.Logger: The root logger.
+    """
+    logger = logging.getLogger(logger_name)
+    # if the logger has been initialized, just return it
+    if logger_name in initialized_logger:
+        return logger
+
+    format_str = '%(asctime)s %(levelname)s: %(message)s'
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(logging.Formatter(format_str))
+    logger.addHandler(stream_handler)
+    logger.propagate = False
+    rank, _ = get_dist_info()
+    if rank != 0:
+        logger.setLevel('ERROR')
+    elif log_file is not None:
+        logger.setLevel(log_level)
+        # add file handler
+        file_handler = logging.FileHandler(log_file, 'w')
+        file_handler.setFormatter(logging.Formatter(format_str))
+        file_handler.setLevel(log_level)
+        logger.addHandler(file_handler)
+    initialized_logger[logger_name] = True
+    return logger
+
+
+def get_env_info():
+    """Get environment information.
+
+    Currently, only log the software version.
+    """
+    import torch
+    import torchvision
+
+    from basicsr.version import __version__
+    msg = r"""
+                ____                _       _____  ____
+               / __ ) ____ _ _____ (_)_____/ ___/ / __ \
+              / __  |/ __ `// ___// // ___/\__ \ / /_/ /
+             / /_/ // /_/ /(__  )/ // /__ ___/ // _, _/
+            /_____/ \__,_//____//_/ \___//____//_/ |_|
+     ______                   __   __                 __      __
+    / ____/____   ____   ____/ /  / /   __  __ _____ / /__   / /
+   / / __ / __ \ / __ \ / __  /  / /   / / / // ___// //_/  / /
+  / /_/ // /_/ // /_/ // /_/ /  / /___/ /_/ // /__ / /<    /_/
+  \____/ \____/ \____/ \____/  /_____/\____/ \___//_/|_|  (_)
+    """
+    msg += ('\nVersion Information: '
+            f'\n\tBasicSR: {__version__}'
+            f'\n\tPyTorch: {torch.__version__}'
+            f'\n\tTorchVision: {torchvision.__version__}')
+    return msg
diff --git a/opensora/models/super_resolution/basicsr/utils/matlab_functions.py b/opensora/models/super_resolution/basicsr/utils/matlab_functions.py
new file mode 100644
index 000000000..f9f1a83bc
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/matlab_functions.py
@@ -0,0 +1,359 @@
+import math
+import numpy as np
+import torch
+
+
+def cubic(x):
+    """cubic function used for calculate_weights_indices."""
+    absx = torch.abs(x)
+    absx2 = absx**2
+    absx3 = absx**3
+    return (1.5 * absx3 - 2.5 * absx2 + 1) * (
+        (absx <= 1).type_as(absx)) + (-0.5 * absx3 + 2.5 * absx2 - 4 * absx + 2) * (((absx > 1) *
+                                                                                     (absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+    """Calculate weights and indices, used for imresize function.
+
+    Args:
+        in_length (int): Input length.
+        out_length (int): Output length.
+        scale (float): Scale factor.
+        kernel_width (int): Kernel width.
+        antialisaing (bool): Whether to apply anti-aliasing when downsampling.
+    """
+
+    if (scale < 1) and antialiasing:
+        # Use a modified kernel (larger kernel width) to simultaneously
+        # interpolate and antialias
+        kernel_width = kernel_width / scale
+
+    # Output-space coordinates
+    x = torch.linspace(1, out_length, out_length)
+
+    # Input-space coordinates. Calculate the inverse mapping such that 0.5
+    # in output space maps to 0.5 in input space, and 0.5 + scale in output
+    # space maps to 1.5 in input space.
+    u = x / scale + 0.5 * (1 - 1 / scale)
+
+    # What is the left-most pixel that can be involved in the computation?
+    left = torch.floor(u - kernel_width / 2)
+
+    # What is the maximum number of pixels that can be involved in the
+    # computation?  Note: it's OK to use an extra pixel here; if the
+    # corresponding weights are all zero, it will be eliminated at the end
+    # of this function.
+    p = math.ceil(kernel_width) + 2
+
+    # The indices of the input pixels involved in computing the k-th output
+    # pixel are in row k of the indices matrix.
+    indices = left.view(out_length, 1).expand(out_length, p) + torch.linspace(0, p - 1, p).view(1, p).expand(
+        out_length, p)
+
+    # The weights used to compute the k-th output pixel are in row k of the
+    # weights matrix.
+    distance_to_center = u.view(out_length, 1).expand(out_length, p) - indices
+
+    # apply cubic kernel
+    if (scale < 1) and antialiasing:
+        weights = scale * cubic(distance_to_center * scale)
+    else:
+        weights = cubic(distance_to_center)
+
+    # Normalize the weights matrix so that each row sums to 1.
+    weights_sum = torch.sum(weights, 1).view(out_length, 1)
+    weights = weights / weights_sum.expand(out_length, p)
+
+    # If a column in weights is all zero, get rid of it. only consider the
+    # first and last column.
+    weights_zero_tmp = torch.sum((weights == 0), 0)
+    if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 1, p - 2)
+        weights = weights.narrow(1, 1, p - 2)
+    if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+        indices = indices.narrow(1, 0, p - 2)
+        weights = weights.narrow(1, 0, p - 2)
+    weights = weights.contiguous()
+    indices = indices.contiguous()
+    sym_len_s = -indices.min() + 1
+    sym_len_e = indices.max() - in_length
+    indices = indices + sym_len_s - 1
+    return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+@torch.no_grad()
+def imresize(img, scale, antialiasing=True):
+    """imresize function same as MATLAB.
+
+    It now only supports bicubic.
+    The same scale applies for both height and width.
+
+    Args:
+        img (Tensor | Numpy array):
+            Tensor: Input image with shape (c, h, w), [0, 1] range.
+            Numpy: Input image with shape (h, w, c), [0, 1] range.
+        scale (float): Scale factor. The same scale applies for both height
+            and width.
+        antialisaing (bool): Whether to apply anti-aliasing when downsampling.
+            Default: True.
+
+    Returns:
+        Tensor: Output image with shape (c, h, w), [0, 1] range, w/o round.
+    """
+    squeeze_flag = False
+    if type(img).__module__ == np.__name__:  # numpy type
+        numpy_type = True
+        if img.ndim == 2:
+            img = img[:, :, None]
+            squeeze_flag = True
+        img = torch.from_numpy(img.transpose(2, 0, 1)).float()
+    else:
+        numpy_type = False
+        if img.ndim == 2:
+            img = img.unsqueeze(0)
+            squeeze_flag = True
+
+    in_c, in_h, in_w = img.size()
+    out_h, out_w = math.ceil(in_h * scale), math.ceil(in_w * scale)
+    kernel_width = 4
+    kernel = 'cubic'
+
+    # get weights and indices
+    weights_h, indices_h, sym_len_hs, sym_len_he = calculate_weights_indices(in_h, out_h, scale, kernel, kernel_width,
+                                                                             antialiasing)
+    weights_w, indices_w, sym_len_ws, sym_len_we = calculate_weights_indices(in_w, out_w, scale, kernel, kernel_width,
+                                                                             antialiasing)
+    # process H dimension
+    # symmetric copying
+    img_aug = torch.FloatTensor(in_c, in_h + sym_len_hs + sym_len_he, in_w)
+    img_aug.narrow(1, sym_len_hs, in_h).copy_(img)
+
+    sym_patch = img[:, :sym_len_hs, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, 0, sym_len_hs).copy_(sym_patch_inv)
+
+    sym_patch = img[:, -sym_len_he:, :]
+    inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(1, inv_idx)
+    img_aug.narrow(1, sym_len_hs + in_h, sym_len_he).copy_(sym_patch_inv)
+
+    out_1 = torch.FloatTensor(in_c, out_h, in_w)
+    kernel_width = weights_h.size(1)
+    for i in range(out_h):
+        idx = int(indices_h[i][0])
+        for j in range(in_c):
+            out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_h[i])
+
+    # process W dimension
+    # symmetric copying
+    out_1_aug = torch.FloatTensor(in_c, out_h, in_w + sym_len_ws + sym_len_we)
+    out_1_aug.narrow(2, sym_len_ws, in_w).copy_(out_1)
+
+    sym_patch = out_1[:, :, :sym_len_ws]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, 0, sym_len_ws).copy_(sym_patch_inv)
+
+    sym_patch = out_1[:, :, -sym_len_we:]
+    inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+    sym_patch_inv = sym_patch.index_select(2, inv_idx)
+    out_1_aug.narrow(2, sym_len_ws + in_w, sym_len_we).copy_(sym_patch_inv)
+
+    out_2 = torch.FloatTensor(in_c, out_h, out_w)
+    kernel_width = weights_w.size(1)
+    for i in range(out_w):
+        idx = int(indices_w[i][0])
+        for j in range(in_c):
+            out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_w[i])
+
+    if squeeze_flag:
+        out_2 = out_2.squeeze(0)
+    if numpy_type:
+        out_2 = out_2.numpy()
+        if not squeeze_flag:
+            out_2 = out_2.transpose(1, 2, 0)
+
+    return out_2
+
+
+def rgb2ycbcr(img, y_only=False):
+    """Convert a RGB image to YCbCr image.
+
+    This function produces the same results as Matlab's `rgb2ycbcr` function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `RGB <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [65.481, 128.553, 24.966]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def bgr2ycbcr(img, y_only=False):
+    """Convert a BGR image to YCbCr image.
+
+    The bgr version of rgb2ycbcr.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `BGR <-> YCrCb`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+        y_only (bool): Whether to only return Y channel. Default: False.
+
+    Returns:
+        ndarray: The converted YCbCr image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img)
+    if y_only:
+        out_img = np.dot(img, [24.966, 128.553, 65.481]) + 16.0
+    else:
+        out_img = np.matmul(
+            img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786], [65.481, -37.797, 112.0]]) + [16, 128, 128]
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def ycbcr2rgb(img):
+    """Convert a YCbCr image to RGB image.
+
+    This function produces the same results as Matlab's ycbcr2rgb function.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> RGB`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        ndarray: The converted RGB image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+                              [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]  # noqa: E126
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def ycbcr2bgr(img):
+    """Convert a YCbCr image to BGR image.
+
+    The bgr version of ycbcr2rgb.
+    It implements the ITU-R BT.601 conversion for standard-definition
+    television. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion.
+
+    It differs from a similar function in cv2.cvtColor: `YCrCb <-> BGR`.
+    In OpenCV, it implements a JPEG conversion. See more details in
+    https://en.wikipedia.org/wiki/YCbCr#JPEG_conversion.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        ndarray: The converted BGR image. The output image has the same type
+            and range as input image.
+    """
+    img_type = img.dtype
+    img = _convert_input_type_range(img) * 255
+    out_img = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0.00791071, -0.00153632, 0],
+                              [0, -0.00318811, 0.00625893]]) * 255.0 + [-276.836, 135.576, -222.921]  # noqa: E126
+    out_img = _convert_output_type_range(out_img, img_type)
+    return out_img
+
+
+def _convert_input_type_range(img):
+    """Convert the type and range of the input image.
+
+    It converts the input image to np.float32 type and range of [0, 1].
+    It is mainly used for pre-processing the input image in colorspace
+    conversion functions such as rgb2ycbcr and ycbcr2rgb.
+
+    Args:
+        img (ndarray): The input image. It accepts:
+            1. np.uint8 type with range [0, 255];
+            2. np.float32 type with range [0, 1].
+
+    Returns:
+        (ndarray): The converted image with type of np.float32 and range of
+            [0, 1].
+    """
+    img_type = img.dtype
+    img = img.astype(np.float32)
+    if img_type == np.float32:
+        pass
+    elif img_type == np.uint8:
+        img /= 255.
+    else:
+        raise TypeError(f'The img type should be np.float32 or np.uint8, but got {img_type}')
+    return img
+
+
+def _convert_output_type_range(img, dst_type):
+    """Convert the type and range of the image according to dst_type.
+
+    It converts the image to desired type and range. If `dst_type` is np.uint8,
+    images will be converted to np.uint8 type with range [0, 255]. If
+    `dst_type` is np.float32, it converts the image to np.float32 type with
+    range [0, 1].
+    It is mainly used for post-processing images in colorspace conversion
+    functions such as rgb2ycbcr and ycbcr2rgb.
+
+    Args:
+        img (ndarray): The image to be converted with np.float32 type and
+            range [0, 255].
+        dst_type (np.uint8 | np.float32): If dst_type is np.uint8, it
+            converts the image to np.uint8 type with range [0, 255]. If
+            dst_type is np.float32, it converts the image to np.float32 type
+            with range [0, 1].
+
+    Returns:
+        (ndarray): The converted image with desired type and range.
+    """
+    if dst_type not in (np.uint8, np.float32):
+        raise TypeError(f'The dst_type should be np.float32 or np.uint8, but got {dst_type}')
+    if dst_type == np.uint8:
+        img = img.round()
+    else:
+        img /= 255.
+    return img.astype(dst_type)
diff --git a/opensora/models/super_resolution/basicsr/utils/misc.py b/opensora/models/super_resolution/basicsr/utils/misc.py
new file mode 100644
index 000000000..728fef857
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/misc.py
@@ -0,0 +1,141 @@
+import numpy as np
+import os
+import random
+import time
+import torch
+from os import path as osp
+
+from .dist_util import master_only
+
+
+def set_random_seed(seed):
+    """Set random seeds."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def get_time_str():
+    return time.strftime('%Y%m%d_%H%M%S', time.localtime())
+
+
+def mkdir_and_rename(path):
+    """mkdirs. If path exists, rename it with timestamp and create a new one.
+
+    Args:
+        path (str): Folder path.
+    """
+    if osp.exists(path):
+        new_name = path + '_archived_' + get_time_str()
+        print(f'Path already exists. Rename it to {new_name}', flush=True)
+        os.rename(path, new_name)
+    os.makedirs(path, exist_ok=True)
+
+
+@master_only
+def make_exp_dirs(opt):
+    """Make dirs for experiments."""
+    path_opt = opt['path'].copy()
+    if opt['is_train']:
+        mkdir_and_rename(path_opt.pop('experiments_root'))
+    else:
+        mkdir_and_rename(path_opt.pop('results_root'))
+    for key, path in path_opt.items():
+        if ('strict_load' in key) or ('pretrain_network' in key) or ('resume' in key) or ('param_key' in key):
+            continue
+        else:
+            os.makedirs(path, exist_ok=True)
+
+
+def scandir(dir_path, suffix=None, recursive=False, full_path=False):
+    """Scan a directory to find the interested files.
+
+    Args:
+        dir_path (str): Path of the directory.
+        suffix (str | tuple(str), optional): File suffix that we are
+            interested in. Default: None.
+        recursive (bool, optional): If set to True, recursively scan the
+            directory. Default: False.
+        full_path (bool, optional): If set to True, include the dir_path.
+            Default: False.
+
+    Returns:
+        A generator for all the interested files with relative paths.
+    """
+
+    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
+        raise TypeError('"suffix" must be a string or tuple of strings')
+
+    root = dir_path
+
+    def _scandir(dir_path, suffix, recursive):
+        for entry in os.scandir(dir_path):
+            if not entry.name.startswith('.') and entry.is_file():
+                if full_path:
+                    return_path = entry.path
+                else:
+                    return_path = osp.relpath(entry.path, root)
+
+                if suffix is None:
+                    yield return_path
+                elif return_path.endswith(suffix):
+                    yield return_path
+            else:
+                if recursive:
+                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
+                else:
+                    continue
+
+    return _scandir(dir_path, suffix=suffix, recursive=recursive)
+
+
+def check_resume(opt, resume_iter):
+    """Check resume states and pretrain_network paths.
+
+    Args:
+        opt (dict): Options.
+        resume_iter (int): Resume iteration.
+    """
+    if opt['path']['resume_state']:
+        # get all the networks
+        networks = [key for key in opt.keys() if key.startswith('network_')]
+        flag_pretrain = False
+        for network in networks:
+            if opt['path'].get(f'pretrain_{network}') is not None:
+                flag_pretrain = True
+        if flag_pretrain:
+            print('pretrain_network path will be ignored during resuming.')
+        # set pretrained model paths
+        for network in networks:
+            name = f'pretrain_{network}'
+            basename = network.replace('network_', '')
+            if opt['path'].get('ignore_resume_networks') is None or (network
+                                                                     not in opt['path']['ignore_resume_networks']):
+                opt['path'][name] = osp.join(opt['path']['models'], f'net_{basename}_{resume_iter}.pth')
+                print(f"Set {name} to {opt['path'][name]}")
+
+        # change param_key to params in resume
+        param_keys = [key for key in opt['path'].keys() if key.startswith('param_key')]
+        for param_key in param_keys:
+            if opt['path'][param_key] == 'params_ema':
+                opt['path'][param_key] = 'params'
+                print(f'Set {param_key} to params')
+
+
+def sizeof_fmt(size, suffix='B'):
+    """Get human readable file size.
+
+    Args:
+        size (int): File size.
+        suffix (str): Suffix. Default: 'B'.
+
+    Return:
+        str: Formatted file siz.
+    """
+    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
+        if abs(size) < 1024.0:
+            return f'{size:3.1f} {unit}{suffix}'
+        size /= 1024.0
+    return f'{size:3.1f} Y{suffix}'
diff --git a/opensora/models/super_resolution/basicsr/utils/options.py b/opensora/models/super_resolution/basicsr/utils/options.py
new file mode 100644
index 000000000..1925644dd
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/options.py
@@ -0,0 +1,200 @@
+import argparse
+import random
+import torch
+import yaml
+from collections import OrderedDict
+from os import path as osp
+
+from basicsr.utils import set_random_seed
+from basicsr.utils.dist_util import get_dist_info, init_dist, master_only
+
+
+def ordered_yaml():
+    """Support OrderedDict for yaml.
+
+    Returns:
+        yaml Loader and Dumper.
+    """
+    try:
+        from yaml import CDumper as Dumper
+        from yaml import CLoader as Loader
+    except ImportError:
+        from yaml import Dumper, Loader
+
+    _mapping_tag = yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG
+
+    def dict_representer(dumper, data):
+        return dumper.represent_dict(data.items())
+
+    def dict_constructor(loader, node):
+        return OrderedDict(loader.construct_pairs(node))
+
+    Dumper.add_representer(OrderedDict, dict_representer)
+    Loader.add_constructor(_mapping_tag, dict_constructor)
+    return Loader, Dumper
+
+
+def dict2str(opt, indent_level=1):
+    """dict to string for printing options.
+
+    Args:
+        opt (dict): Option dict.
+        indent_level (int): Indent level. Default: 1.
+
+    Return:
+        (str): Option string for printing.
+    """
+    msg = '\n'
+    for k, v in opt.items():
+        if isinstance(v, dict):
+            msg += ' ' * (indent_level * 2) + k + ':['
+            msg += dict2str(v, indent_level + 1)
+            msg += ' ' * (indent_level * 2) + ']\n'
+        else:
+            msg += ' ' * (indent_level * 2) + k + ': ' + str(v) + '\n'
+    return msg
+
+
+def _postprocess_yml_value(value):
+    # None
+    if value == '~' or value.lower() == 'none':
+        return None
+    # bool
+    if value.lower() == 'true':
+        return True
+    elif value.lower() == 'false':
+        return False
+    # !!float number
+    if value.startswith('!!float'):
+        return float(value.replace('!!float', ''))
+    # number
+    if value.isdigit():
+        return int(value)
+    elif value.replace('.', '', 1).isdigit() and value.count('.') < 2:
+        return float(value)
+    # list
+    if value.startswith('['):
+        return eval(value)
+    # str
+    return value
+
+
+def parse_options(root_path, SR, is_train=True):
+    parser = argparse.ArgumentParser()
+    # parser.add_argument('-opt', type=str, default = 'options/test/test_RGT_S_x2.yml',required=True, help='Path to option YAML file.')
+    if SR == 'x4': 
+        file_path = osp.join(root_path,'options/test/test_RGT_x4.yml')
+    if SR == 'x2': 
+        file_path = osp.join(root_path,'options/test/test_RGT_x2.yml')
+    parser.add_argument('-opt', type=str, default = file_path, help='Path to option YAML file.')
+    parser.add_argument('--launcher', choices=['none', 'pytorch', 'slurm'], default='none', help='job launcher')
+    parser.add_argument('--auto_resume', action='store_true')
+    parser.add_argument('--debug', action='store_true')
+    parser.add_argument('--local_rank', type=int, default=0)
+    parser.add_argument(
+        '--force_yml', nargs='+', default=None, help='Force to update yml files. Examples: train:ema_decay=0.999')
+    args = parser.parse_args()
+
+    # parse yml to dict
+    with open(args.opt, mode='r') as f:
+        # print(args.opt)
+        opt = yaml.load(f, Loader=ordered_yaml()[0])
+
+    # distributed settings
+    if args.launcher == 'none':
+        opt['dist'] = False
+        print('Disable distributed.', flush=True)
+    else:
+        opt['dist'] = True
+        if args.launcher == 'slurm' and 'dist_params' in opt:
+            init_dist(args.launcher, **opt['dist_params'])
+        else:
+            init_dist(args.launcher)
+    opt['rank'], opt['world_size'] = get_dist_info()
+
+    # random seed
+    seed = opt.get('manual_seed')
+    if seed is None:
+        seed = random.randint(1, 10000)
+        opt['manual_seed'] = seed
+    set_random_seed(seed + opt['rank'])
+
+    # force to update yml options
+    if args.force_yml is not None:
+        for entry in args.force_yml:
+            # now do not support creating new keys
+            keys, value = entry.split('=')
+            keys, value = keys.strip(), value.strip()
+            value = _postprocess_yml_value(value)
+            eval_str = 'opt'
+            for key in keys.split(':'):
+                eval_str += f'["{key}"]'
+            eval_str += '=value'
+            # using exec function
+            exec(eval_str)
+
+    opt['auto_resume'] = args.auto_resume
+    opt['is_train'] = is_train
+
+    # debug setting
+    if args.debug and not opt['name'].startswith('debug'):
+        opt['name'] = 'debug_' + opt['name']
+
+    if opt['num_gpu'] == 'auto':
+        opt['num_gpu'] = torch.cuda.device_count()
+
+    # datasets
+    for phase, dataset in opt['datasets'].items():
+        # for multiple datasets, e.g., val_1, val_2; test_1, test_2
+        phase = phase.split('_')[0]
+        dataset['phase'] = phase
+        if 'scale' in opt:
+            dataset['scale'] = opt['scale']
+        if dataset.get('dataroot_gt') is not None:
+            dataset['dataroot_gt'] = osp.expanduser(dataset['dataroot_gt'])
+        if dataset.get('dataroot_lq') is not None:
+            dataset['dataroot_lq'] = osp.expanduser(dataset['dataroot_lq'])
+
+    # paths
+    for key, val in opt['path'].items():
+        if (val is not None) and ('resume_state' in key or 'pretrain_network' in key):
+            opt['path'][key] = osp.expanduser(val)
+
+    if is_train:
+        experiments_root = osp.join(root_path, 'experiments', opt['name'])
+        opt['path']['experiments_root'] = experiments_root
+        opt['path']['models'] = osp.join(experiments_root, 'models')
+        opt['path']['training_states'] = osp.join(experiments_root, 'training_states')
+        opt['path']['log'] = experiments_root
+        opt['path']['visualization'] = osp.join(experiments_root, 'visualization')
+
+        # change some options for debug mode
+        if 'debug' in opt['name']:
+            if 'val' in opt:
+                opt['val']['val_freq'] = 8
+            opt['logger']['print_freq'] = 1
+            opt['logger']['save_checkpoint_freq'] = 8
+    else:  # test
+        results_root = osp.join(root_path, 'results', opt['name'])
+        opt['path']['results_root'] = results_root
+        opt['path']['log'] = results_root
+        opt['path']['visualization'] = osp.join(results_root, 'visualization')
+
+    return opt, args
+
+
+@master_only
+def copy_opt_file(opt_file, experiments_root):
+    # copy the yml file to the experiment root
+    import sys
+    import time
+    from shutil import copyfile
+    cmd = ' '.join(sys.argv)
+    filename = osp.join(experiments_root, osp.basename(opt_file))
+    copyfile(opt_file, filename)
+
+    with open(filename, 'r+') as f:
+        lines = f.readlines()
+        lines.insert(0, f'# GENERATE TIME: {time.asctime()}\n# CMD:\n# {cmd}\n\n')
+        f.seek(0)
+        f.writelines(lines)
diff --git a/opensora/models/super_resolution/basicsr/utils/registry.py b/opensora/models/super_resolution/basicsr/utils/registry.py
new file mode 100644
index 000000000..655753b3b
--- /dev/null
+++ b/opensora/models/super_resolution/basicsr/utils/registry.py
@@ -0,0 +1,82 @@
+# Modified from: https://github.com/facebookresearch/fvcore/blob/master/fvcore/common/registry.py  # noqa: E501
+
+
+class Registry():
+    """
+    The registry that provides name -> object mapping, to support third-party
+    users' custom modules.
+
+    To create a registry (e.g. a backbone registry):
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY = Registry('BACKBONE')
+
+    To register an object:
+
+    .. code-block:: python
+
+        @BACKBONE_REGISTRY.register()
+        class MyBackbone():
+            ...
+
+    Or:
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY.register(MyBackbone)
+    """
+
+    def __init__(self, name):
+        """
+        Args:
+            name (str): the name of this registry
+        """
+        self._name = name
+        self._obj_map = {}
+
+    def _do_register(self, name, obj):
+        assert (name not in self._obj_map), (f"An object named '{name}' was already registered "
+                                             f"in '{self._name}' registry!")
+        self._obj_map[name] = obj
+
+    def register(self, obj=None):
+        """
+        Register the given object under the the name `obj.__name__`.
+        Can be used as either a decorator or not.
+        See docstring of this class for usage.
+        """
+        if obj is None:
+            # used as a decorator
+            def deco(func_or_class):
+                name = func_or_class.__name__
+                self._do_register(name, func_or_class)
+                return func_or_class
+
+            return deco
+
+        # used as a function call
+        name = obj.__name__
+        self._do_register(name, obj)
+
+    def get(self, name):
+        ret = self._obj_map.get(name)
+        if ret is None:
+            raise KeyError(f"No object named '{name}' found in '{self._name}' registry!")
+        return ret
+
+    def __contains__(self, name):
+        return name in self._obj_map
+
+    def __iter__(self):
+        return iter(self._obj_map.items())
+
+    def keys(self):
+        return self._obj_map.keys()
+
+
+DATASET_REGISTRY = Registry('dataset')
+ARCH_REGISTRY = Registry('arch')
+MODEL_REGISTRY = Registry('model')
+LOSS_REGISTRY = Registry('loss')
+METRIC_REGISTRY = Registry('metric')
diff --git a/opensora/models/super_resolution/options/test/test_RGT_x2.yml b/opensora/models/super_resolution/options/test/test_RGT_x2.yml
new file mode 100644
index 000000000..4f88a04d6
--- /dev/null
+++ b/opensora/models/super_resolution/options/test/test_RGT_x2.yml
@@ -0,0 +1,94 @@
+# general settings
+name: test_RGT_x2
+model_type: RGTModel
+scale: 2
+num_gpu: 1
+manual_seed: 10
+
+datasets:
+  test_1:  # the 1st test dataset
+    task: SR
+    name: Set5
+    type: PairedImageDataset
+    dataroot_gt: datasets/benchmark/Set5/HR
+    dataroot_lq: datasets/benchmark/Set5/LR_bicubic/X2
+    filename_tmpl: '{}x2'
+    io_backend:
+      type: disk
+
+  # test_2:  # the 2st test dataset
+  #   task: SR
+  #   name: Set14
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/Set14/HR
+  #   dataroot_lq: datasets/benchmark/Set14/LR_bicubic/X2
+  #   filename_tmpl: '{}x2'
+  #   io_backend:
+  #     type: disk
+
+  # test_3:  # the 3st test dataset
+  #   task: SR
+  #   name: B100
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/B100/HR
+  #   dataroot_lq: datasets/benchmark/B100/LR_bicubic/X2
+  #   filename_tmpl: '{}x2'
+  #   io_backend:
+  #     type: disk
+
+  # test_4:  # the 4st test dataset
+  #   task: SR
+  #   name: Urban100
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/Urban100/HR
+  #   dataroot_lq: datasets/benchmark/Urban100/LR_bicubic/X2
+  #   filename_tmpl: '{}x2'
+  #   io_backend:
+  #     type: disk
+
+  # test_5:  # the 5st test dataset
+  #   task: SR
+  #   name: Manga109
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/Manga109/HR
+  #   dataroot_lq: datasets/benchmark/Manga109/LR_bicubic/X2
+  #   filename_tmpl: '{}_LRBI_x2'
+  #   io_backend:
+  #     type: disk
+
+
+# network structures
+network_g:
+  type: RGT
+  upscale: 2
+  in_chans: 3
+  img_size: 64
+  img_range: 1.
+  depth: [6,6,6,6,6,6,6,6]
+  embed_dim: 180
+  num_heads: [6,6,6,6,6,6,6,6]
+  mlp_ratio: 2
+  resi_connection: '1conv'
+  split_size: [8,32]
+  c_ratio: 0.5
+
+# path
+path:
+  pretrain_network_g: /remote-home/lzy/RGT/experiments/pretrained_models/RGT_x2.pth
+  strict_load_g: True
+
+# validation settings
+val:
+  save_img: True
+  suffix: ~  # add suffix to saved images, if None, use exp name
+  use_chop: False  # True to save memory, if img too large
+
+  metrics:
+    psnr: # metric name, can be arbitrary
+      type: calculate_psnr
+      crop_border: 2
+      test_y_channel: True
+    ssim:
+      type: calculate_ssim
+      crop_border: 2
+      test_y_channel: True
\ No newline at end of file
diff --git a/opensora/models/super_resolution/options/test/test_RGT_x4.yml b/opensora/models/super_resolution/options/test/test_RGT_x4.yml
new file mode 100644
index 000000000..a776fa5f8
--- /dev/null
+++ b/opensora/models/super_resolution/options/test/test_RGT_x4.yml
@@ -0,0 +1,94 @@
+# general settings
+name: test_RGT_x4
+model_type: RGTModel
+scale: 4
+num_gpu: 1
+manual_seed: 10
+
+datasets:
+  test_1:  # the 1st test dataset
+    task: SR
+    name: Set5
+    type: PairedImageDataset
+    dataroot_gt: datasets/benchmark/Set5/HR
+    dataroot_lq: datasets/benchmark/Set5/LR_bicubic/X4
+    filename_tmpl: '{}x4'
+    io_backend:
+      type: disk
+
+  # test_2:  # the 2st test dataset
+  #   task: SR
+  #   name: Set14
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/Set14/HR
+  #   dataroot_lq: datasets/benchmark/Set14/LR_bicubic/X4
+  #   filename_tmpl: '{}x4'
+  #   io_backend:
+  #     type: disk
+
+  # test_3:  # the 3st test dataset
+  #   task: SR
+  #   name: B100
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/B100/HR
+  #   dataroot_lq: datasets/benchmark/B100/LR_bicubic/X4
+  #   filename_tmpl: '{}x4'
+  #   io_backend:
+  #     type: disk
+
+  # test_4:  # the 4st test dataset
+  #   task: SR
+  #   name: Urban100
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/Urban100/HR
+  #   dataroot_lq: datasets/benchmark/Urban100/LR_bicubic/X4
+  #   filename_tmpl: '{}x4'
+  #   io_backend:
+  #     type: disk
+
+  # test_5:  # the 5st test dataset
+  #   task: SR
+  #   name: Manga109
+  #   type: PairedImageDataset
+  #   dataroot_gt: datasets/benchmark/Manga109/HR
+  #   dataroot_lq: datasets/benchmark/Manga109/LR_bicubic/X4
+  #   filename_tmpl: '{}_LRBI_x4'
+  #   io_backend:
+  #     type: disk
+
+
+# network structures
+network_g:
+  type: RGT
+  upscale: 4
+  in_chans: 3
+  img_size: 64
+  img_range: 1.
+  depth: [6,6,6,6,6,6,6,6]
+  embed_dim: 180
+  num_heads: [6,6,6,6,6,6,6,6]
+  mlp_ratio: 2
+  resi_connection: '1conv'
+  split_size: [8,32]
+  c_ratio: 0.5
+
+# path
+path:
+  pretrain_network_g: /remote-home/lzy/RGT/experiments/pretrained_models/RGT_x4.pth
+  strict_load_g: True
+
+# validation settings
+val:
+  save_img: True
+  suffix: ~  # add suffix to saved images, if None, use exp name
+  use_chop: False  # True to save memory, if img too large
+
+  metrics:
+    psnr: # metric name, can be arbitrary
+      type: calculate_psnr
+      crop_border: 4
+      test_y_channel: True
+    ssim:
+      type: calculate_ssim
+      crop_border: 4
+      test_y_channel: True
\ No newline at end of file
diff --git a/opensora/models/super_resolution/options/test/test_single_config.yml b/opensora/models/super_resolution/options/test/test_single_config.yml
new file mode 100644
index 000000000..58a0753db
--- /dev/null
+++ b/opensora/models/super_resolution/options/test/test_single_config.yml
@@ -0,0 +1,41 @@
+# general settings
+name: test_single
+model_type: RGTModel
+scale: 2
+num_gpu: 1
+manual_seed: 10
+
+datasets:
+  test_1:  # the 1st test dataset
+    name: Single
+    type: SingleImageDataset
+    dataroot_lq: /test
+    io_backend:
+      type: disk
+
+
+# network structures
+network_g:
+  type: RGT
+  upscale: 2
+  in_chans: 3
+  img_size: 64
+  img_range: 1.
+  depth: [6,6,6,6,6,6,6,6]
+  embed_dim: 180
+  num_heads: [6,6,6,6,6,6,6,6]
+  mlp_ratio: 2
+  resi_connection: '1conv'
+  split_size: [8,32]
+  c_ratio: 0.5
+
+# path
+path:
+  pretrain_network_g: /test
+  strict_load_g: True
+
+# validation settings
+val:
+  save_img: True
+  suffix: ~  # add suffix to saved images, if None, use exp name
+  use_chop: False  # True to save memory, if img too large
\ No newline at end of file
diff --git a/opensora/models/super_resolution/run.py b/opensora/models/super_resolution/run.py
new file mode 100644
index 000000000..ac8e2bfcb
--- /dev/null
+++ b/opensora/models/super_resolution/run.py
@@ -0,0 +1,138 @@
+import cv2
+import argparse
+from basicsr.test_img import image_sr  
+from os import path as osp
+import os
+import shutil
+from PIL import Image
+import re
+import imageio.v2 as imageio
+import threading
+from concurrent.futures import ThreadPoolExecutor
+import time
+
+def replace_filename(original_path, suffix):
+
+    directory = os.path.dirname(original_path)
+    old_filename = os.path.basename(original_path)
+    name_part, file_extension = os.path.splitext(old_filename)
+    new_filename = f"{name_part}{suffix}{file_extension}"
+    new_path = os.path.join(directory, new_filename)
+
+    return new_path
+
+def create_temp_folder(folder_path):
+    
+    if os.path.exists(folder_path):
+        shutil.rmtree(folder_path)
+    os.makedirs(folder_path)
+
+def delete_temp_folder(folder_path):
+    shutil.rmtree(folder_path)
+
+def extract_number(filename):
+    s = re.findall(r'\d+', filename)
+    return int(s[0]) if s else -1
+
+def bicubic_upsample_opencv(input_image_path, output_image_path, scale_factor):
+    
+    img = cv2.imread(input_image_path)
+    
+    original_height, original_width = img.shape[:2]
+    
+    new_width = int(original_width * scale_factor)
+    new_height = int(original_height * scale_factor)
+    
+    upsampled_img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
+    cv2.imwrite(output_image_path, upsampled_img)
+
+
+def process_frame(frame_count, frame, temp_LR_folder_path, temp_HR_folder_path, SR):
+    frame_path = os.path.join(temp_LR_folder_path, f"frame_{frame_count}{SR}.png")
+    cv2.imwrite(frame_path, frame)
+    HR_frame_path = os.path.join(temp_HR_folder_path, f"frame_{frame_count}.png")
+
+    if SR == 'x4': 
+        bicubic_upsample_opencv(frame_path, HR_frame_path, 4)
+    elif SR == 'x2': 
+        bicubic_upsample_opencv(frame_path, HR_frame_path, 2)
+
+def video_sr(args):
+    file_name = os.path.basename(args.input_dir)
+    video_output_path = os.path.join(args.output_dir,file_name)
+
+    if args.SR == 'x4': 
+        temp_LR_folder_path = os.path.join(args.output_dir, f'temp_LR/X4')
+        video_output_path = replace_filename(video_output_path, '_x4')
+        result_temp = osp.join(args.root_path, f'results/test_RGT_x4/visualization/Set5')
+    if args.SR == 'x2': 
+        temp_LR_folder_path = os.path.join(args.output_dir, f'temp_LR/X2')
+        video_output_path = replace_filename(video_output_path, '_x2')
+        result_temp = osp.join(args.root_path, f'results/test_RGT_x2/visualization/Set5')
+    
+    temp_HR_folder_path = os.path.join(args.output_dir, f'temp_HR')
+    
+    # create_temp_folder(result_temp)
+    create_temp_folder(temp_LR_folder_path)
+    create_temp_folder(temp_HR_folder_path) 
+
+    cap = cv2.VideoCapture(args.input_dir)
+    if not cap.isOpened():
+        print("Error opening video file.")
+        return
+
+    t1 = time.time()
+    frame_count = 0
+    frames_to_process = []
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames_to_process.append((frame_count, frame))
+        frame_count += 1
+
+    with ThreadPoolExecutor(max_workers = args.mul_numwork) as executor:
+        for frame_count, frame in frames_to_process:
+            executor.submit(process_frame, frame_count, frame, temp_LR_folder_path, temp_HR_folder_path, args.SR)
+    
+    print("total frames:",frame_count)
+    print("fps :",cap.get(cv2.CAP_PROP_FPS))
+    
+    t2 = time.time()
+    print('mul threads: ',t2 - t1,'s')
+    # progress all frames in video 
+    image_sr(args)
+
+    t3 = time.time()
+    print('image super resolution: ',t3 - t2,'s')
+    # recover video form all frames
+    frame_files = sorted(os.listdir(result_temp), key=extract_number)
+    video_frames = [imageio.imread(os.path.join(result_temp, frame_file)) for frame_file in frame_files]
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    imageio.mimwrite(video_output_path, video_frames, fps=fps, quality=9)  
+    
+    t4 = time.time()
+    print('tranformer frames to  video: ',t4 - t3,'s')
+    # release all resources
+    cap.release()
+    delete_temp_folder(os.path.dirname(temp_LR_folder_path))
+    delete_temp_folder(temp_HR_folder_path)
+    delete_temp_folder(os.path.join(args.root_path, f'results'))
+    
+    t5 = time.time()
+    print('delete time: ',t5 - t4,'s')
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="RGT for Video Super-Resolution")
+    # make sure you SR is match with the ckpt_path
+    parser.add_argument("--SR", type=str, choices=['x2', 'x4'], default='x4', help='image resolution')
+    parser.add_argument("--ckpt_path", type=str, default = "/remote-home/lzy/RGT/experiments/pretrained_models/RGT_x4.pth")
+    
+    parser.add_argument("--root_path", type=str, default = "/remote-home/lzy/RGT")
+    parser.add_argument("--input_dir", type=str, default= "/remote-home/lzy/RGT/datasets/video/video_test1.mp4")
+    parser.add_argument("--output_dir", type=str, default= "/remote-home/lzy/RGT/datasets/video_output")
+    
+    parser.add_argument("--mul_numwork", type=int, default = 16, help ='max_workers to execute Multi')
+    parser.add_argument("--use_chop", type= bool, default = True,  help ='use_chop: True  # True to save memory, if img too large')
+    args = parser.parse_args()
+    video_sr(args)
diff --git a/opensora/models/text_encoder/__init__.py b/opensora/models/text_encoder/__init__.py
new file mode 100644
index 000000000..16428cf4f
--- /dev/null
+++ b/opensora/models/text_encoder/__init__.py
@@ -0,0 +1,76 @@
+import torch
+from torch import nn
+from transformers import CLIPModel, CLIPProcessor
+
+from opensora.utils.utils import get_precision
+
+
+class T5Wrapper(nn.Module):
+    def __init__(self, args, **kwargs):
+        super(T5Wrapper, self).__init__()
+        self.model_name = args.text_encoder_name
+        if 'mt5' in self.model_name:
+            from transformers import MT5EncoderModel
+            self.text_enc = MT5EncoderModel.from_pretrained(self.model_name, cache_dir=args.cache_dir, **kwargs).eval()
+            # self.text_enc = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir, **kwargs).eval()
+            
+        elif 't5' in self.model_name:
+            from transformers import T5EncoderModel
+            self.text_enc = T5EncoderModel.from_pretrained(self.model_name, cache_dir=args.cache_dir, **kwargs).eval()
+            # self.text_enc = T5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/models--DeepFloyd--t5-v1_1-xxl/snapshots/c9c625d2ec93667ec579ede125fd3811d1f81d37", cache_dir=args.cache_dir, **kwargs).eval()
+
+    def forward(self, input_ids, attention_mask):
+        text_encoder_embs = self.text_enc(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
+        return text_encoder_embs.detach()
+
+class CLIPWrapper(nn.Module):
+    def __init__(self, args):
+        super(CLIPWrapper, self).__init__()
+        self.model_name = args.text_encoder_name
+        dtype = get_precision(args)
+        model_kwargs = {'cache_dir': args.cache_dir, 'low_cpu_mem_usage': True, 'torch_dtype': dtype}
+        self.text_enc = CLIPModel.from_pretrained(self.model_name, **model_kwargs).eval()
+
+    def forward(self, prompts): # prompts = ['I am a test caption', 'Test twice']
+        # Process the prompts and move to the correct device
+        inputs = self.clip_processor(text=prompts, return_tensors="pt", padding=True, truncation=True).to(self.device)
+        
+        # Generate text features
+        with torch.no_grad():
+            text_features = self.clip_model.get_text_features(**inputs)
+        
+        return text_features.detach()
+
+
+
+text_encoder = {
+    'google/mt5-xl': T5Wrapper,
+    'google/mt5-xxl': T5Wrapper,
+    'google/umt5-xl': T5Wrapper,
+    'google/umt5-xxl': T5Wrapper,
+    'DeepFloyd/t5-v1_1-xxl': T5Wrapper,
+    'clip': CLIPWrapper
+}
+
+
+def get_text_enc(args):
+    """deprecation"""
+    encoder_key = None
+    for key in text_encoder.keys():
+        if key in args.text_encoder_name:
+            encoder_key = key
+            break
+    text_enc = text_encoder.get(encoder_key, None)
+    assert text_enc is not None
+    return text_enc(args)
+
+def get_text_warpper(text_encoder_name):
+    """deprecation"""
+    encoder_key = None
+    for key in text_encoder.keys():
+        if key in text_encoder_name:
+            encoder_key = key
+            break
+    text_enc = text_encoder.get(encoder_key, None)
+    assert text_enc is not None
+    return text_enc
diff --git a/opensora/models/text_encoder/clip.py b/opensora/models/text_encoder/clip.py
new file mode 100644
index 000000000..11bf4041e
--- /dev/null
+++ b/opensora/models/text_encoder/clip.py
@@ -0,0 +1,222 @@
+# -*- coding: utf-8 -*-
+import os
+import re
+import ftfy
+import torch
+import html
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer, CLIPTextModel
+
+class CLIPEmbedder:
+    """
+    A class for embedding texts and images using a pretrained CLIP model.
+    """
+    
+    def __init__(self, device='cuda', model_name='openai/clip-vit-base-patch32', cache_dir='./cache_dir', use_text_preprocessing=True, max_length=77):
+        """
+        Initializes the CLIPEmbedder with specified model and configurations.
+        """
+        self.device = torch.device(device)
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.use_text_preprocessing = use_text_preprocessing
+        self.max_length = max_length
+        
+        os.makedirs(self.cache_dir, exist_ok=True)
+        
+        self.processor = CLIPProcessor.from_pretrained(model_name, cache_dir=self.cache_dir)
+        self.model = CLIPModel.from_pretrained(model_name, cache_dir=self.cache_dir).to(self.device).eval()
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
+        self.text_model = CLIPTextModel.from_pretrained(model_name, cache_dir=self.cache_dir).to(self.device).eval()
+        
+        for param in self.text_model.parameters():
+            param.requires_grad = False
+
+    def get_text_embeddings(self, texts):
+        """
+        Generates embeddings for a list of text prompts.
+        """
+        self._validate_input_list(texts, str)
+        
+        if self.use_text_preprocessing:
+            texts = [self._clean_text(text) for text in texts]
+        
+        inputs = self.processor(text=texts, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length).to(self.device)
+        
+        with torch.no_grad():
+            embeddings = self.model.get_text_features(**inputs)
+        
+        return embeddings
+
+    def encode_text(self, texts):
+        """
+        Encodes texts into embeddings and returns the last hidden state and pooled output.
+        """
+        self._validate_input_list(texts, str)
+        
+        batch_encoding = self.tokenizer(texts, return_tensors="pt", truncation=True, max_length=self.max_length, padding="max_length").to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.text_model(**batch_encoding)
+        
+        return outputs.last_hidden_state, outputs.pooler_output
+
+    def get_image_embeddings(self, image_paths):
+        """
+        Generates embeddings for a list of image file paths.
+        """
+        self._validate_input_list(image_paths, str)
+        images = [self._load_image(path) for path in image_paths]
+        
+        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
+        
+        with torch.no_grad():
+            embeddings = self.model.get_image_features(**inputs)
+        
+        return embeddings
+
+    def _validate_input_list(self, input_list, expected_type):
+        """
+        Validates that the input is a list of expected type.
+        """
+        if not isinstance(input_list, list) or not all(isinstance(item, expected_type) for item in input_list):
+            raise ValueError(f"Input must be a list of {expected_type.__name__}.")
+
+    def _clean_text(self, text):
+        """
+        Applies basic cleaning and formatting to a text string.
+        """
+        text = ftfy.fix_text(text)
+        text = html.unescape(text)
+        return text.strip()
+
+    def _load_image(self, image_path):
+        """
+        Loads and preprocesses an image from a file path.
+        """
+        try:
+            image = Image.open(image_path).convert("RGB")
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        except Exception as e:
+            raise Exception(f"Error loading image {image_path}: {e}")
+        return image
+
+    def clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub('<person>', 'person', caption)
+        # urls:
+        caption = re.sub(
+            r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+            '', caption)  # regex for urls
+        caption = re.sub(
+            r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+            '', caption)  # regex for urls
+
+        caption = BeautifulSoup(caption, features='html.parser').text
+
+
+        caption = re.sub(r'@[\w\d]+\b', '', caption)
+
+        caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
+        caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
+        caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
+        caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
+        caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
+        caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
+        caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
+
+        caption = re.sub(
+            r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
+            '-', caption)
+
+
+        caption = re.sub(r'[`´«»“”¨]', '"', caption)
+        caption = re.sub(r'[‘’]', "'", caption)
+
+
+        caption = re.sub(r'&quot;?', '', caption)
+
+        caption = re.sub(r'&amp', '', caption)
+
+
+        caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
+
+
+        caption = re.sub(r'\d:\d\d\s+$', '', caption)
+
+
+        caption = re.sub(r'\\n', ' ', caption)
+
+
+        caption = re.sub(r'#\d{1,3}\b', '', caption)
+
+        caption = re.sub(r'#\d{5,}\b', '', caption)
+        caption = re.sub(r'\b\d{6,}\b', '', caption)
+        caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
+        caption = re.sub(r'[\"\']{2,}', r'"', caption)  
+        caption = re.sub(r'[\.]{2,}', r' ', caption)  
+
+        caption = re.sub(self.bad_punct_regex, r' ', caption)  
+        caption = re.sub(r'\s+\.\s+', r' ', caption)  
+        regex2 = re.compile(r'(?:\-|\_)')
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, ' ', caption)
+        caption = self.basic_clean(caption)
+        caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
+        caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
+        caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
+
+        caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
+        caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
+        caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
+        caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
+        caption = re.sub(r'\bpage\s+\d+\b', '', caption)
+
+        caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
+
+        caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
+
+        caption = re.sub(r'\b\s+\:\s+', r': ', caption)
+        caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
+        caption = re.sub(r'\s+', ' ', caption)
+
+        caption.strip()
+
+        caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
+        caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
+        caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
+        caption = re.sub(r'^\.\S+$', '', caption)
+
+        return caption.strip()
+
+    @staticmethod
+    def basic_clean(text):
+        text = ftfy.fix_text(text)
+        text = html.unescape(html.unescape(text))
+        return text.strip()
+
+if __name__ == '__main__':
+
+    clip_embedder = CLIPEmbedder()
+
+    # Example
+    text_prompts = [
+        'A photo of a cute puppy playing with a ball.',
+        'An image of a beautiful sunset over the ocean.',
+        'A scene depicting a busy city street.'
+    ]
+    text_embeddings = clip_embedder.get_text_embeddings(text_prompts)
+    print(f"Text embeddings shape: {text_embeddings.shape}")
+
+    image_paths = ['image1.jpg', 'image2.png']
+    try:
+        image_embeddings = clip_embedder.get_image_embeddings(image_paths)
+        print(f"Image embeddings shape: {image_embeddings.shape}")
+    except FileNotFoundError as e:
+        print(e)
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
diff --git a/opensora/models/text_encoder/t5.py b/opensora/models/text_encoder/t5.py
new file mode 100644
index 000000000..cbc782072
--- /dev/null
+++ b/opensora/models/text_encoder/t5.py
@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+import os
+import re
+import html
+import urllib.parse as ul
+
+import ftfy
+import torch
+from bs4 import BeautifulSoup
+from transformers import T5EncoderModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+class T5Embedder:
+
+    available_models = ['t5-v1_1-xxl']
+    bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
+
+    def __init__(self, device, dir_or_name='t5-v1_1-xxl', *, cache_dir='./cache_dir', hf_token=None, use_text_preprocessing=True,
+                 t5_model_kwargs=None, torch_dtype=None, model_max_length=120):
+        self.device = torch.device(device)
+        self.torch_dtype = torch_dtype or torch.bfloat16
+        if t5_model_kwargs is None:
+            t5_model_kwargs = {'low_cpu_mem_usage': True, 'torch_dtype': self.torch_dtype}
+            t5_model_kwargs['device_map'] = {'shared': self.device, 'encoder': self.device}
+
+        self.use_text_preprocessing = use_text_preprocessing
+        self.hf_token = hf_token
+        self.cache_dir = cache_dir
+        self.dir_or_name = dir_or_name
+        cache_dir = os.path.join(self.cache_dir, 't5-v1_1-xxl')
+        for filename in ['config.json', 'special_tokens_map.json', 'spiece.model', 'tokenizer_config.json',
+                         'pytorch_model-00001-of-00002.bin', 'pytorch_model-00002-of-00002.bin', 'pytorch_model.bin.index.json']:
+            hf_hub_download(repo_id='DeepFloyd/t5-v1_1-xxl', filename=filename, cache_dir=cache_dir,
+                            force_filename=filename, token=self.hf_token)
+
+        print(cache_dir)
+        self.tokenizer = AutoTokenizer.from_pretrained(cache_dir)
+        self.model = T5EncoderModel.from_pretrained(cache_dir, **t5_model_kwargs).eval()
+        self.model_max_length = model_max_length
+
+    def get_text_embeddings(self, texts):
+        texts = [self.text_preprocessing(text) for text in texts]
+
+        text_tokens_and_mask = self.tokenizer(
+            texts,
+            max_length=self.model_max_length,
+            padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors='pt'
+        )
+
+        text_tokens_and_mask['input_ids'] = text_tokens_and_mask['input_ids']
+        text_tokens_and_mask['attention_mask'] = text_tokens_and_mask['attention_mask']
+
+        with torch.no_grad():
+            text_encoder_embs = self.model(
+                input_ids=text_tokens_and_mask['input_ids'].to(self.device),
+                attention_mask=text_tokens_and_mask['attention_mask'].to(self.device),
+            )['last_hidden_state'].detach()
+        return text_encoder_embs, text_tokens_and_mask['attention_mask'].to(self.device)
+
+    def text_preprocessing(self, text):
+        if self.use_text_preprocessing:
+            # The exact text cleaning as was in the training stage:
+            text = self.clean_caption(text)
+            text = self.clean_caption(text)
+            return text
+        else:
+            return text.lower().strip()
+
+    @staticmethod
+    def basic_clean(text):
+        text = ftfy.fix_text(text)
+        text = html.unescape(html.unescape(text))
+        return text.strip()
+
+    def clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub('<person>', 'person', caption)
+        # urls:
+        caption = re.sub(
+            r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+            '', caption)  # regex for urls
+        caption = re.sub(
+            r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+            '', caption)  # regex for urls
+        # html:
+        caption = BeautifulSoup(caption, features='html.parser').text
+
+        # @<nickname>
+        caption = re.sub(r'@[\w\d]+\b', '', caption)
+
+        # 31C0—31EF CJK Strokes
+        # 31F0—31FF Katakana Phonetic Extensions
+        # 3200—32FF Enclosed CJK Letters and Months
+        # 3300—33FF CJK Compatibility
+        # 3400—4DBF CJK Unified Ideographs Extension A
+        # 4DC0—4DFF Yijing Hexagram Symbols
+        # 4E00—9FFF CJK Unified Ideographs
+        caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
+        caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
+        caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
+        caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
+        caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
+        caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
+        caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
+        #######################################################
+
+        # все виды тире / all types of dash --> "-"
+        caption = re.sub(
+            r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
+            '-', caption)
+
+        # кавычки к одному стандарту
+        caption = re.sub(r'[`´«»“”¨]', '"', caption)
+        caption = re.sub(r'[‘’]', "'", caption)
+
+        # &quot;
+        caption = re.sub(r'&quot;?', '', caption)
+        # &amp
+        caption = re.sub(r'&amp', '', caption)
+
+        # ip adresses:
+        caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
+
+        # article ids:
+        caption = re.sub(r'\d:\d\d\s+$', '', caption)
+
+        # \n
+        caption = re.sub(r'\\n', ' ', caption)
+
+        # "#123"
+        caption = re.sub(r'#\d{1,3}\b', '', caption)
+        # "#12345.."
+        caption = re.sub(r'#\d{5,}\b', '', caption)
+        # "123456.."
+        caption = re.sub(r'\b\d{6,}\b', '', caption)
+        # filenames:
+        caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
+
+        #
+        caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
+        caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
+
+        caption = re.sub(self.bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+        caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
+
+        # this-is-my-cute-cat / this_is_my_cute_cat
+        regex2 = re.compile(r'(?:\-|\_)')
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, ' ', caption)
+
+        caption = self.basic_clean(caption)
+
+        caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
+        caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
+        caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
+
+        caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
+        caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
+        caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
+        caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
+        caption = re.sub(r'\bpage\s+\d+\b', '', caption)
+
+        caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
+
+        caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
+
+        caption = re.sub(r'\b\s+\:\s+', r': ', caption)
+        caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
+        caption = re.sub(r'\s+', ' ', caption)
+
+        caption.strip()
+
+        caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
+        caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
+        caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
+        caption = re.sub(r'^\.\S+$', '', caption)
+
+        return caption.strip()
+
+if __name__ == '__main__':
+    t5 = T5Embedder(device="cuda", cache_dir='./cache_dir', torch_dtype=torch.float)
+    device = t5.device
+    prompts = ['I am a test caption', 'Test twice']
+    with torch.no_grad():
+        caption_embs, emb_masks = t5.get_text_embeddings(prompts)
+        emb_dict = {
+            'caption_feature': caption_embs.float().cpu().data.numpy(),
+            'attention_mask': emb_masks.cpu().data.numpy(),
+        }
+    import ipdb;ipdb.set_trace()
+    print()
\ No newline at end of file
diff --git a/opensora/npu_config.py b/opensora/npu_config.py
new file mode 100644
index 000000000..68b27f7a8
--- /dev/null
+++ b/opensora/npu_config.py
@@ -0,0 +1,442 @@
+import math
+import mmap
+import os
+import pickle
+import random
+import numpy as np
+import torch
+import subprocess
+import sys
+import threading
+import gc
+import torch.distributed as dist
+import re
+from opensora.adaptor.zp_manager import zp_manager
+
+try:
+    import torch_npu
+
+    npu_is_available = True
+    from torch_npu.contrib import transfer_to_npu
+except:
+    npu_is_available = False
+
+from contextlib import contextmanager
+import types
+
+time = 0
+def compress_video(input_file, output_file, out_size):
+    """使用 ffmpeg 压缩视频文件。"""
+    command = [
+        'ffmpeg',
+        '-i', input_file,
+        '-vf', f"scale='min({out_size},iw)':'min({out_size},ih)':force_original_aspect_ratio=decrease",
+        '-c:v', 'libx264',
+        '-crf', '18',
+        '-preset', 'slow',
+        '-c:a', 'copy',
+        output_file
+    ]
+    subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+
+@contextmanager
+def set_run_dtype(x, dtype=None):
+    # 保存原始环境变量的值（如果存在）
+    npu_config.original_run_dtype = x.dtype
+    # 设置环境变量为指定的值
+    npu_config.current_run_dtype = dtype
+    try:
+        # Yield control back to the body of the `with` statement
+        yield
+    finally:
+        # 恢复原始的环境变量值
+        npu_config.current_run_dtype = None
+        npu_config.original_run_dtype = None
+
+
+class NPUConfig:
+    N_NPU_PER_NODE = 8
+
+    def __init__(self):
+        self.on_npu = npu_is_available
+        self.node_world_size = self.N_NPU_PER_NODE
+        self.profiling = False
+        self.profiling_step = 5
+        self.enable_FA = True
+        self.enable_FP32 = False
+        self.load_pickle = True
+        self.use_small_dataset = False
+        self.current_run_dtype = None
+        self.original_run_dtype = None
+        self.zp_manager = zp_manager
+        self.replaced_type = torch.float32
+        self.conv_dtype = torch.float16
+        self.agent = None
+        self.sparse = None
+        if self.enable_FA and self.enable_FP32:
+            self.inf_float = -10000.0
+        else:
+            self.inf_float = -10000.0
+
+        if self.use_small_dataset:
+            self.load_pickle = False
+
+        self._loss = []
+        self.work_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+        self.pickle_save_path = f"{self.work_path}/pickles"
+        self.mm = dict()
+
+        if self.on_npu:
+            import deepspeed
+            import sys
+            torch_npu.npu.set_compile_mode(jit_compile=False)
+
+            import deepspeed.runtime.utils as utils
+            from opensora.adaptor.utils import all_gather_dp_groups, all_gather_into_tensor_dp_groups
+            utils.all_gather_dp_groups = all_gather_dp_groups
+
+            import deepspeed.runtime.bf16_optimizer as bf16_optimizer
+            from opensora.adaptor.bf16_optimizer import BF16_Optimizer
+            self.replace_methods(bf16_optimizer.BF16_Optimizer, BF16_Optimizer)
+
+            from opensora.adaptor.stage_1_and_2 import DeepSpeedZeroOptimizer
+            import deepspeed.runtime.zero.stage_1_and_2 as stage_1_and_2
+            self.replace_methods(stage_1_and_2.DeepSpeedZeroOptimizer, DeepSpeedZeroOptimizer, ['_has_inf_or_nan'])
+
+            import deepspeed.runtime.engine as engine
+            from opensora.adaptor.engine import DeepSpeedEngine
+            self.replace_methods(engine.DeepSpeedEngine, DeepSpeedEngine, skip_fcns=['__init__', '_copy_recovery_script', '_change_recovery_script_permissions'])
+
+        if "RANK" in os.environ:
+            self.rank = int(os.environ["RANK"])
+            self.world_size = int(os.environ["WORLD_SIZE"])
+            torch_npu.npu.set_device(self.get_local_rank())
+        else:
+            self.rank = torch.cuda.current_device()
+            self.world_size = self.N_NPU_PER_NODE
+        self.print_with_rank(f"The npu_config.on_npu is {self.on_npu}")
+        self.bind_thread_to_cpu()
+        gc.set_threshold(700, 10, 10000)
+
+    def get_tokens(self, input_string):
+        input_string = input_string.strip('[]')
+        items = input_string.split(', ')
+        tokens = {}
+        for item in items:
+            range_part, value_part = item.split(': ')
+            start, end = map(int, range_part.split('-'))
+            if value_part == "None":
+                value = None
+            else:
+                values = tuple(map(int, re.findall(r'\d+', value_part)))
+                if len(values) == 6:
+                    value = (values[:3], values[3:]) # conv3d
+                else:
+                    value = values  # adaptive_avg_pool
+            for i in range(start, end):
+                tokens[i] = value
+                if self.rank == 0:
+                    print(i, value)
+        return tokens
+
+
+    def get_total_cores(self):
+        try:
+            total_cores = os.sysconf('SC_NPROCESSORS_ONLN')
+        except (AttributeError, ValueError):
+            total_cores = os.cpu_count()
+        return total_cores
+
+
+    def bind_thread_to_cpu(self):
+        total_cores = self.get_total_cores()
+        # 每个卡的核心数量
+        cores_per_rank = total_cores // 8
+        # 计算本地rank
+        local_rank = self.rank % 8
+        # 计算当前 rank 的 CPU 核范围
+        start_core = local_rank * cores_per_rank
+        end_core = start_core + cores_per_rank - 1
+        # 构建 CPU 核范围字符串
+        cpu_cores_range = f"{start_core}-{end_core}"
+        pid = os.getpid()
+        command = f"taskset -cp {cpu_cores_range} {pid}"
+
+        subprocess.run(command, shell=True, check=True)
+        return f"Binding Cores:{self.rank}:{pid}:{cpu_cores_range}"
+
+    def replace_methods(self, target_class, source_class, skip_fcns=[], only_include_fcns=None):
+        for attr_name in dir(source_class):
+            attr_value = getattr(source_class, attr_name)
+            if attr_name in source_class.__dict__:
+                attr_class_value = source_class.__dict__[attr_name]
+            else:
+                attr_class_value = attr_value
+            if (isinstance(attr_class_value, staticmethod) or isinstance(attr_class_value, classmethod)
+                    or attr_name in skip_fcns):
+                print(f"skip replace {attr_name}")
+                continue
+
+            if only_include_fcns is not None and attr_name not in only_include_fcns:
+                continue
+
+            elif isinstance(attr_value, types.FunctionType):
+                setattr(target_class, attr_name, attr_value)
+
+    def get_attention_mask(self, attention_mask, repeat_num):
+        if self.on_npu and attention_mask is not None:
+            if npu_config.enable_FA:
+                attention_mask = attention_mask.to(torch.bool)
+            attention_mask = attention_mask.repeat(1, repeat_num, 1)
+        return attention_mask
+    def set_current_run_dtype(self, variables):
+        if variables[0].dtype != self.current_run_dtype and self.current_run_dtype is not None:
+            for index, var in enumerate(variables):
+                variables[index] = var.to(self.current_run_dtype)
+        return tuple(variables)
+
+    def restore_dtype(self, x):
+        if x.dtype != self.original_run_dtype and self.original_run_dtype is not None:
+            x = x.to(self.original_run_dtype)
+        return x
+
+    def get_output_video_path(self, name):
+        os.makedirs(f"{self.work_path}/output_videos", exist_ok=True)
+        return f"{self.work_path}/output_videos/{name}"
+
+    def get_node_id(self):
+        return self.rank // self.node_world_size
+
+    def get_node_size(self):
+        return self.world_size // self.node_world_size
+
+    def get_local_rank(self):
+        return self.rank % self.N_NPU_PER_NODE
+
+    def get_pickle_path(self, file_name):
+        return f"{self.pickle_save_path}/{file_name}_local_n63"
+
+    def free_mm(self):
+        for key, value in self.mm.items():
+            value.close()
+        self.mm.clear()
+
+    def __del__(self):
+        self.free_mm()
+
+    def try_load_pickle(self, file_name, function):
+        file_name = self.get_pickle_path(file_name)
+        if os.path.exists(file_name) and self.load_pickle:
+            with open(file_name, 'rb') as file:
+                # self.mm[file_name] = mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)
+                # # 使用 mmap 进行数据读取
+                # loaded_data = pickle.loads(self.mm[file_name][:])
+                loaded_data = pickle.load(file)
+                return loaded_data
+        else:
+            data = function()
+            if not self.use_small_dataset:
+                if self.rank % self.N_NPU_PER_NODE == 0:
+                    # 只需要rank0保存文件
+                    os.makedirs(self.pickle_save_path, exist_ok=True)
+                    with open(file_name, 'wb') as file:
+                        pickle.dump(data, file, pickle.HIGHEST_PROTOCOL)
+            return data
+
+    def try_get_vid_path(self, file, out_size=1024):
+        output_file = file.rsplit(".", 1)[0] + f"_resize{out_size}.mp4"
+        if not os.path.exists(output_file):
+            return file
+        #     compress_video(file, output_file, out_size)
+        return output_file
+
+    def npu_format_cast(self, x):
+        return torch_npu.npu_format_cast(x, 2)
+
+    def calc_grad_norm(self, model):
+        # 计算并打印梯度范数
+        # model_engine = accelerator.deepspeed_engine_wrapped.engine
+        # gradients = model_engine.get_gradients()
+        # grad_norm = get_grad_norm(gradients)
+        # 计算并打印梯度范数
+        grad_norm = 0
+        n_grad = 0
+        # for name, param in model.named_parameters():
+        #     grad_data = deepspeed.utils.safe_get_full_grad(param)
+        #     # self.print_tensor_stats(grad_data, name=name)
+        #
+        #     if grad_data is not None:
+        #         param_norm = grad_data.norm(2)
+        #         grad_norm += param_norm.item() ** 2
+        #         n_grad += 1
+        # grad_norm = (grad_norm / n_grad) ** (1. / 2)
+
+        return grad_norm
+
+    def _run(self, operator, x, tmp_dtype, out_dtype=None, out_nd_format=False):
+        if self.on_npu:
+            if out_dtype is None:
+                out_dtype = x.dtype
+
+            with torch.cuda.amp.autocast(enabled=False):
+                x = operator.to(device=x.device, dtype=tmp_dtype)(x.to(tmp_dtype))
+                x = x.to(out_dtype)
+                if out_nd_format:
+                    return self.npu_format_cast(x)
+                else:
+                    return x
+        else:
+            return operator(x)
+
+    def run_group_norm(self, operator, x):
+        return self._run(operator, x, torch.float32)
+
+    def run_layer_norm(self, operator, x):
+        return self._run(operator, x, torch.float32)
+
+    def print_tensor_stats(self, tensor, name="Tensor", rank=None):
+        if rank and rank != self.rank:
+            return
+
+        if tensor is None:
+            self.print_msg(f"Tensor {name} is None.")
+            return
+
+        x_dtype = tensor.dtype
+        tensor = tensor.to(torch.bfloat16)
+        max_val = tensor.max().item()
+        min_val = tensor.min().item()
+        abs_max_val = min(abs(max_val), abs(min_val))
+        mean_val = tensor.mean().item()
+        median_val = tensor.median().item()
+        std_val = tensor.std().item()
+        shape = tensor.shape
+        self.print_msg(
+            f"{name} - Max: {max_val}, Min: {min_val}, Mean: {mean_val}, AbsMax: {abs_max_val},"
+            f"Median: {median_val}, Std: {std_val}, Shape: {shape}, Type: {x_dtype}")
+
+    def run_conv3d(self, operator, x, out_dtype):
+        return self._run(operator, x, self.conv_dtype, out_dtype, out_nd_format=True)
+
+    def run_pool_2d(self, operator, x):
+        return self._run(operator, x, self.replaced_type)
+
+    def run_pad_2d(self, operator, x, pad, mode="constant"):
+        if self.on_npu:
+            x_dtype = x.dtype
+            x = x.to(self.replaced_type)
+            x = operator(x, pad, mode)
+            x = x.to(x_dtype)
+        else:
+            x = operator(x, pad, mode)
+        return x
+
+    def seed_everything(self, seed=100):
+        seed += self.rank
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    def print_with_rank(self, msg, rank=0, save=False):
+        if self.rank == rank:
+            print(f"{msg}", flush=True)
+            if save:
+                self._loss.append(msg)
+
+    def print_msg(self, msg, on=True, rank=None):
+        if on:
+            if self.rank == rank or rank is None:
+                print(f"[RANK-{self.rank}]: {msg}", flush=True)
+
+    def save_loss(self, filename, rank=0):
+        if self.rank == rank:
+            import json
+            with open(filename, 'w') as file:
+                json.dump(self._loss, file, indent=4)
+
+    def run_attention(self, query, key, value, atten_mask, input_layout, head_dim, head_num):
+        if self.enable_FA:
+            hidden_states = torch_npu.npu_fusion_attention(query, key, value,
+                                                           atten_mask=atten_mask,
+                                                           input_layout=input_layout,
+                                                           scale=1 / math.sqrt(head_dim),
+                                                           head_num=head_num)[0]
+        else:
+            hidden_states = self.scaled_dot_product_attention(query, key, value,
+                                                              atten_mask=atten_mask,
+                                                              input_layout=input_layout,
+                                                              scale=1 / math.sqrt(head_dim),
+                                                              head_num=head_num)
+        return hidden_states
+
+    def save_weight(self, weight):
+        if weight.shape[-1] == weight.shape[-2]:
+            global time
+            filename = f"weight_/attn_weight_{time}.pt"
+            torch.save(weight, filename)
+            time += 1
+
+    def scaled_dot_product_attention(self, query, key, value, input_layout, head_num=None,
+                                     atten_mask=None, scale=None, dropout_p=0.0, is_causal=False) -> torch.Tensor:
+        # L, S = query.size(-2), key.size(-2)
+        def trans_tensor_shape(x, layout, head_num):
+            if layout == "BSH":
+                batch = x.shape[0]
+                x = x.view(batch, -1, head_num, x.shape[-1] // head_num).transpose(1, 2).contiguous()
+            elif layout == "SBH":
+                batch = x.shape[1]
+                x = x.view(-1, batch * head_num, x.shape[-1] // head_num).transpose(0, 1).contiguous()
+                x = x.view(batch, head_num, -1, x.shape[-1])
+            return x
+
+        query = trans_tensor_shape(query, input_layout, head_num)
+        key = trans_tensor_shape(key, input_layout, head_num)
+        value = trans_tensor_shape(value, input_layout, head_num)
+
+        attn_weight = query @ key.transpose(-2, -1) * scale
+        attn_bias = torch.zeros_like(attn_weight, dtype=query.dtype, device=query.device)
+        if is_causal:
+            assert atten_mask is None
+            temp_mask = torch.zeros_like(attn_weight, dtype=torch.bool, device=query.device).tril(diagonal=0)
+            attn_bias.masked_fill_(temp_mask.logical_not(), npu_config.inf_float)
+            attn_bias.to(query.dtype)
+
+        if atten_mask is not None:
+            assert (not self.enable_FA) and atten_mask.dtype != torch.bool, \
+                "attention_mask must not be bool type when use this function"
+
+        attn_weight += attn_bias
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+        self.save_weight(attn_weight)
+        attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+        output = attn_weight @ value
+        if input_layout == "BSH":
+            output = output.transpose(1, 2).contiguous().view(output.shape[0], -1, head_num * output.shape[-1])
+        else:
+            output = output.view(output.shape[0] * head_num, -1, output.shape[-1]).transpose(0, 1).contiguous()
+            output = output.view(output.shape[0], -1, head_num * output.shape[-1])
+        return output
+
+    def print_tensor_with_rank(self, name, tensor, rank=[0], dim_print_cnt=[]):
+        if type(rank) is not list:
+            rank = [rank]
+        if self.rank in rank:
+            def print_dim(tensor_, indices):
+                if tensor_.dim() == len(indices):
+                    return '{0:10.5f} '.format(tensor[tuple(indices)].detach().item())
+                else:
+                    cur_dim = len(indices)
+                    ret = ''
+                    for x in range(0, tensor_.size(cur_dim), tensor_.size(cur_dim) // dim_print_cnt[cur_dim]):
+                        ret += print_dim(tensor_, indices + [x])
+                    return ret + '\n'
+
+            print(name, tensor.size(), self.rank, '\n', print_dim(tensor, []))
+
+
+npu_config = NPUConfig()
diff --git a/opensora/sample/pipeline_inpaint.py b/opensora/sample/pipeline_inpaint.py
new file mode 100644
index 000000000..32464b26c
--- /dev/null
+++ b/opensora/sample/pipeline_inpaint.py
@@ -0,0 +1,904 @@
+# Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch.nn.functional as F
+import html
+import inspect
+import re
+import urllib.parse as ul
+from typing import Callable, List, Optional, Tuple, Union
+import math
+import torch
+from transformers import T5EncoderModel, T5Tokenizer
+
+from diffusers.models import AutoencoderKL, Transformer2DModel
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.utils import (
+    BACKENDS_MAPPING,
+    deprecate,
+    is_bs4_available,
+    is_ftfy_available,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+
+from einops import rearrange
+
+try:
+    from opensora.npu_config import npu_config
+except:
+    npu_config = None
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+if is_bs4_available():
+    from bs4 import BeautifulSoup
+
+if is_ftfy_available():
+    import ftfy
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import PixArtSigmaPipeline
+
+        >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
+        >>> pipe = PixArtSigmaPipeline.from_pretrained(
+        ...     "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
+        ... )
+        >>> # Enable memory optimizations.
+        >>> # pipe.enable_model_cpu_offload()
+
+        >>> prompt = "A small cactus with a happy face in the Sahara desert."
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+                must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class OpenSoraInpaintPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using PixArt-Sigma.
+    """
+
+    bad_punct_regex = re.compile(
+        r"["
+        + "#®•©™&@·º½¾¿¡§~"
+        + r"\)"
+        + r"\("
+        + r"\]"
+        + r"\["
+        + r"\}"
+        + r"\{"
+        + r"\|"
+        + "\\"
+        + r"\/"
+        + r"\*"
+        + r"]{1,}"
+    )  # noqa
+
+    _optional_components = ["tokenizer", "text_encoder"]
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKL,
+        transformer: Transformer2DModel,
+        scheduler: DPMSolverMultistepScheduler,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
+        )
+
+        # self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        clean_caption: bool = False,
+        max_sequence_length: int = 120,
+        **kwargs,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
+                instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
+                PixArt-Alpha, this should be "".
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                whether to use classifier free guidance or not
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                number of images that should be generated per prompt
+            device: (`torch.device`, *optional*):
+                torch device to place the resulting embeddings on
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
+                string.
+            clean_caption (`bool`, defaults to `False`):
+                If `True`, the function will preprocess and clean the provided caption before encoding.
+            max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt.
+        """
+
+        if "mask_feature" in kwargs:
+            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
+            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
+
+        if device is None:
+            device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # See Section 3.1. of the paper.
+        max_length = max_sequence_length
+
+        if prompt_embeds is None:
+            prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {max_length} tokens: {removed_text}"
+                )
+
+            prompt_attention_mask = text_inputs.attention_mask
+            prompt_attention_mask = prompt_attention_mask.to(device)
+
+            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            dtype = self.text_encoder.dtype
+        elif self.transformer is not None:
+            dtype = self.transformer.dtype
+        else:
+            dtype = None
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens = [negative_prompt] * batch_size
+            uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            negative_prompt_attention_mask = uncond_input.attention_mask
+            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        num_frames, 
+        height,
+        width,
+        negative_prompt,
+        callback_steps,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if num_frames <= 0:
+            raise ValueError(f"`num_frames` have to be positive but is {num_frames}.")
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
+    def _text_preprocessing(self, text, clean_caption=False):
+        if clean_caption and not is_bs4_available():
+            logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if clean_caption and not is_ftfy_available():
+            logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if not isinstance(text, (tuple, list)):
+            text = [text]
+
+        def process(text: str):
+            if clean_caption:
+                text = self._clean_caption(text)
+                text = self._clean_caption(text)
+            else:
+                text = text.lower().strip()
+            return text
+
+        return [process(t) for t in text]
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
+    def _clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub("<person>", "person", caption)
+        # urls:
+        caption = re.sub(
+            r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        caption = re.sub(
+            r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        # html:
+        caption = BeautifulSoup(caption, features="html.parser").text
+
+        # @<nickname>
+        caption = re.sub(r"@[\w\d]+\b", "", caption)
+
+        # 31C0—31EF CJK Strokes
+        # 31F0—31FF Katakana Phonetic Extensions
+        # 3200—32FF Enclosed CJK Letters and Months
+        # 3300—33FF CJK Compatibility
+        # 3400—4DBF CJK Unified Ideographs Extension A
+        # 4DC0—4DFF Yijing Hexagram Symbols
+        # 4E00—9FFF CJK Unified Ideographs
+        caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
+        caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
+        caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
+        caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
+        caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
+        caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
+        # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
+        #######################################################
+
+        # все виды тире / all types of dash --> "-"
+        caption = re.sub(
+            r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",  # noqa
+            "-",
+            caption,
+        )
+
+        # кавычки к одному стандарту
+        caption = re.sub(r"[`´«»“”¨]", '"', caption)
+        caption = re.sub(r"[‘’]", "'", caption)
+
+        # &quot;
+        caption = re.sub(r"&quot;?", "", caption)
+        # &amp
+        caption = re.sub(r"&amp", "", caption)
+
+        # ip adresses:
+        caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
+
+        # article ids:
+        caption = re.sub(r"\d:\d\d\s+$", "", caption)
+
+        # \n
+        caption = re.sub(r"\\n", " ", caption)
+
+        # "#123"
+        caption = re.sub(r"#\d{1,3}\b", "", caption)
+        # "#12345.."
+        caption = re.sub(r"#\d{5,}\b", "", caption)
+        # "123456.."
+        caption = re.sub(r"\b\d{6,}\b", "", caption)
+        # filenames:
+        caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
+
+        #
+        caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
+        caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
+
+        caption = re.sub(self.bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+        caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
+
+        # this-is-my-cute-cat / this_is_my_cute_cat
+        regex2 = re.compile(r"(?:\-|\_)")
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, " ", caption)
+
+        caption = ftfy.fix_text(caption)
+        caption = html.unescape(html.unescape(caption))
+
+        caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
+        caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
+        caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
+
+        caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
+        caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
+        caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
+        caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
+        caption = re.sub(r"\bpage\s+\d+\b", "", caption)
+
+        caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
+
+        caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
+
+        caption = re.sub(r"\b\s+\:\s+", r": ", caption)
+        caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
+        caption = re.sub(r"\s+", " ", caption)
+
+        caption.strip()
+
+        caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
+        caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
+        caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
+        caption = re.sub(r"^\.\S+$", "", caption)
+        return caption.strip()
+    
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int(num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), 
+            math.ceil(int(height) / self.vae.vae_scale_factor[1]),
+            math.ceil(int(width) / self.vae.vae_scale_factor[2]),
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+
+
+        return latents
+
+    def prepare_mask_masked_video(
+        self, 
+        conditional_images, 
+        conditional_images_indices, 
+        num_frames, 
+        batch_size,
+        height,
+        width,
+        num_images_per_prompt=1, 
+        use_vae_preprocessed_mask=False, 
+        device="cuda"
+    ):
+        
+        # NOTE inpaint
+        assert isinstance(conditional_images_indices, list) and len(conditional_images_indices) == len(conditional_images) and isinstance(conditional_images_indices[0], int), "conditional_images_indices should be a list of int" 
+        if isinstance(conditional_images, list) and isinstance(conditional_images[0], torch.Tensor):
+            if len(conditional_images[0].shape) == 3:
+                conditional_images = [conditional_image.unsqueeze(1) for conditional_image in conditional_images] # C H W -> C 1 H W
+            elif len(conditional_images[0].shape) == 4:
+                conditional_images = [conditional_image.transpose(0, 1) for conditional_image in conditional_images] # 1 C H W -> C 1 H W
+            conditional_images = torch.cat(conditional_images, dim=1).to(device=device) # C F H W
+        elif isinstance(conditional_images, torch.Tensor):
+            assert len(conditional_images.shape) == 4, "The shape of conditional_images should be a tensor with 4 dim"
+            conditional_images = conditional_images.transpose(0, 1) # F C H W -> C F H W
+            conditional_images = conditional_images.to(device=device)
+        else:
+            raise ValueError("conditional_images should be a list of torch.Tensor")
+
+        input_video = torch.zeros([3, num_frames, height, width], dtype=self.vae.vae.dtype, device=device)
+        input_video[:, conditional_images_indices] = conditional_images.to(input_video.dtype)
+
+        print(f"conditional_images_indices: {conditional_images_indices}")
+
+        input_video = input_video.unsqueeze(0).repeat(batch_size * num_images_per_prompt, 1, 1, 1, 1)
+        
+        # default mode
+        if not use_vae_preprocessed_mask:
+            B, C, T, H, W = input_video.shape
+            mask = torch.ones([B, 1, T, H, W], device=device)
+            mask[:, :, conditional_images_indices] = 0
+            masked_video = input_video * (mask < 0.5) 
+            masked_video = self.vae.encode(masked_video).to(device)
+
+            mask = rearrange(mask, 'b c t h w -> (b c t) 1 h w')
+            latent_size = (height // self.vae.vae_scale_factor[1], width // self.vae.vae_scale_factor[2])
+            if num_frames % 2 == 1:
+                latent_size_t = (num_frames - 1) // self.vae.vae_scale_factor[0] + 1
+            else:
+                latent_size_t = num_frames // self.vae.vae_scale_factor[0]
+            mask = F.interpolate(mask, size=latent_size, mode='bilinear')
+            mask = rearrange(mask, '(b c t) 1 h w -> b c t h w', t=T, b=B)
+            mask_first_frame = mask[:, :, 0:1].repeat(1, 1, self.vae.vae_scale_factor[0], 1, 1).contiguous()
+            mask = torch.cat([mask_first_frame, mask[:, :, 1:]], dim=2)
+            mask = mask.view(batch_size, self.vae.vae_scale_factor[0], latent_size_t, *latent_size).contiguous()
+        else:
+            mask = torch.ones_like(input_video, device=device)
+            mask[:, :, conditional_images_indices] = 0
+            masked_video = input_video * (mask < 0.5) 
+            masked_video = self.vae.encode(masked_video).to(device)
+
+            mask = self.vae.encode(mask).to(device)
+        
+        return mask, masked_video
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        # NOTE inpaint
+        conditional_images: Optional[List[torch.FloatTensor]] = None,
+        conditional_images_indices: Optional[List[int]] = None,
+        use_vae_preprocessed_mask: bool = False,
+
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        timesteps: List[int] = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: Optional[int] = 1,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        clean_caption: bool = True,
+        use_resolution_binning: bool = True,
+        max_sequence_length: int = 300,
+        **kwargs,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        """
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_inference_steps (`int`, *optional*, defaults to 100):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
+                timesteps are used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 4.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The width in pixels of the generated image.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
+                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            clean_caption (`bool`, *optional*, defaults to `True`):
+                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
+                be installed. If the dependencies are not installed, the embeddings will be created from the raw
+                prompt.
+            use_resolution_binning (`bool` defaults to `True`):
+                If set to `True`, the requested height and width are first mapped to the closest resolutions using
+                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
+                the requested resolution. Useful for generating non-square images.
+            max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
+                returned where the first element is a list with the generated images
+        """
+        # 1. Check inputs. Raise error if not correct
+        num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0]
+        height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1]
+        width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2]
+        self.check_inputs(
+            prompt,
+            num_frames, 
+            height,
+            width,
+            negative_prompt,
+            callback_steps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_attention_mask,
+        )
+
+        # 2. Default height and width to transformer
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        # import ipdb;ipdb.set_trace()
+        device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            clean_caption=clean_caption,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+
+        # 5. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            latent_channels,
+            num_frames, 
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 6.1 Prepare micro-conditions.
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+
+        # 7. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        mask, masked_video = self.prepare_mask_masked_video(
+            conditional_images, 
+            conditional_images_indices, 
+            num_frames,
+            batch_size,
+            height,
+            width,
+            num_images_per_prompt, 
+            use_vae_preprocessed_mask,
+            device=latents.device,
+        )
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents, masked_video, mask], dim=1)
+                latent_model_input = torch.cat([latent_model_input] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                current_timestep = t
+                if not torch.is_tensor(current_timestep):
+                    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                    # This would be a good case for the `match` statement (Python 3.10+)
+                    is_mps = latent_model_input.device.type == "mps"
+                    if isinstance(current_timestep, float):
+                        dtype = torch.float32 if is_mps else torch.float64
+                    else:
+                        dtype = torch.int32 if is_mps else torch.int64
+                    current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
+                elif len(current_timestep.shape) == 0:
+                    current_timestep = current_timestep[None].to(latent_model_input.device)
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                current_timestep = current_timestep.expand(latent_model_input.shape[0])
+
+                # import ipdb;ipdb.set_trace()
+                if prompt_embeds.ndim == 3:
+                    prompt_embeds = prompt_embeds.unsqueeze(1)  # b l d -> b 1 l d
+                if prompt_attention_mask.ndim == 2:
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(1)  # b l -> b 1 l
+                # prepare attention_mask.
+                # b c t h w -> b t h w
+                attention_mask = torch.ones_like(latent_model_input)[:, 0]
+                # predict noise model_output
+                noise_pred = self.transformer(
+                    latent_model_input,
+                    attention_mask=attention_mask, 
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_attention_mask=prompt_attention_mask,
+                    timestep=current_timestep,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # learned sigma
+                if self.transformer.config.out_channels // 2 == latent_channels:
+                    noise_pred = noise_pred.chunk(2, dim=1)[0]
+                else:
+                    noise_pred = noise_pred
+
+                # compute previous image: x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                # print(f'latents_{i}_{t}', torch.max(latents), torch.min(latents), torch.mean(latents), torch.std(latents))
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+        # import ipdb;ipdb.set_trace()
+        # latents = latents.squeeze(2)
+        if not output_type == "latent":
+            # b t h w c
+            image = self.decode_latents(latents)
+            image = image[:, :num_frames, :height, :width]
+        else:
+            image = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+
+    def decode_latents(self, latents):
+        # print(f'before vae decode', torch.max(latents).item(), torch.min(latents).item(), torch.mean(latents).item(), torch.std(latents).item())
+        video = self.vae.decode(latents.to(self.vae.vae.dtype))
+        # print(f'after vae decode', torch.max(video).item(), torch.min(video).item(), torch.mean(video).item(), torch.std(video).item())
+        video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous() # b t h w c
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        return video
diff --git a/opensora/sample/pipeline_inpaint_sp.py b/opensora/sample/pipeline_inpaint_sp.py
new file mode 100644
index 000000000..9d4e6ac27
--- /dev/null
+++ b/opensora/sample/pipeline_inpaint_sp.py
@@ -0,0 +1,936 @@
+# Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch.nn.functional as F
+import html
+import inspect
+import re
+import urllib.parse as ul
+from typing import Callable, List, Optional, Tuple, Union
+import math
+import torch
+from transformers import T5EncoderModel, T5Tokenizer
+from einops import rearrange
+from diffusers.models import AutoencoderKL, Transformer2DModel
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.utils import (
+    BACKENDS_MAPPING,
+    deprecate,
+    is_bs4_available,
+    is_ftfy_available,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+
+try:
+    import torch_npu
+    from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info
+except:
+    torch_npu = None
+    from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+from PIL import Image
+import numpy as np
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+if is_bs4_available():
+    from bs4 import BeautifulSoup
+
+if is_ftfy_available():
+    import ftfy
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import PixArtSigmaPipeline
+
+        >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
+        >>> pipe = PixArtSigmaPipeline.from_pretrained(
+        ...     "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
+        ... )
+        >>> # Enable memory optimizations.
+        >>> # pipe.enable_model_cpu_offload()
+
+        >>> prompt = "A small cactus with a happy face in the Sahara desert."
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+                must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class OpenSoraInpaintPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using PixArt-Sigma.
+    """
+
+    bad_punct_regex = re.compile(
+        r"["
+        + "#®•©™&@·º½¾¿¡§~"
+        + r"\)"
+        + r"\("
+        + r"\]"
+        + r"\["
+        + r"\}"
+        + r"\{"
+        + r"\|"
+        + "\\"
+        + r"\/"
+        + r"\*"
+        + r"]{1,}"
+    )  # noqa
+
+    _optional_components = ["tokenizer", "text_encoder"]
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKL,
+        transformer: Transformer2DModel,
+        scheduler: DPMSolverMultistepScheduler,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
+        )
+
+        # self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        clean_caption: bool = False,
+        max_sequence_length: int = 120,
+        **kwargs,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
+                instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
+                PixArt-Alpha, this should be "".
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                whether to use classifier free guidance or not
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                number of images that should be generated per prompt
+            device: (`torch.device`, *optional*):
+                torch device to place the resulting embeddings on
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
+                string.
+            clean_caption (`bool`, defaults to `False`):
+                If `True`, the function will preprocess and clean the provided caption before encoding.
+            max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt.
+        """
+
+        if "mask_feature" in kwargs:
+            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
+            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
+
+        if device is None:
+            device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # See Section 3.1. of the paper.
+        max_length = max_sequence_length
+
+        if prompt_embeds is None:
+            prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {max_length} tokens: {removed_text}"
+                )
+
+            prompt_attention_mask = text_inputs.attention_mask
+            prompt_attention_mask = prompt_attention_mask.to(device)
+
+            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            dtype = self.text_encoder.dtype
+        elif self.transformer is not None:
+            dtype = self.transformer.dtype
+        else:
+            dtype = None
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens = [negative_prompt] * batch_size
+            uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            negative_prompt_attention_mask = uncond_input.attention_mask
+            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        num_frames, 
+        height,
+        width,
+        negative_prompt,
+        callback_steps,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if num_frames <= 0:
+            raise ValueError(f"`num_frames` have to be positive but is {num_frames}.")
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
+    def _text_preprocessing(self, text, clean_caption=False):
+        if clean_caption and not is_bs4_available():
+            logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if clean_caption and not is_ftfy_available():
+            logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if not isinstance(text, (tuple, list)):
+            text = [text]
+
+        def process(text: str):
+            if clean_caption:
+                text = self._clean_caption(text)
+                text = self._clean_caption(text)
+            else:
+                text = text.lower().strip()
+            return text
+
+        return [process(t) for t in text]
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
+    def _clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub("<person>", "person", caption)
+        # urls:
+        caption = re.sub(
+            r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        caption = re.sub(
+            r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        # html:
+        caption = BeautifulSoup(caption, features="html.parser").text
+
+        # @<nickname>
+        caption = re.sub(r"@[\w\d]+\b", "", caption)
+
+        # 31C0—31EF CJK Strokes
+        # 31F0—31FF Katakana Phonetic Extensions
+        # 3200—32FF Enclosed CJK Letters and Months
+        # 3300—33FF CJK Compatibility
+        # 3400—4DBF CJK Unified Ideographs Extension A
+        # 4DC0—4DFF Yijing Hexagram Symbols
+        # 4E00—9FFF CJK Unified Ideographs
+        caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
+        caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
+        caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
+        caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
+        caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
+        caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
+        # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
+        #######################################################
+
+        # все виды тире / all types of dash --> "-"
+        caption = re.sub(
+            r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",  # noqa
+            "-",
+            caption,
+        )
+
+        # кавычки к одному стандарту
+        caption = re.sub(r"[`´«»“”¨]", '"', caption)
+        caption = re.sub(r"[‘’]", "'", caption)
+
+        # &quot;
+        caption = re.sub(r"&quot;?", "", caption)
+        # &amp
+        caption = re.sub(r"&amp", "", caption)
+
+        # ip adresses:
+        caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
+
+        # article ids:
+        caption = re.sub(r"\d:\d\d\s+$", "", caption)
+
+        # \n
+        caption = re.sub(r"\\n", " ", caption)
+
+        # "#123"
+        caption = re.sub(r"#\d{1,3}\b", "", caption)
+        # "#12345.."
+        caption = re.sub(r"#\d{5,}\b", "", caption)
+        # "123456.."
+        caption = re.sub(r"\b\d{6,}\b", "", caption)
+        # filenames:
+        caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
+
+        #
+        caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
+        caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
+
+        caption = re.sub(self.bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+        caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
+
+        # this-is-my-cute-cat / this_is_my_cute_cat
+        regex2 = re.compile(r"(?:\-|\_)")
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, " ", caption)
+
+        caption = ftfy.fix_text(caption)
+        caption = html.unescape(html.unescape(caption))
+
+        caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
+        caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
+        caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
+
+        caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
+        caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
+        caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
+        caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
+        caption = re.sub(r"\bpage\s+\d+\b", "", caption)
+
+        caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
+
+        caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
+
+        caption = re.sub(r"\b\s+\:\s+", r": ", caption)
+        caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
+        caption = re.sub(r"\s+", " ", caption)
+
+        caption.strip()
+
+        caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
+        caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
+        caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
+        caption = re.sub(r"^\.\S+$", "", caption)
+        return caption.strip()
+    
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int(num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), 
+            math.ceil(int(height) / self.vae.vae_scale_factor[1]),
+            math.ceil(int(width) / self.vae.vae_scale_factor[2]),
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+
+
+        return latents
+
+    def prepare_mask_masked_video(
+        self, 
+        conditional_images, 
+        conditional_images_indices, 
+        num_frames, 
+        batch_size,
+        height,
+        width,
+        num_images_per_prompt=1, 
+        use_vae_preprocessed_mask=False, 
+        weight_dtype=torch.bfloat16,
+        device="cuda"
+    ):
+        # NOTE inpaint
+        assert isinstance(conditional_images_indices, list) and len(conditional_images_indices) == len(conditional_images) and isinstance(conditional_images_indices[0], int), "conditional_images_indices should be a list of int" 
+        if isinstance(conditional_images, list) and isinstance(conditional_images[0], torch.Tensor):
+            if len(conditional_images[0].shape) == 3:
+                conditional_images = [conditional_image.unsqueeze(1) for conditional_image in conditional_images] # C H W -> C 1 H W
+            elif len(conditional_images[0].shape) == 4:
+                conditional_images = [conditional_image.transpose(0, 1) for conditional_image in conditional_images] # 1 C H W -> C 1 H W
+            conditional_images = torch.cat(conditional_images, dim=1).to(device=device) # C F H W
+        elif isinstance(conditional_images, torch.Tensor):
+            assert len(conditional_images.shape) == 4, "The shape of conditional_images should be a tensor with 4 dim"
+            conditional_images = conditional_images.transpose(0, 1) # F C H W -> C F H W
+            conditional_images = conditional_images.to(device=device)
+        else:
+            raise ValueError("conditional_images should be a list of torch.Tensor")
+
+        input_video = torch.zeros([3, num_frames, height, width], dtype=self.vae.vae.dtype, device=device)
+        input_video[:, conditional_images_indices] = conditional_images.to(input_video.dtype)
+
+        print(f"conditional_images_indices: {conditional_images_indices}")
+
+        input_video = input_video.unsqueeze(0).repeat(batch_size * num_images_per_prompt, 1, 1, 1, 1)
+        
+        # default mode
+        if not use_vae_preprocessed_mask:
+            B, C, T, H, W = input_video.shape
+            mask = torch.ones([B, 1, T, H, W], device=device)
+            mask[:, :, conditional_images_indices] = 0
+            masked_video = input_video * (mask < 0.5) 
+            masked_video = self.vae.encode(masked_video).to(device)
+
+            mask = rearrange(mask, 'b c t h w -> (b c t) 1 h w')
+            latent_size = (height // self.vae.vae_scale_factor[1], width // self.vae.vae_scale_factor[2])
+            if num_frames % 2 == 1:
+                latent_size_t = (num_frames - 1) // self.vae.vae_scale_factor[0] + 1
+            else:
+                latent_size_t = num_frames // self.vae.vae_scale_factor[0]
+            mask = F.interpolate(mask, size=latent_size, mode='bilinear')
+            mask = rearrange(mask, '(b c t) 1 h w -> b c t h w', t=T, b=B)
+            mask_first_frame = mask[:, :, 0:1].repeat(1, 1, self.vae.vae_scale_factor[0], 1, 1).contiguous()
+            mask = torch.cat([mask_first_frame, mask[:, :, 1:]], dim=2)
+            mask = mask.view(batch_size, self.vae.vae_scale_factor[0], latent_size_t, *latent_size).contiguous()
+        else:
+            mask = torch.ones_like(input_video, device=device)
+            mask[:, :, conditional_images_indices] = 0
+            masked_video = input_video * (mask < 0.5) 
+            masked_video = self.vae.encode(masked_video).to(device)
+
+            mask = self.vae.encode(mask).to(device)
+        
+        mask = mask.to(dtype=weight_dtype)
+        masked_video = masked_video.to(dtype=weight_dtype)
+
+        return mask, masked_video
+
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        # NOTE inpaint
+        conditional_images: Optional[List[torch.FloatTensor]] = None,
+        conditional_images_indices: Optional[List[int]] = None,
+        use_vae_preprocessed_mask: bool = False,
+
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        timesteps: List[int] = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: Optional[int] = 1,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        clean_caption: bool = True,
+        use_resolution_binning: bool = True,
+        max_sequence_length: int = 300,
+        **kwargs,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        """
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_inference_steps (`int`, *optional*, defaults to 100):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
+                timesteps are used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 4.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The width in pixels of the generated image.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
+                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            clean_caption (`bool`, *optional*, defaults to `True`):
+                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
+                be installed. If the dependencies are not installed, the embeddings will be created from the raw
+                prompt.
+            use_resolution_binning (`bool` defaults to `True`):
+                If set to `True`, the requested height and width are first mapped to the closest resolutions using
+                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
+                the requested resolution. Useful for generating non-square images.
+            max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
+                returned where the first element is a list with the generated images
+        """
+        # 1. Check inputs. Raise error if not correct
+        num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0]
+        height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1]
+        width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2]
+        self.check_inputs(
+            prompt,
+            num_frames, 
+            height,
+            width,
+            negative_prompt,
+            callback_steps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_attention_mask,
+        )
+
+        # 2. Default height and width to transformer
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        # import ipdb;ipdb.set_trace()
+        device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            clean_caption=clean_caption,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+
+        # 5. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        world_size = hccl_info.world_size if torch_npu is not None else nccl_info.world_size
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            latent_channels,
+            (num_frames + world_size - 1) // world_size if get_sequence_parallel_state() else num_frames,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        if get_sequence_parallel_state():
+            prompt_embeds = rearrange(prompt_embeds, 'b (n x) h -> b n x h', n=world_size,
+                                      x=prompt_embeds.shape[1] // world_size).contiguous()
+            rank = hccl_info.rank if torch_npu is not None else nccl_info.rank
+            prompt_embeds = prompt_embeds[:, rank, :, :]
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 6.1 Prepare micro-conditions.
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+
+        mask, masked_video = self.prepare_mask_masked_video(
+            conditional_images, 
+            conditional_images_indices, 
+            num_frames,
+            batch_size,
+            height,
+            width,
+            num_images_per_prompt, 
+            use_vae_preprocessed_mask,
+            weight_dtype=latents.dtype,
+            device=latents.device,
+        )
+
+        if get_sequence_parallel_state():
+            latents_frames = latents.shape[2]
+            rank = hccl_info.rank if torch_npu is not None else nccl_info.rank
+            mask = mask[:, :, latents_frames * rank: latents_frames * (rank + 1)]
+            masked_video = masked_video[:, :, latents_frames * rank: latents_frames * (rank + 1)]
+
+        # 7. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents, masked_video, mask], dim=1)
+                latent_model_input = torch.cat([latent_model_input] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                current_timestep = t
+                if not torch.is_tensor(current_timestep):
+                    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                    # This would be a good case for the `match` statement (Python 3.10+)
+                    is_mps = latent_model_input.device.type == "mps"
+                    if isinstance(current_timestep, float):
+                        dtype = torch.float32 if is_mps else torch.float64
+                    else:
+                        dtype = torch.int32 if is_mps else torch.int64
+                    current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
+                elif len(current_timestep.shape) == 0:
+                    current_timestep = current_timestep[None].to(latent_model_input.device)
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                current_timestep = current_timestep.expand(latent_model_input.shape[0])
+
+                # import ipdb;ipdb.set_trace()
+                if prompt_embeds.ndim == 3:
+                    prompt_embeds = prompt_embeds.unsqueeze(1)  # b l d -> b 1 l d
+                if prompt_attention_mask.ndim == 2:
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(1)  # b l -> b 1 l
+                # prepare attention_mask.
+                # b c t h w -> b t h w
+                attention_mask = torch.ones_like(latent_model_input)[:, 0]
+                if get_sequence_parallel_state():
+                    attention_mask = attention_mask.repeat(1, world_size, 1, 1)
+                # predict noise model_output
+                noise_pred = self.transformer(
+                    latent_model_input,
+                    attention_mask=attention_mask, 
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_attention_mask=prompt_attention_mask,
+                    timestep=current_timestep,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # learned sigma
+                if self.transformer.config.out_channels // 2 == latent_channels:
+                    noise_pred = noise_pred.chunk(2, dim=1)[0]
+                else:
+                    noise_pred = noise_pred
+
+                # compute previous image: x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                # print(f'latents_{i}_{t}', torch.max(latents), torch.min(latents), torch.mean(latents), torch.std(latents))
+                
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+        # import ipdb;ipdb.set_trace()
+        # latents = latents.squeeze(2)
+        world_size = hccl_info.world_size if torch_npu is not None else nccl_info.world_size
+        if get_sequence_parallel_state():
+            latents_shape = list(latents.shape)
+            full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
+            all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
+            torch.distributed.all_gather_into_tensor(all_latents, latents)
+            latents_list = list(all_latents.chunk(world_size, dim=0))
+            latents = torch.cat(latents_list, dim=2)
+
+        if not output_type == "latent":
+            # b t h w c
+            image = self.decode_latents(latents)
+            image = image[:, :num_frames, :height, :width]
+        else:
+            image = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+    
+    
+    def decode_latents(self, latents):
+        # print(f'before vae decode', torch.max(latents).item(), torch.min(latents).item(), torch.mean(latents).item(), torch.std(latents).item())
+        video = self.vae.decode(latents.to(self.vae.vae.dtype))
+        # print(f'after vae decode', torch.max(video).item(), torch.min(video).item(), torch.mean(video).item(), torch.std(video).item())
+        video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous() # b t h w c
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        return video
diff --git a/opensora/sample/pipeline_opensora.py b/opensora/sample/pipeline_opensora.py
new file mode 100644
index 000000000..86ba919d8
--- /dev/null
+++ b/opensora/sample/pipeline_opensora.py
@@ -0,0 +1,820 @@
+# Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import html
+import inspect
+import re
+import urllib.parse as ul
+from typing import Callable, List, Optional, Tuple, Union
+import math
+import torch
+from transformers import T5EncoderModel, T5Tokenizer
+
+from diffusers.models import AutoencoderKL, Transformer2DModel
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.utils import (
+    BACKENDS_MAPPING,
+    deprecate,
+    is_bs4_available,
+    is_ftfy_available,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+
+try:
+    from opensora.npu_config import npu_config
+except:
+    npu_config = None
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+if is_bs4_available():
+    from bs4 import BeautifulSoup
+
+if is_ftfy_available():
+    import ftfy
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import PixArtSigmaPipeline
+
+        >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
+        >>> pipe = PixArtSigmaPipeline.from_pretrained(
+        ...     "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
+        ... )
+        >>> # Enable memory optimizations.
+        >>> # pipe.enable_model_cpu_offload()
+
+        >>> prompt = "A small cactus with a happy face in the Sahara desert."
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+                must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class OpenSoraPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using PixArt-Sigma.
+    """
+
+    bad_punct_regex = re.compile(
+        r"["
+        + "#®•©™&@·º½¾¿¡§~"
+        + r"\)"
+        + r"\("
+        + r"\]"
+        + r"\["
+        + r"\}"
+        + r"\{"
+        + r"\|"
+        + "\\"
+        + r"\/"
+        + r"\*"
+        + r"]{1,}"
+    )  # noqa
+
+    _optional_components = ["tokenizer", "text_encoder"]
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKL,
+        transformer: Transformer2DModel,
+        scheduler: DPMSolverMultistepScheduler,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
+        )
+
+        # self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        clean_caption: bool = False,
+        max_sequence_length: int = 120,
+        **kwargs,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
+                instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
+                PixArt-Alpha, this should be "".
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                whether to use classifier free guidance or not
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                number of images that should be generated per prompt
+            device: (`torch.device`, *optional*):
+                torch device to place the resulting embeddings on
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
+                string.
+            clean_caption (`bool`, defaults to `False`):
+                If `True`, the function will preprocess and clean the provided caption before encoding.
+            max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt.
+        """
+
+        if "mask_feature" in kwargs:
+            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
+            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
+
+        if device is None:
+            device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # See Section 3.1. of the paper.
+        max_length = max_sequence_length
+
+        if prompt_embeds is None:
+            prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {max_length} tokens: {removed_text}"
+                )
+
+            prompt_attention_mask = text_inputs.attention_mask
+            prompt_attention_mask = prompt_attention_mask.to(device)
+
+            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            dtype = self.text_encoder.dtype
+        elif self.transformer is not None:
+            dtype = self.transformer.dtype
+        else:
+            dtype = None
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens = [negative_prompt] * batch_size
+            uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            negative_prompt_attention_mask = uncond_input.attention_mask
+            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        num_frames, 
+        height,
+        width,
+        negative_prompt,
+        callback_steps,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if num_frames <= 0:
+            raise ValueError(f"`num_frames` have to be positive but is {num_frames}.")
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
+    def _text_preprocessing(self, text, clean_caption=False):
+        if clean_caption and not is_bs4_available():
+            logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if clean_caption and not is_ftfy_available():
+            logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if not isinstance(text, (tuple, list)):
+            text = [text]
+
+        def process(text: str):
+            if clean_caption:
+                text = self._clean_caption(text)
+                text = self._clean_caption(text)
+            else:
+                text = text.lower().strip()
+            return text
+
+        return [process(t) for t in text]
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
+    def _clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub("<person>", "person", caption)
+        # urls:
+        caption = re.sub(
+            r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        caption = re.sub(
+            r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        # html:
+        caption = BeautifulSoup(caption, features="html.parser").text
+
+        # @<nickname>
+        caption = re.sub(r"@[\w\d]+\b", "", caption)
+
+        # 31C0—31EF CJK Strokes
+        # 31F0—31FF Katakana Phonetic Extensions
+        # 3200—32FF Enclosed CJK Letters and Months
+        # 3300—33FF CJK Compatibility
+        # 3400—4DBF CJK Unified Ideographs Extension A
+        # 4DC0—4DFF Yijing Hexagram Symbols
+        # 4E00—9FFF CJK Unified Ideographs
+        caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
+        caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
+        caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
+        caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
+        caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
+        caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
+        # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
+        #######################################################
+
+        # все виды тире / all types of dash --> "-"
+        caption = re.sub(
+            r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",  # noqa
+            "-",
+            caption,
+        )
+
+        # кавычки к одному стандарту
+        caption = re.sub(r"[`´«»“”¨]", '"', caption)
+        caption = re.sub(r"[‘’]", "'", caption)
+
+        # &quot;
+        caption = re.sub(r"&quot;?", "", caption)
+        # &amp
+        caption = re.sub(r"&amp", "", caption)
+
+        # ip adresses:
+        caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
+
+        # article ids:
+        caption = re.sub(r"\d:\d\d\s+$", "", caption)
+
+        # \n
+        caption = re.sub(r"\\n", " ", caption)
+
+        # "#123"
+        caption = re.sub(r"#\d{1,3}\b", "", caption)
+        # "#12345.."
+        caption = re.sub(r"#\d{5,}\b", "", caption)
+        # "123456.."
+        caption = re.sub(r"\b\d{6,}\b", "", caption)
+        # filenames:
+        caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
+
+        #
+        caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
+        caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
+
+        caption = re.sub(self.bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+        caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
+
+        # this-is-my-cute-cat / this_is_my_cute_cat
+        regex2 = re.compile(r"(?:\-|\_)")
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, " ", caption)
+
+        caption = ftfy.fix_text(caption)
+        caption = html.unescape(html.unescape(caption))
+
+        caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
+        caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
+        caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
+
+        caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
+        caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
+        caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
+        caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
+        caption = re.sub(r"\bpage\s+\d+\b", "", caption)
+
+        caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
+
+        caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
+
+        caption = re.sub(r"\b\s+\:\s+", r": ", caption)
+        caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
+        caption = re.sub(r"\s+", " ", caption)
+
+        caption.strip()
+
+        caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
+        caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
+        caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
+        caption = re.sub(r"^\.\S+$", "", caption)
+        return caption.strip()
+    
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int(num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), 
+            math.ceil(int(height) / self.vae.vae_scale_factor[1]),
+            math.ceil(int(width) / self.vae.vae_scale_factor[2]),
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+
+
+        return latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        timesteps: List[int] = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: Optional[int] = 1,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        clean_caption: bool = True,
+        use_resolution_binning: bool = True,
+        max_sequence_length: int = 300,
+        **kwargs,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        """
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_inference_steps (`int`, *optional*, defaults to 100):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
+                timesteps are used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 4.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The width in pixels of the generated image.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
+                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            clean_caption (`bool`, *optional*, defaults to `True`):
+                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
+                be installed. If the dependencies are not installed, the embeddings will be created from the raw
+                prompt.
+            use_resolution_binning (`bool` defaults to `True`):
+                If set to `True`, the requested height and width are first mapped to the closest resolutions using
+                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
+                the requested resolution. Useful for generating non-square images.
+            max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
+                returned where the first element is a list with the generated images
+        """
+        # 1. Check inputs. Raise error if not correct
+        num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0]
+        height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1]
+        width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2]
+        self.check_inputs(
+            prompt,
+            num_frames, 
+            height,
+            width,
+            negative_prompt,
+            callback_steps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_attention_mask,
+        )
+
+        # 2. Default height and width to transformer
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        # import ipdb;ipdb.set_trace()
+        device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            clean_caption=clean_caption,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+
+        # 5. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            latent_channels,
+            num_frames, 
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 6.1 Prepare micro-conditions.
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+
+        # 7. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                current_timestep = t
+                if not torch.is_tensor(current_timestep):
+                    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                    # This would be a good case for the `match` statement (Python 3.10+)
+                    is_mps = latent_model_input.device.type == "mps"
+                    if isinstance(current_timestep, float):
+                        dtype = torch.float32 if is_mps else torch.float64
+                    else:
+                        dtype = torch.int32 if is_mps else torch.int64
+                    current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
+                elif len(current_timestep.shape) == 0:
+                    current_timestep = current_timestep[None].to(latent_model_input.device)
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                current_timestep = current_timestep.expand(latent_model_input.shape[0])
+
+                # import ipdb;ipdb.set_trace()
+                if prompt_embeds.ndim == 3:
+                    prompt_embeds = prompt_embeds.unsqueeze(1)  # b l d -> b 1 l d
+                if prompt_attention_mask.ndim == 2:
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(1)  # b l -> b 1 l
+                # prepare attention_mask.
+                # b c t h w -> b t h w
+                attention_mask = torch.ones_like(latent_model_input)[:, 0]
+                # predict noise model_output
+                noise_pred = self.transformer(
+                    latent_model_input,
+                    attention_mask=attention_mask, 
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_attention_mask=prompt_attention_mask,
+                    timestep=current_timestep,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # learned sigma
+                if self.transformer.config.out_channels // 2 == latent_channels:
+                    noise_pred = noise_pred.chunk(2, dim=1)[0]
+                else:
+                    noise_pred = noise_pred
+
+                # compute previous image: x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                # print(f'latents_{i}_{t}', torch.max(latents), torch.min(latents), torch.mean(latents), torch.std(latents))
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+        # import ipdb;ipdb.set_trace()
+        # latents = latents.squeeze(2)
+        if not output_type == "latent":
+            # b t h w c
+            image = self.decode_latents(latents)
+            image = image[:, :num_frames, :height, :width]
+        else:
+            image = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+
+    def decode_latents(self, latents):
+        # print(f'before vae decode', torch.max(latents).item(), torch.min(latents).item(), torch.mean(latents).item(), torch.std(latents).item())
+        video = self.vae.decode(latents.to(self.vae.vae.dtype))
+        # print(f'after vae decode', torch.max(video).item(), torch.min(video).item(), torch.mean(video).item(), torch.std(video).item())
+        video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous() # b t h w c
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        return video
diff --git a/opensora/sample/pipeline_opensora_sp.py b/opensora/sample/pipeline_opensora_sp.py
new file mode 100644
index 000000000..03a84599e
--- /dev/null
+++ b/opensora/sample/pipeline_opensora_sp.py
@@ -0,0 +1,841 @@
+# Copyright 2024 PixArt-Sigma Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import html
+import inspect
+import re
+import urllib.parse as ul
+from typing import Callable, List, Optional, Tuple, Union
+import math
+import torch
+from transformers import T5EncoderModel, T5Tokenizer
+from einops import rearrange
+from diffusers.models import AutoencoderKL, Transformer2DModel
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.utils import (
+    BACKENDS_MAPPING,
+    deprecate,
+    is_bs4_available,
+    is_ftfy_available,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+
+try:
+    import torch_npu
+    from opensora.acceleration.parallel_states import get_sequence_parallel_state, hccl_info
+except:
+    torch_npu = None
+    from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+if is_bs4_available():
+    from bs4 import BeautifulSoup
+
+if is_ftfy_available():
+    import ftfy
+
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import PixArtSigmaPipeline
+
+        >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
+        >>> pipe = PixArtSigmaPipeline.from_pretrained(
+        ...     "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
+        ... )
+        >>> # Enable memory optimizations.
+        >>> # pipe.enable_model_cpu_offload()
+
+        >>> prompt = "A small cactus with a happy face in the Sahara desert."
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+                must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class OpenSoraPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using PixArt-Sigma.
+    """
+
+    bad_punct_regex = re.compile(
+        r"["
+        + "#®•©™&@·º½¾¿¡§~"
+        + r"\)"
+        + r"\("
+        + r"\]"
+        + r"\["
+        + r"\}"
+        + r"\{"
+        + r"\|"
+        + "\\"
+        + r"\/"
+        + r"\*"
+        + r"]{1,}"
+    )  # noqa
+
+    _optional_components = ["tokenizer", "text_encoder"]
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKL,
+        transformer: Transformer2DModel,
+        scheduler: DPMSolverMultistepScheduler,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
+        )
+
+        # self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        clean_caption: bool = False,
+        max_sequence_length: int = 120,
+        **kwargs,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
+                instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
+                PixArt-Alpha, this should be "".
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                whether to use classifier free guidance or not
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                number of images that should be generated per prompt
+            device: (`torch.device`, *optional*):
+                torch device to place the resulting embeddings on
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
+                string.
+            clean_caption (`bool`, defaults to `False`):
+                If `True`, the function will preprocess and clean the provided caption before encoding.
+            max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt.
+        """
+
+        if "mask_feature" in kwargs:
+            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
+            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
+
+        if device is None:
+            device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # See Section 3.1. of the paper.
+        max_length = max_sequence_length
+
+        if prompt_embeds is None:
+            prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {max_length} tokens: {removed_text}"
+                )
+
+            prompt_attention_mask = text_inputs.attention_mask
+            prompt_attention_mask = prompt_attention_mask.to(device)
+
+            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            dtype = self.text_encoder.dtype
+        elif self.transformer is not None:
+            dtype = self.transformer.dtype
+        else:
+            dtype = None
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
+        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens = [negative_prompt] * batch_size
+            uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            negative_prompt_attention_mask = uncond_input.attention_mask
+            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        num_frames, 
+        height,
+        width,
+        negative_prompt,
+        callback_steps,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if num_frames <= 0:
+            raise ValueError(f"`num_frames` have to be positive but is {num_frames}.")
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+
+        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
+            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
+    def _text_preprocessing(self, text, clean_caption=False):
+        if clean_caption and not is_bs4_available():
+            logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if clean_caption and not is_ftfy_available():
+            logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
+            logger.warning("Setting `clean_caption` to False...")
+            clean_caption = False
+
+        if not isinstance(text, (tuple, list)):
+            text = [text]
+
+        def process(text: str):
+            if clean_caption:
+                text = self._clean_caption(text)
+                text = self._clean_caption(text)
+            else:
+                text = text.lower().strip()
+            return text
+
+        return [process(t) for t in text]
+
+    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
+    def _clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub("<person>", "person", caption)
+        # urls:
+        caption = re.sub(
+            r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        caption = re.sub(
+            r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
+            "",
+            caption,
+        )  # regex for urls
+        # html:
+        caption = BeautifulSoup(caption, features="html.parser").text
+
+        # @<nickname>
+        caption = re.sub(r"@[\w\d]+\b", "", caption)
+
+        # 31C0—31EF CJK Strokes
+        # 31F0—31FF Katakana Phonetic Extensions
+        # 3200—32FF Enclosed CJK Letters and Months
+        # 3300—33FF CJK Compatibility
+        # 3400—4DBF CJK Unified Ideographs Extension A
+        # 4DC0—4DFF Yijing Hexagram Symbols
+        # 4E00—9FFF CJK Unified Ideographs
+        caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
+        caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
+        caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
+        caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
+        caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
+        caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
+        # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
+        #######################################################
+
+        # все виды тире / all types of dash --> "-"
+        caption = re.sub(
+            r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",  # noqa
+            "-",
+            caption,
+        )
+
+        # кавычки к одному стандарту
+        caption = re.sub(r"[`´«»“”¨]", '"', caption)
+        caption = re.sub(r"[‘’]", "'", caption)
+
+        # &quot;
+        caption = re.sub(r"&quot;?", "", caption)
+        # &amp
+        caption = re.sub(r"&amp", "", caption)
+
+        # ip adresses:
+        caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
+
+        # article ids:
+        caption = re.sub(r"\d:\d\d\s+$", "", caption)
+
+        # \n
+        caption = re.sub(r"\\n", " ", caption)
+
+        # "#123"
+        caption = re.sub(r"#\d{1,3}\b", "", caption)
+        # "#12345.."
+        caption = re.sub(r"#\d{5,}\b", "", caption)
+        # "123456.."
+        caption = re.sub(r"\b\d{6,}\b", "", caption)
+        # filenames:
+        caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
+
+        #
+        caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
+        caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
+
+        caption = re.sub(self.bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+        caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
+
+        # this-is-my-cute-cat / this_is_my_cute_cat
+        regex2 = re.compile(r"(?:\-|\_)")
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, " ", caption)
+
+        caption = ftfy.fix_text(caption)
+        caption = html.unescape(html.unescape(caption))
+
+        caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
+        caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
+        caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
+
+        caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
+        caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
+        caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
+        caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
+        caption = re.sub(r"\bpage\s+\d+\b", "", caption)
+
+        caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
+
+        caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
+
+        caption = re.sub(r"\b\s+\:\s+", r": ", caption)
+        caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
+        caption = re.sub(r"\s+", " ", caption)
+
+        caption.strip()
+
+        caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
+        caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
+        caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
+        caption = re.sub(r"^\.\S+$", "", caption)
+        return caption.strip()
+    
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int(num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), 
+            math.ceil(int(height) / self.vae.vae_scale_factor[1]),
+            math.ceil(int(width) / self.vae.vae_scale_factor[2]),
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+
+
+        return latents
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        timesteps: List[int] = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: Optional[int] = 1,
+        num_frames: Optional[int] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        clean_caption: bool = True,
+        use_resolution_binning: bool = True,
+        max_sequence_length: int = 300,
+        **kwargs,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        """
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_inference_steps (`int`, *optional*, defaults to 100):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
+                timesteps are used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 4.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The width in pixels of the generated image.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
+                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            clean_caption (`bool`, *optional*, defaults to `True`):
+                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
+                be installed. If the dependencies are not installed, the embeddings will be created from the raw
+                prompt.
+            use_resolution_binning (`bool` defaults to `True`):
+                If set to `True`, the requested height and width are first mapped to the closest resolutions using
+                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
+                the requested resolution. Useful for generating non-square images.
+            max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
+                returned where the first element is a list with the generated images
+        """
+        # 1. Check inputs. Raise error if not correct
+        num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0]
+        height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1]
+        width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2]
+        self.check_inputs(
+            prompt,
+            num_frames, 
+            height,
+            width,
+            negative_prompt,
+            callback_steps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_attention_mask,
+        )
+
+        # 2. Default height and width to transformer
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+        # import ipdb;ipdb.set_trace()
+        device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            clean_caption=clean_caption,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+
+        # 5. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        world_size = hccl_info.world_size if torch_npu is not None else nccl_info.world_size
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            latent_channels,
+            (num_frames + world_size - 1) // world_size if get_sequence_parallel_state() else num_frames,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        if get_sequence_parallel_state():
+            prompt_embeds = rearrange(prompt_embeds, 'b (n x) h -> b n x h', n=world_size,
+                                      x=prompt_embeds.shape[1] // world_size).contiguous()
+            rank = hccl_info.rank if torch_npu is not None else nccl_info.rank
+            prompt_embeds = prompt_embeds[:, rank, :, :]
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 6.1 Prepare micro-conditions.
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+
+        # 7. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                current_timestep = t
+                if not torch.is_tensor(current_timestep):
+                    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                    # This would be a good case for the `match` statement (Python 3.10+)
+                    is_mps = latent_model_input.device.type == "mps"
+                    if isinstance(current_timestep, float):
+                        dtype = torch.float32 if is_mps else torch.float64
+                    else:
+                        dtype = torch.int32 if is_mps else torch.int64
+                    current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
+                elif len(current_timestep.shape) == 0:
+                    current_timestep = current_timestep[None].to(latent_model_input.device)
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                current_timestep = current_timestep.expand(latent_model_input.shape[0])
+
+                # import ipdb;ipdb.set_trace()
+                if prompt_embeds.ndim == 3:
+                    prompt_embeds = prompt_embeds.unsqueeze(1)  # b l d -> b 1 l d
+                if prompt_attention_mask.ndim == 2:
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(1)  # b l -> b 1 l
+                # prepare attention_mask.
+                # b c t h w -> b t h w
+                attention_mask = torch.ones_like(latent_model_input)[:, 0]
+                if get_sequence_parallel_state():
+                    attention_mask = attention_mask.repeat(1, world_size, 1, 1)
+                # predict noise model_output
+                noise_pred = self.transformer(
+                    latent_model_input,
+                    attention_mask=attention_mask, 
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_attention_mask=prompt_attention_mask,
+                    timestep=current_timestep,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # learned sigma
+                if self.transformer.config.out_channels // 2 == latent_channels:
+                    noise_pred = noise_pred.chunk(2, dim=1)[0]
+                else:
+                    noise_pred = noise_pred
+
+                # compute previous image: x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                # print(f'latents_{i}_{t}', torch.max(latents), torch.min(latents), torch.mean(latents), torch.std(latents))
+                
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+        # import ipdb;ipdb.set_trace()
+        # latents = latents.squeeze(2)
+        world_size = hccl_info.world_size if torch_npu is not None else nccl_info.world_size
+        if get_sequence_parallel_state():
+            latents_shape = list(latents.shape)
+            full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
+            all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
+            torch.distributed.all_gather_into_tensor(all_latents, latents)
+            latents_list = list(all_latents.chunk(world_size, dim=0))
+            latents = torch.cat(latents_list, dim=2)
+
+        if not output_type == "latent":
+            # b t h w c
+            image = self.decode_latents(latents)
+            image = image[:, :num_frames, :height, :width]
+        else:
+            image = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+    
+    
+    def decode_latents(self, latents):
+        # print(f'before vae decode', torch.max(latents).item(), torch.min(latents).item(), torch.mean(latents).item(), torch.std(latents).item())
+        video = self.vae.decode(latents.to(self.vae.vae.dtype))
+        # print(f'after vae decode', torch.max(video).item(), torch.min(video).item(), torch.mean(video).item(), torch.std(video).item())
+        video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous() # b t h w c
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        return video
diff --git a/opensora/sample/sample_inpaint.py b/opensora/sample/sample_inpaint.py
new file mode 100644
index 000000000..43b0fe3ff
--- /dev/null
+++ b/opensora/sample/sample_inpaint.py
@@ -0,0 +1,264 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, MT5EncoderModel, UMT5EncoderModel, AutoTokenizer
+
+
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.opensora.modeling_inpaint import OpenSoraInpaint
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_inpaint import OpenSoraInpaintPipeline
+
+import imageio
+
+# for validation
+import glob
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms import Lambda
+from opensora.dataset.transform import ToTensorVideo, CenterCropResizeVideo, TemporalRandomCrop, LongSideResizeVideo, SpatialStrideCropVideo
+import numpy as np
+from einops import rearrange
+
+def main(args):
+    # torch.manual_seed(args.seed)
+    # torch.backends.cuda.matmul.allow_tf32 = False
+    weight_dtype = torch.bfloat16
+    device = torch.device(args.device)
+
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    norm_fun = Lambda(lambda x: 2. * x - 1.)
+
+    resize = [CenterCropResizeVideo((args.height, args.width)), ]
+    transform = transforms.Compose([
+        ToTensorVideo(),
+        *resize, 
+        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+        norm_fun
+    ])
+
+    def preprocess_images(images):
+        print(images)
+        if len(images) == 1:
+            conditional_images_indices = [0]
+        elif len(images) == 2:
+            conditional_images_indices = [0, -1]
+        else:
+            print("Only support 1 or 2 condition images!")
+            raise NotImplementedError
+        
+        try:
+            conditional_images = [Image.open(image).convert("RGB") for image in images]
+            conditional_images = [torch.from_numpy(np.copy(np.array(image))) for image in conditional_images]
+            conditional_images = [rearrange(image, 'h w c -> c h w').unsqueeze(0) for image in conditional_images]
+            conditional_images = [transform(image).to(device=device, dtype=weight_dtype) for image in conditional_images]
+        except Exception as e:
+            print('Error when loading images')
+            print(f'condition images are {images}')
+            raise e
+        return dict(conditional_images=conditional_images, conditional_images_indices=conditional_images_indices)
+
+
+    transformer_model = OpenSoraInpaint.from_pretrained(args.model_path, cache_dir=args.cache_dir, 
+                                                        low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
+
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+    # set eval mode
+    transformer_model.eval()
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler(clip_sample=False)
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler(clip_sample=False)
+    elif args.sample_method == 'DPMSolverMultistep':
+        '''
+        DPM++ 2M	        DPMSolverMultistepScheduler	
+        DPM++ 2M Karras	    DPMSolverMultistepScheduler	init with use_karras_sigmas=True
+        DPM++ 2M SDE	    DPMSolverMultistepScheduler	init with algorithm_type="sde-dpmsolver++"
+        DPM++ 2M SDE Karras	DPMSolverMultistepScheduler	init with use_karras_sigmas=True and algorithm_type="sde-dpmsolver++"
+        
+        DPM++ SDE	        DPMSolverSinglestepScheduler	
+        DPM++ SDE Karras	DPMSolverSinglestepScheduler	init with use_karras_sigmas=True
+        DPM2	            KDPM2DiscreteScheduler	
+        DPM2 Karras	        KDPM2DiscreteScheduler	init with use_karras_sigmas=True
+        DPM2 a	            KDPM2AncestralDiscreteScheduler	
+        DPM2 a Karras	    KDPM2AncestralDiscreteScheduler	init with use_karras_sigmas=True
+        '''
+        # scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+    elif args.sample_method == 'EulerDiscreteSVD':
+        scheduler = EulerDiscreteScheduler.from_pretrained("stabilityai/stable-video-diffusion-img2vid", 
+                                                        subfolder="scheduler", cache_dir=args.cache_dir)
+    pipeline = OpenSoraInpaintPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model)
+    pipeline.to(dtype=weight_dtype, device=device)
+    if args.compile:
+        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+        options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
+                    "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
+                    "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
+        #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
+        from onediffx import compile_pipe
+        pipeline = compile_pipe(
+                pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
+            )
+
+        # 4%
+        # pipeline.transformer = torch.compile(pipeline.transformer)
+    
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path)
+        
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        text_prompt = [i.strip() for i in text_prompt]
+
+    if not isinstance(args.conditional_images_path, list):
+        args.conditional_images_path = [args.conditional_images_path]
+    if len(args.conditional_images_path) == 1 and args.conditional_images_path[0].endswith('txt'):
+        temp = open(args.conditional_images_path[0], 'r').readlines()
+        conditional_images = [i.strip().split(',') for i in temp]
+        
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), 
+    {}. 
+    """
+    
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+
+    video_grids = []
+    for idx, (prompt, images) in enumerate(zip(text_prompt, conditional_images)):
+
+        pre_results = preprocess_images(images)
+        cond_imgs = pre_results['conditional_images']
+        cond_imgs_indices = pre_results['conditional_images_indices']
+
+        videos = pipeline(
+            conditional_images=cond_imgs,
+            conditional_images_indices=cond_imgs_indices,
+            prompt=positive_prompt.format(prompt),
+            negative_prompt=negative_prompt, 
+            num_frames=args.num_frames,
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_sampling_steps,
+            guidance_scale=args.guidance_scale,
+            num_images_per_prompt=1,
+            mask_feature=True,
+            device=args.device, 
+            max_sequence_length=args.max_sequence_length, 
+        ).images
+        try:
+            if args.num_frames == 1:
+                ext = 'jpg'
+                videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                save_image(videos / 255.0, os.path.join(args.save_img_path, f'{idx}.{ext}'), nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+            else:
+                ext = 'mp4'
+                imageio.mimwrite(
+                    os.path.join(args.save_img_path, f'{idx}.{ext}'), videos[0], fps=args.fps, quality=6)  # highest quality is 10, lowest is 0
+            save_path = os.path.join(args.save_img_path, f'{idx}.{ext}')
+            print(f'image or video is saved at {save_path}')
+        except:
+            print('Error when saving {}'.format(prompt))
+        video_grids.append(videos)
+    video_grids = torch.cat(video_grids, dim=0)
+
+    # torchvision.io.write_video(args.save_img_path + '_%04d' % args.run_time + '-.mp4', video_grids, fps=6)
+    if args.num_frames == 1:
+        save_image(video_grids / 255.0, os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
+                   nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+    else:
+        video_grids = save_video_grid(video_grids)
+        imageio.mimwrite(os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'), video_grids, fps=args.fps, quality=6)
+
+    print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=2.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--max_sequence_length", type=int, default=300)
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--run_time", type=int, default=0)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--compile', action='store_true')
+    parser.add_argument('--save_memory', action='store_true')
+
+    parser.add_argument('--conditional_images_path', type=str, help='the path of txt file containing the paths of condition images')
+    args = parser.parse_args()
+
+    main(args)
diff --git a/opensora/sample/sample_inpaint_ddp.py b/opensora/sample/sample_inpaint_ddp.py
new file mode 100644
index 000000000..8260d147a
--- /dev/null
+++ b/opensora/sample/sample_inpaint_ddp.py
@@ -0,0 +1,357 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+import torch.distributed as dist
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderModel
+
+import os, sys
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+from opensora.models.diffusion.opensora.modeling_inpaint import OpenSoraInpaint
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_inpaint import OpenSoraInpaintPipeline
+
+import imageio
+
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+    pass
+import time
+
+# for validation
+import glob
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms import Lambda
+from opensora.dataset.transform import ToTensorVideo, CenterCropResizeVideo, TemporalRandomCrop, LongSideResizeVideo, SpatialStrideCropVideo
+import numpy as np
+from einops import rearrange
+
+
+def load_t2v_checkpoint(model_path):
+
+    transformer_model = OpenSoraInpaint.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                    low_cpu_mem_usage=False, device_map=None,
+                                                    torch_dtype=weight_dtype)
+    # print(transformer_model.config)
+
+    # set eval mode
+    transformer_model.eval()
+    pipeline = OpenSoraInpaintPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model).to(device)
+
+    if args.compile:
+        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+        # options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
+        #             "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
+        #             "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
+        # #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # from onediffx import compile_pipe
+        # pipeline = compile_pipe(
+        #         pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
+        #     )
+
+        # 4%
+        pipeline.transformer = torch.compile(pipeline.transformer)
+    return pipeline
+
+
+def get_latest_path():
+    # Get the most recent checkpoint
+    dirs = os.listdir(args.model_path)
+    dirs = [d for d in dirs if d.startswith("checkpoint")]
+    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+    path = dirs[-1] if len(dirs) > 0 else None
+
+    return path
+
+
+def run_model_and_save_images(pipeline, model_path):
+    norm_fun = Lambda(lambda x: 2. * x - 1.)
+
+    resize = [CenterCropResizeVideo((args.height, args.width)), ]
+    transform = transforms.Compose([
+        ToTensorVideo(),
+        *resize, 
+        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+        norm_fun
+    ])
+
+    def preprocess_images(images):
+        print(images)
+        if len(images) == 1:
+            conditional_images_indices = [0]
+        elif len(images) == 2:
+            conditional_images_indices = [0, -1]
+        else:
+            print("Only support 1 or 2 condition images!")
+            raise NotImplementedError
+        
+        try:
+            conditional_images = [Image.open(image).convert("RGB") for image in images]
+            conditional_images = [torch.from_numpy(np.copy(np.array(image))) for image in conditional_images]
+            conditional_images = [rearrange(image, 'h w c -> c h w').unsqueeze(0) for image in conditional_images]
+            conditional_images = [transform(image).to(device=device, dtype=weight_dtype) for image in conditional_images]
+        except Exception as e:
+            print('Error when loading images')
+            print(f'condition images are {images}')
+            raise e
+        return dict(conditional_images=conditional_images, conditional_images_indices=conditional_images_indices)
+    
+
+    checkpoint_name = f"{os.path.basename(model_path)}"
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), (unwatermarked), 
+    {}. 
+    emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
+    sharp focus, high budget, cinemascope, moody, epic, gorgeous
+    """
+    
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path)
+        
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        text_prompt = [i.strip() for i in text_prompt]
+
+    if not isinstance(args.conditional_images_path, list):
+        args.conditional_images_path = [args.conditional_images_path]
+    if len(args.conditional_images_path) == 1 and args.conditional_images_path[0].endswith('txt'):
+        temp = open(args.conditional_images_path[0], 'r').readlines()
+        conditional_images = [i.strip().split(',') for i in temp]
+    
+    assert len(text_prompt) % world_size == 0, "The sample num must be a multiple of the world size; otherwise, it may cause an all_gather error."
+
+    video_grids = []
+    for idx, (prompt, images) in enumerate(zip(text_prompt, conditional_images)):
+        if idx % world_size != local_rank:
+            continue
+
+        pre_results = preprocess_images(images)
+        cond_imgs = pre_results['conditional_images']
+        cond_imgs_indices = pre_results['conditional_images_indices']
+
+        videos = pipeline(
+            conditional_images=cond_imgs,
+            conditional_images_indices=cond_imgs_indices,
+            prompt=positive_prompt.format(prompt),
+            negative_prompt=negative_prompt, 
+            num_frames=args.num_frames,
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_sampling_steps,
+            guidance_scale=args.guidance_scale,
+            num_images_per_prompt=1,
+            mask_feature=True,
+            device=args.device, 
+            max_sequence_length=args.max_sequence_length, 
+        ).images
+
+        try:
+            if args.num_frames == 1:
+                ext = 'jpg'
+                videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                save_image(videos / 255.0, os.path.join(args.save_img_path, f'{idx}.{ext}'), nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+            else:
+                ext = 'mp4'
+                imageio.mimwrite(
+                    os.path.join(args.save_img_path, f'{idx}.{ext}'), videos[0], fps=args.fps, quality=6)  # highest quality is 10, lowest is 0
+                save_path = os.path.join(args.save_img_path, f'{idx}.{ext}')
+                print(f'image or video is saved at {save_path}')
+        except:
+            print('Error when saving {}'.format(prompt))
+        video_grids.append(videos)
+
+    dist.barrier()
+    video_grids = torch.cat(video_grids, dim=0).cuda()
+    shape = list(video_grids.shape)
+    shape[0] *= world_size
+    gathered_tensor = torch.zeros(shape, dtype=video_grids.dtype, device=device)
+    dist.all_gather_into_tensor(gathered_tensor, video_grids.contiguous())
+    video_grids = gathered_tensor.cpu()
+
+    def get_file_name():
+        return os.path.join(args.save_img_path,
+                            f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
+    
+    if local_rank == 0:
+        if args.num_frames == 1:
+            save_image(video_grids / 255.0, get_file_name(),
+                    nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+        else:
+            video_grids = save_video_grid(video_grids)
+            imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6, codec='libx264',
+                    output_params=['-threads', '20'])
+
+    print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--max_sequence_length", type=int, default=512)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--refine_caption', action='store_true')
+    parser.add_argument('--compile', action='store_true')
+    parser.add_argument('--model_type', type=str, default="dit", choices=['dit', 'udit', 'latte'])
+    parser.add_argument('--save_memory', action='store_true')
+
+    parser.add_argument('--conditional_images_path', type=str, help='the path of txt file containing the paths of condition images')
+    args = parser.parse_args()
+
+    if torch_npu is not None:
+        npu_config.print_msg(args)
+
+    # 初始化分布式环境
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print('world_size', world_size)
+    if torch_npu is not None and npu_config.on_npu:
+        torch_npu.npu.set_device(local_rank)
+    else:
+        torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+
+    torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+    # print(11111111111111111111, local_rank, device)
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
+                                                  low_cpu_mem_usage=True, torch_dtype=weight_dtype).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+
+    # set eval mode
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler(clip_sample=False)
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler(clip_sample=False)
+    elif args.sample_method == 'DPMSolverMultistep':
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path, exist_ok=True)
+
+    if args.num_frames == 1:
+        video_length = 1
+        ext = 'jpg'
+    else:
+        ext = 'mp4'
+
+    latest_path = None
+    save_img_path = args.save_img_path
+
+    if latest_path is None:
+        latest_path = ''
+
+    full_path = f"{args.model_path}"
+    # full_path = f"{args.model_path}/{latest_path}/model_ema"
+    # full_path = f"{args.model_path}/{latest_path}/model"
+    pipeline = load_t2v_checkpoint(full_path)
+    pipeline = pipeline.to(weight_dtype)
+
+    print('load model')
+    if npu_config is not None and npu_config.on_npu and npu_config.profiling:
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+        )
+        profile_output_path = "/home/image_data/npu_profiling_t2v"
+        os.makedirs(profile_output_path, exist_ok=True)
+
+        with torch_npu.profiler.profile(
+                activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                with_stack=True,
+                record_shapes=True,
+                profile_memory=True,
+                experimental_config=experimental_config,
+                schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1,
+                                                        skip_first=0),
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+        ) as prof:
+            run_model_and_save_images(pipeline, latest_path)
+            prof.step()
+    else:
+        # print('gpu')
+        run_model_and_save_images(pipeline, latest_path)
diff --git a/opensora/sample/sample_inpaint_sp.py b/opensora/sample/sample_inpaint_sp.py
new file mode 100644
index 000000000..cc38c2cd3
--- /dev/null
+++ b/opensora/sample/sample_inpaint_sp.py
@@ -0,0 +1,347 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+import torch.distributed as dist
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderModel
+
+import os, sys
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+from opensora.models.diffusion.opensora.modeling_inpaint import OpenSoraInpaint
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_inpaint_sp import OpenSoraInpaintPipeline
+
+import imageio
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+    from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, hccl_info
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
+    pass
+import time
+
+# for validation
+import glob
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms import Lambda
+from opensora.dataset.transform import ToTensorVideo, CenterCropResizeVideo, TemporalRandomCrop, LongSideResizeVideo, SpatialStrideCropVideo
+import numpy as np
+from einops import rearrange
+
+
+def load_t2v_checkpoint(model_path):
+    print('load_t2v_checkpoint, ', model_path)
+    transformer_model = OpenSoraInpaint.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,)
+
+    # print(transformer_model.config)
+
+    # set eval mode
+    transformer_model.eval()
+    pipeline = OpenSoraInpaintPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model).to(device)
+
+
+    if args.compile:
+        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+        # options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
+        #             "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
+        #             "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
+        # #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # from onediffx import compile_pipe
+        # pipeline = compile_pipe(
+        #         pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
+        #     )
+
+        # 4%
+        pipeline.transformer = torch.compile(pipeline.transformer)
+    return pipeline
+
+
+def get_latest_path():
+    # Get the most recent checkpoint
+    dirs = os.listdir(args.model_path)
+    dirs = [d for d in dirs if d.startswith("checkpoint")]
+    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+    path = dirs[-1] if len(dirs) > 0 else None
+
+    return path
+
+
+def run_model_and_save_images(pipeline, model_path):
+
+    norm_fun = Lambda(lambda x: 2. * x - 1.)
+
+    resize = [CenterCropResizeVideo((args.height, args.width)), ]
+    transform = transforms.Compose([
+        ToTensorVideo(),
+        *resize, 
+        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+        norm_fun
+    ])
+
+    def preprocess_images(images):
+        print(images)
+        if len(images) == 1:
+            conditional_images_indices = [0]
+        elif len(images) == 2:
+            conditional_images_indices = [0, -1]
+        else:
+            print("Only support 1 or 2 condition images!")
+            raise NotImplementedError
+        
+        try:
+            conditional_images = [Image.open(image).convert("RGB") for image in images]
+            conditional_images = [torch.from_numpy(np.copy(np.array(image))) for image in conditional_images]
+            conditional_images = [rearrange(image, 'h w c -> c h w').unsqueeze(0) for image in conditional_images]
+            conditional_images = [transform(image).to(device=device, dtype=weight_dtype) for image in conditional_images]
+        except Exception as e:
+            print('Error when loading images')
+            print(f'condition images are {images}')
+            raise e
+        return dict(conditional_images=conditional_images, conditional_images_indices=conditional_images_indices)
+    
+
+    checkpoint_name = f"{os.path.basename(model_path)}"
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), (unwatermarked), 
+    {}. 
+    emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
+    sharp focus, high budget, cinemascope, moody, epic, gorgeous
+    """
+    
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path)
+        
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        text_prompt = [i.strip() for i in text_prompt]
+
+    if not isinstance(args.conditional_images_path, list):
+        args.conditional_images_path = [args.conditional_images_path]
+    if len(args.conditional_images_path) == 1 and args.conditional_images_path[0].endswith('txt'):
+        temp = open(args.conditional_images_path[0], 'r').readlines()
+        conditional_images = [i.strip().split(',') for i in temp]
+    
+    video_grids = []
+    for idx, (prompt, images) in enumerate(zip(text_prompt, conditional_images)):
+        pre_results = preprocess_images(images)
+        cond_imgs = pre_results['conditional_images']
+        cond_imgs_indices = pre_results['conditional_images_indices']
+
+        videos = pipeline(
+            conditional_images=cond_imgs,
+            conditional_images_indices=cond_imgs_indices,
+            prompt=positive_prompt.format(prompt),
+            negative_prompt=negative_prompt, 
+            num_frames=args.num_frames,
+            height=args.height,
+            width=args.width,
+            num_inference_steps=args.num_sampling_steps,
+            guidance_scale=args.guidance_scale,
+            num_images_per_prompt=1,
+            mask_feature=True,
+            device=args.device, 
+            max_sequence_length=args.max_sequence_length, 
+        ).images
+        
+        if nccl_info.rank <= 0:
+            try:
+                if args.num_frames == 1:
+                    ext = 'jpg'
+                    videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                    save_image(videos / 255.0, os.path.join(args.save_img_path, f'{idx}.{ext}'), nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+                else:
+                    ext = 'mp4'
+                    imageio.mimwrite(
+                        os.path.join(args.save_img_path, f'{idx}.{ext}'), videos[0], fps=args.fps, quality=6)  # highest quality is 10, lowest is 0
+                    save_path = os.path.join(args.save_img_path, f'{idx}.{ext}')
+                    print(f'image or video is saved at {save_path}')
+            except:
+                print('Error when saving {}'.format(prompt))
+            video_grids.append(videos)
+    if nccl_info.rank <= 0:
+        video_grids = torch.cat(video_grids, dim=0)
+
+        def get_file_name():
+            return os.path.join(args.save_img_path,
+                                f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
+
+        if args.num_frames == 1:
+            save_image(video_grids / 255.0, get_file_name(),
+                    nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+        else:
+            video_grids = save_video_grid(video_grids)
+            imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6)
+
+        print('save path {}'.format(args.save_img_path))
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--max_sequence_length", type=int, default=512)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--model_type', type=str, default="dit", choices=['dit', 'udit', 'latte'])
+    parser.add_argument('--compile', action='store_true')
+    parser.add_argument('--save_memory', action='store_true')
+
+    parser.add_argument('--conditional_images_path', type=str, help='the path of txt file containing the paths of condition images')
+    args = parser.parse_args()
+    args = parser.parse_args()
+
+    if torch_npu is not None:
+        npu_config.print_msg(args)
+
+    # 初始化分布式环境
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print('world_size', world_size)
+    if torch_npu is not None and npu_config.on_npu:
+        torch_npu.npu.set_device(local_rank)
+    else:
+        torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+    initialize_sequence_parallel_state(world_size)
+    # torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+    # print(11111111111111111111, local_rank, device)
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
+                                                  low_cpu_mem_usage=True, torch_dtype=weight_dtype).to(device)
+    tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+    # set eval mode
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler(clip_sample=False)
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler(clip_sample=False)
+    elif args.sample_method == 'DPMSolverMultistep':
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path, exist_ok=True)
+
+    if args.num_frames == 1:
+        video_length = 1
+        ext = 'jpg'
+    else:
+        ext = 'mp4'
+
+    latest_path = None
+    save_img_path = args.save_img_path
+
+    if latest_path is None:
+        latest_path = ''
+    # full_path = f"{args.model_path}/{latest_path}/model_ema"
+    full_path = f"{args.model_path}"
+    # full_path = f"{args.model_path}/{latest_path}/model"
+    pipeline = load_t2v_checkpoint(full_path)
+    pipeline = pipeline.to(weight_dtype)
+    print('load model')
+    if npu_config is not None and npu_config.on_npu and npu_config.profiling:
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+        )
+        profile_output_path = "/home/image_data/npu_profiling_t2v"
+        os.makedirs(profile_output_path, exist_ok=True)
+
+        with torch_npu.profiler.profile(
+                activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                with_stack=True,
+                record_shapes=True,
+                profile_memory=True,
+                experimental_config=experimental_config,
+                schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1,
+                                                        skip_first=0),
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+        ) as prof:
+            run_model_and_save_images(pipeline, latest_path)
+            prof.step()
+    else:
+        # print('gpu')
+        run_model_and_save_images(pipeline, latest_path)
diff --git a/opensora/sample/sample_t2v.py b/opensora/sample/sample_t2v.py
new file mode 100644
index 000000000..955c44bbf
--- /dev/null
+++ b/opensora/sample/sample_t2v.py
@@ -0,0 +1,236 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, MT5EncoderModel, UMT5EncoderModel, AutoTokenizer
+
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_opensora import OpenSoraPipeline
+
+import imageio
+
+
+def main(args):
+    # torch.manual_seed(args.seed)
+    # torch.backends.cuda.matmul.allow_tf32 = False
+    weight_dtype = torch.bfloat16
+    device = torch.device(args.device)
+
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    # if args.model_3d:
+    #     transformer_model = OpenSoraT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
+    # else:
+    #     transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
+    
+    if args.model_type == 'dit':
+        transformer_model = OpenSoraT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, 
+                                                        low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
+    elif args.model_type == 'udit':
+        transformer_model = UDiTT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, ignore_mismatched_sizes=True, 
+                                                        low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
+    else:
+        transformer_model = LatteT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, low_cpu_mem_usage=False, 
+                                                     device_map=None, torch_dtype=weight_dtype)
+        
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+    # text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
+    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
+    
+    
+    # set eval mode
+    transformer_model.eval()
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler(clip_sample=False)
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler(clip_sample=False)
+    elif args.sample_method == 'DPMSolverMultistep':
+        '''
+        DPM++ 2M	        DPMSolverMultistepScheduler	
+        DPM++ 2M Karras	    DPMSolverMultistepScheduler	init with use_karras_sigmas=True
+        DPM++ 2M SDE	    DPMSolverMultistepScheduler	init with algorithm_type="sde-dpmsolver++"
+        DPM++ 2M SDE Karras	DPMSolverMultistepScheduler	init with use_karras_sigmas=True and algorithm_type="sde-dpmsolver++"
+        
+        DPM++ SDE	        DPMSolverSinglestepScheduler	
+        DPM++ SDE Karras	DPMSolverSinglestepScheduler	init with use_karras_sigmas=True
+        DPM2	            KDPM2DiscreteScheduler	
+        DPM2 Karras	        KDPM2DiscreteScheduler	init with use_karras_sigmas=True
+        DPM2 a	            KDPM2AncestralDiscreteScheduler	
+        DPM2 a Karras	    KDPM2AncestralDiscreteScheduler	init with use_karras_sigmas=True
+        '''
+        # scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+    elif args.sample_method == 'EulerDiscreteSVD':
+        scheduler = EulerDiscreteScheduler.from_pretrained("stabilityai/stable-video-diffusion-img2vid", 
+                                                        subfolder="scheduler", cache_dir=args.cache_dir)
+    pipeline = OpenSoraPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model)
+    pipeline.to(device)
+    if args.compile:
+        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+        options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
+                    "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
+                    "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
+        #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
+        from onediffx import compile_pipe
+        pipeline = compile_pipe(
+                pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
+            )
+
+        # 4%
+        # pipeline.transformer = torch.compile(pipeline.transformer)
+    
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path)
+        
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        text_prompt = [i.strip() for i in text_prompt]
+
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), 
+    {}. 
+    """
+    
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+
+    # positive_prompt = """
+    # (masterpiece), (best quality), (ultra-detailed), 
+    # {}. 
+    # emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
+    # sharp focus, high budget, cinemascope, moody, epic, gorgeous
+    # """
+    
+    # negative_prompt = """
+    # disfigured, poorly drawn face, longbody, lowres, bad anatomy, bad hands, missing fingers, cropped, worst quality, low quality
+    # """
+
+    video_grids = []
+    for idx, prompt in enumerate(text_prompt):
+        videos = pipeline(positive_prompt.format(prompt),
+                          negative_prompt=negative_prompt, 
+                          num_frames=args.num_frames,
+                          height=args.height,
+                          width=args.width,
+                          num_inference_steps=args.num_sampling_steps,
+                          guidance_scale=args.guidance_scale,
+                          num_images_per_prompt=1,
+                          mask_feature=True,
+                          device=args.device, 
+                          max_sequence_length=args.max_sequence_length, 
+                          ).images
+        try:
+            if args.num_frames == 1:
+                ext = 'jpg'
+                videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                save_image(videos / 255.0, os.path.join(args.save_img_path, f'{idx}.{ext}'), nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+
+            else:
+                ext = 'mp4'
+                imageio.mimwrite(
+                    os.path.join(args.save_img_path, f'{idx}.{ext}'), videos[0], fps=args.fps, quality=6)  # highest quality is 10, lowest is 0
+        except:
+            print('Error when saving {}'.format(prompt))
+        video_grids.append(videos)
+    video_grids = torch.cat(video_grids, dim=0)
+
+
+    # torchvision.io.write_video(args.save_img_path + '_%04d' % args.run_time + '-.mp4', video_grids, fps=6)
+    if args.num_frames == 1:
+        save_image(video_grids / 255.0, os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
+                   nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+    else:
+        video_grids = save_video_grid(video_grids)
+        imageio.mimwrite(os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'), video_grids, fps=args.fps, quality=6)
+
+    print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=2.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--max_sequence_length", type=int, default=300)
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--run_time", type=int, default=0)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--compile', action='store_true')
+    parser.add_argument('--model_type', type=str, default="udit", choices=['dit', 'udit', 'latte'])
+    parser.add_argument('--save_memory', action='store_true')
+    args = parser.parse_args()
+
+    main(args)
\ No newline at end of file
diff --git a/opensora/sample/sample_t2v_ddp.py b/opensora/sample/sample_t2v_ddp.py
new file mode 100644
index 000000000..78ded3c3a
--- /dev/null
+++ b/opensora/sample/sample_t2v_ddp.py
@@ -0,0 +1,343 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+import torch.distributed as dist
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderModel
+
+import os, sys
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_opensora import OpenSoraPipeline
+
+import imageio
+
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+except:
+    torch_npu = None
+    npu_config = None
+    pass
+import time
+
+
+
+def load_t2v_checkpoint(model_path):
+    if args.model_type == 'udit':
+        transformer_model = UDiTT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,
+                                                        torch_dtype=weight_dtype)
+    elif args.model_type == 'dit':
+        transformer_model = OpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,
+                                                        torch_dtype=weight_dtype)
+    else:
+        transformer_model = LatteT2V.from_pretrained(model_path, cache_dir=args.cache_dir, low_cpu_mem_usage=False,
+                                                     device_map=None, torch_dtype=weight_dtype)
+    # print(transformer_model.config)
+
+    # set eval mode
+    transformer_model.eval()
+    pipeline = OpenSoraPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model).to(device)
+
+    if args.compile:
+        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+        # options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
+        #             "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
+        #             "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
+        # #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # from onediffx import compile_pipe
+        # pipeline = compile_pipe(
+        #         pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
+        #     )
+
+        # 4%
+        pipeline.transformer = torch.compile(pipeline.transformer)
+    return pipeline
+
+
+def get_latest_path():
+    # Get the most recent checkpoint
+    dirs = os.listdir(args.model_path)
+    dirs = [d for d in dirs if d.startswith("checkpoint")]
+    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+    path = dirs[-1] if len(dirs) > 0 else None
+
+    return path
+
+
+def run_model_and_save_images(pipeline, model_path):
+    video_grids = []
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        args.text_prompt = [i.strip() for i in text_prompt]
+    checkpoint_name = f"{os.path.basename(model_path)}"
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), (unwatermarked), 
+    {}. 
+    emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
+    sharp focus, high budget, cinemascope, moody, epic, gorgeous
+    """
+    
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+    
+
+    for index, prompt in enumerate(args.text_prompt):
+        if index % world_size != local_rank:
+            continue
+        if args.refine_caption:
+            q = f'Translate this brief generation prompt into a detailed caption: {prompt}'
+            query = f'[UNUSED_TOKEN_146]user\n{q}[UNUSED_TOKEN_145]\n[UNUSED_TOKEN_146]assistant\n'
+            # print(query)
+            with torch.cuda.amp.autocast(): 
+                refine_prompt = model_gen(refiner, query, None)
+            refine_prompt = refine_prompt.replace('<|im_end|>', '').replace('</s>', '')
+            input_prompt = positive_prompt.format(refine_prompt)
+            print(f'Processing the origin prompt({prompt})\n  '
+                  f'refine_prompt ({refine_prompt})\n  input_prompt ({input_prompt})\n  device ({device})')
+        else:
+            input_prompt = positive_prompt.format(prompt)
+            print(f'Processing the origin prompt({prompt})\n  '
+                  f'input_prompt ({input_prompt})\n device ({device})')
+        videos = pipeline(
+            input_prompt, 
+                          negative_prompt=negative_prompt, 
+                          num_frames=args.num_frames,
+                          height=args.height,
+                          width=args.width,
+                          num_inference_steps=args.num_sampling_steps,
+                          guidance_scale=args.guidance_scale,
+                          num_images_per_prompt=1,
+                          mask_feature=True,
+                          device=args.device,
+                          max_sequence_length=args.max_sequence_length,
+                          ).images
+        print(videos.shape)
+        try:
+            if args.num_frames == 1:
+                videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                save_image(videos / 255.0, os.path.join(args.save_img_path,
+                                                        f'{model_path}', f'{args.sample_method}_{index}_{checkpoint_name}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
+                           nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+
+            else:
+                print(f'save {model_path}/{args.sample_method}_{index}_{checkpoint_name}__gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}')
+                imageio.mimwrite(
+                    os.path.join(
+                        args.save_img_path,
+                        f'{model_path}', f'{args.sample_method}_{index}_{checkpoint_name}__gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
+                    ), videos[0],
+                    fps=args.fps, quality=6)  # highest quality is 10, lowest is 0
+        except:
+            print('Error when saving {}'.format(prompt))
+        video_grids.append(videos)
+    dist.barrier()
+    video_grids = torch.cat(video_grids, dim=0).cuda()
+    shape = list(video_grids.shape)
+    shape[0] *= world_size
+    gathered_tensor = torch.zeros(shape, dtype=video_grids.dtype, device=device)
+    dist.all_gather_into_tensor(gathered_tensor, video_grids.contiguous())
+    video_grids = gathered_tensor.cpu()
+
+    def get_file_name():
+        return os.path.join(args.save_img_path,
+                            f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
+    
+    if local_rank == 0:
+        if args.num_frames == 1:
+            save_image(video_grids / 255.0, get_file_name(),
+                    nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+        else:
+            video_grids = save_video_grid(video_grids)
+            imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6, codec='libx264',
+                    output_params=['-threads', '20'])
+
+    print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--max_sequence_length", type=int, default=512)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--refine_caption', action='store_true')
+    parser.add_argument('--compile', action='store_true')
+    parser.add_argument('--model_type', type=str, default="dit", choices=['dit', 'udit', 'latte'])
+    parser.add_argument('--save_memory', action='store_true')
+    args = parser.parse_args()
+
+    if torch_npu is not None:
+        npu_config.print_msg(args)
+
+    # 初始化分布式环境
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print('world_size', world_size)
+    if torch_npu is not None and npu_config.on_npu:
+        torch_npu.npu.set_device(local_rank)
+    else:
+        torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+
+    torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+    # print(11111111111111111111, local_rank, device)
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    # text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", 
+    #                                                cache_dir=args.cache_dir, low_cpu_mem_usage=True, 
+    #                                                torch_dtype=weight_dtype).to(device)
+    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", 
+    #                                           cache_dir=args.cache_dir)
+    
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
+                                                  low_cpu_mem_usage=True, torch_dtype=weight_dtype).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+
+    # set eval mode
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler(clip_sample=False)
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler(clip_sample=False)
+    elif args.sample_method == 'DPMSolverMultistep':
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path, exist_ok=True)
+
+    if args.num_frames == 1:
+        video_length = 1
+        ext = 'jpg'
+    else:
+        ext = 'mp4'
+
+    latest_path = None
+    save_img_path = args.save_img_path
+    # while True:
+        # cur_path = get_latest_path()
+        # print(cur_path, latest_path)
+        # if cur_path == latest_path:
+        #     time.sleep(5)
+        #     continue
+
+        # time.sleep(1)
+        # latest_path = cur_path
+        # os.makedirs(os.path.join(args.save_img_path, latest_path), exist_ok=True)
+
+        # if npu_config is not None:
+        #     npu_config.print_msg(f"The latest_path is {latest_path}")
+        # else:
+        #     print(f"The latest_path is {latest_path}")
+    if latest_path is None:
+        latest_path = ''
+
+    full_path = f"{args.model_path}"
+    # full_path = f"{args.model_path}/{latest_path}/model_ema"
+    # full_path = f"{args.model_path}/{latest_path}/model"
+    try:
+        pipeline = load_t2v_checkpoint(full_path)
+    except:
+        time.sleep(100)
+        pipeline = load_t2v_checkpoint(full_path)
+    # print('load model')
+    if npu_config is not None and npu_config.on_npu and npu_config.profiling:
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+        )
+        profile_output_path = "/home/image_data/npu_profiling_t2v"
+        os.makedirs(profile_output_path, exist_ok=True)
+
+        with torch_npu.profiler.profile(
+                activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                with_stack=True,
+                record_shapes=True,
+                profile_memory=True,
+                experimental_config=experimental_config,
+                schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1,
+                                                        skip_first=0),
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+        ) as prof:
+            run_model_and_save_images(pipeline, latest_path)
+            prof.step()
+    else:
+        # print('gpu')
+        run_model_and_save_images(pipeline, latest_path)
diff --git a/opensora/sample/sample_t2v_on_npu.py b/opensora/sample/sample_t2v_on_npu.py
new file mode 100644
index 000000000..297bd2f41
--- /dev/null
+++ b/opensora/sample/sample_t2v_on_npu.py
@@ -0,0 +1,295 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+import torch.distributed as dist
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, MT5EncoderModel, T5Tokenizer, AutoTokenizer
+
+import os, sys
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+# from opensora.models.diffusion.compress3d.modeling_opensora import CompressOpenSoraT2V
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_opensora import OpenSoraPipeline
+
+import imageio
+
+try:
+    import torch_npu
+except:
+    pass
+import time
+from opensora.npu_config import npu_config
+
+
+def load_t2v_checkpoint(model_path):
+    if args.model_3d:
+        # transformer_model = UDiTT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+        #                                                 low_cpu_mem_usage=False, device_map=None,
+        #                                                 torch_dtype=weight_dtype)
+        # if args.agent is not None or args.sparse:
+        #     transformer_model = CompressOpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+        #                                                     low_cpu_mem_usage=False, device_map=None,
+        #                                                     torch_dtype=weight_dtype)
+        # else:
+        transformer_model = OpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,
+                                                        torch_dtype=weight_dtype)
+    print(transformer_model.config)
+
+    # set eval mode
+    transformer_model.eval()
+    pipeline = OpenSoraPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model).to(device)
+
+    return pipeline
+
+
+def get_latest_path():
+    # Get the most recent checkpoint
+    dirs = os.listdir(args.model_path)
+    dirs = [d for d in dirs if d.startswith("checkpoint")]
+    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+    path = dirs[-1] if len(dirs) > 0 else None
+
+    return path
+
+
+def run_model_and_save_images(pipeline, model_path):
+    video_grids = []
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        args.text_prompt = [i.strip() for i in text_prompt]
+
+    checkpoint_name = f"{os.path.basename(model_path)}"
+
+    positive_prompt = "(masterpiece), (best quality), (ultra-detailed), {}. emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, sharp focus, high budget, cinemascope, moody, epic, gorgeous"
+    negative_prompt = "nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
+    prompts_per_rank = len(args.text_prompt) // npu_config.N_NPU_PER_NODE
+    for index, prompt in enumerate(args.text_prompt):
+        if index < prompts_per_rank * local_rank or index >= prompts_per_rank * (local_rank + 1):
+            continue
+        print('Processing the ({}) prompt'.format(prompt))
+        videos = pipeline(positive_prompt.format(prompt),
+                          negative_prompt=negative_prompt,
+                          num_frames=args.num_frames,
+                          height=args.height,
+                          width=args.width,
+                          num_inference_steps=args.num_sampling_steps,
+                          guidance_scale=args.guidance_scale,
+                          num_images_per_prompt=1,
+                          mask_feature=True,
+                          device=args.device,
+                          max_sequence_length=512,
+                          ).images
+        print(videos.shape)
+        try:
+            if args.num_frames == 1:
+                videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                save_image(videos / 255.0, os.path.join(args.save_img_path,
+                                                        f'{args.model_path.split(os.sep)[-1]}_{args.sample_method}_{index}_{checkpoint_name}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
+                           nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+
+            else:
+                imageio.mimwrite(
+                    os.path.join(
+                        args.save_img_path,
+                        f'{args.model_path.split(os.sep)[-1]}_{args.sample_method}_{index}_{checkpoint_name}__gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
+                    ), videos[0],
+                    fps=args.fps, quality=6, codec='libx264',
+                    output_params=['-threads', '20'])  # highest quality is 10, lowest is 0
+        except:
+            print('Error when saving {}'.format(prompt))
+        video_grids.append(videos)
+
+    video_grids = torch.cat(video_grids, dim=0).cuda()
+    shape = list(video_grids.shape)
+    shape[0] *= world_size
+    gathered_tensor = torch.zeros(shape, dtype=video_grids.dtype, device=device)
+    dist.all_gather_into_tensor(gathered_tensor, video_grids.contiguous())
+    video_grids = gathered_tensor.cpu()
+
+    # video_grids = video_grids.repeat(world_size, 1, 1, 1)
+    # output = torch.zeros(video_grids.shape, dtype=video_grids.dtype, device=device)
+    # dist.all_to_all_single(output, video_grids)
+    # video_grids = output.cpu()
+    def get_file_name():
+        return os.path.join(args.save_img_path,
+                            f'{args.model_path.split(os.sep)[-1]}_{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
+
+    if args.num_frames == 1:
+        save_image(video_grids / 255.0, get_file_name(),
+                   nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+    else:
+        video_grids = save_video_grid(video_grids)
+        imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6)
+
+    print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--run_time", type=int, default=0)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--model_3d', action='store_true')
+    parser.add_argument('--udit', action='store_true')
+    parser.add_argument('--save_memory', action='store_true')
+    parser.add_argument("--agent", type=str, default=None, help="agent attention")
+    parser.add_argument("--sparse", type=str, default=None, help="sparse attention")
+    args = parser.parse_args()
+    if args.agent:
+        npu_config.agent = npu_config.get_tokens(args.agent)
+    if args.sparse:
+        npu_config.sparse = npu_config.get_tokens(args.sparse)
+
+    npu_config.print_msg(args)
+    npu_config.conv_dtype = torch.bfloat16
+    replace_with_fp32_forwards()
+
+    # 初始化分布式环境
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    if npu_config.on_npu:
+        torch_npu.npu.set_device(local_rank)
+    dist.init_process_group(backend='hccl', init_method='env://', world_size=8, rank=local_rank)
+
+    # torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=torch.float32)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
+                                                  low_cpu_mem_usage=True, torch_dtype=torch.float16).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+    # set eval mode
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler()
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler()
+    elif args.sample_method == 'DPMSolverMultistep':
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path, exist_ok=True)
+
+    if args.num_frames == 1:
+        video_length = 1
+        ext = 'jpg'
+    else:
+        ext = 'mp4'
+
+    latest_path = None
+    save_img_path = args.save_img_path
+    first_in = False
+    while True:
+        cur_path = get_latest_path()
+        if cur_path == latest_path:
+            time.sleep(60)
+            continue
+
+        if not first_in:
+            first_in = True
+        else:
+            time.sleep(1200)
+
+        latest_path = cur_path
+
+        npu_config.print_msg(f"The latest_path is {latest_path}")
+        full_path = f"{args.model_path}/{latest_path}/model_ema"
+        # full_path = "/home/opensora/captions/240p_model_ema"
+        pipeline = load_t2v_checkpoint(full_path)
+
+        if npu_config.on_npu and npu_config.profiling:
+            experimental_config = torch_npu.profiler._ExperimentalConfig(
+                profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+                aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+            )
+            profile_output_path = "/home/image_data/npu_profiling_t2v"
+            os.makedirs(profile_output_path, exist_ok=True)
+
+            with torch_npu.profiler.profile(
+                    activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                    with_stack=True,
+                    record_shapes=True,
+                    profile_memory=True,
+                    experimental_config=experimental_config,
+                    schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1,
+                                                         skip_first=0),
+                    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+            ) as prof:
+                run_model_and_save_images(pipeline, latest_path)
+                prof.step()
+        else:
+            run_model_and_save_images(pipeline, latest_path)
diff --git a/opensora/sample/sample_t2v_sp.py b/opensora/sample/sample_t2v_sp.py
new file mode 100644
index 000000000..64cc10d74
--- /dev/null
+++ b/opensora/sample/sample_t2v_sp.py
@@ -0,0 +1,322 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+import torch.distributed as dist
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderModel
+
+import os, sys
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_opensora_sp import OpenSoraPipeline
+
+import imageio
+
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+    from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, hccl_info
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
+    pass
+import time
+
+
+
+def load_t2v_checkpoint(model_path):
+    print('load_t2v_checkpoint, ', model_path)
+    if args.model_type == 'udit':
+        transformer_model = UDiTT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,
+                                                        torch_dtype=weight_dtype)
+    elif args.model_type == 'dit':
+        transformer_model = OpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,
+                                                        torch_dtype=weight_dtype)
+    else:
+        transformer_model = LatteT2V.from_pretrained(model_path, cache_dir=args.cache_dir, low_cpu_mem_usage=False,
+                                                     device_map=None, torch_dtype=weight_dtype)
+    # print(transformer_model.config)
+
+    # set eval mode
+    transformer_model.eval()
+    pipeline = OpenSoraPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model).to(device)
+
+
+    if args.compile:
+        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
+        # options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
+        #             "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
+        #             "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
+        # #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
+        # from onediffx import compile_pipe
+        # pipeline = compile_pipe(
+        #         pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
+        #     )
+
+        # 4%
+        pipeline.transformer = torch.compile(pipeline.transformer)
+    return pipeline
+
+
+def get_latest_path():
+    # Get the most recent checkpoint
+    dirs = os.listdir(args.model_path)
+    dirs = [d for d in dirs if d.startswith("checkpoint")]
+    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+    path = dirs[-1] if len(dirs) > 0 else None
+
+    return path
+
+
+def run_model_and_save_images(pipeline, model_path):
+    video_grids = []
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        args.text_prompt = [i.strip() for i in text_prompt]
+
+    checkpoint_name = f"{os.path.basename(model_path)}"
+    positive_prompt = """
+    (masterpiece), (best quality), (ultra-detailed), (unwatermarked), 
+    {}. 
+    emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
+    sharp focus, high budget, cinemascope, moody, epic, gorgeous
+    """
+    
+    negative_prompt = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
+    """
+
+    for index, prompt in enumerate(args.text_prompt):
+        # if index % world_size != local_rank:
+        #     continue
+        # print('Processing the ({}) prompt, device ({})'.format(prompt, device))
+        videos = pipeline(positive_prompt.format(prompt),
+                        negative_prompt=negative_prompt, 
+                        num_frames=args.num_frames,
+                        height=args.height,
+                        width=args.width,
+                        num_inference_steps=args.num_sampling_steps,
+                        guidance_scale=args.guidance_scale,
+                        num_images_per_prompt=1,
+                        mask_feature=True,
+                        device=args.device,
+                        max_sequence_length=args.max_sequence_length,
+                        ).images
+        print(videos.shape)
+        
+        if nccl_info.rank <= 0:
+            try:
+                if args.num_frames == 1:
+                    videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                    save_image(videos / 255.0, os.path.join(args.save_img_path,
+                                                            f'{args.sample_method}_{index}_{checkpoint_name}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
+                            nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+
+                else:
+                    imageio.mimwrite(
+                        os.path.join(
+                            args.save_img_path,
+                            f'{args.sample_method}_{index}_{checkpoint_name}__gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
+                        ), videos[0],
+                        fps=args.fps, quality=6, codec='libx264',
+                        output_params=['-threads', '20'])  # highest quality is 10, lowest is 0
+            except:
+                print('Error when saving {}'.format(prompt))
+            video_grids.append(videos)
+    if nccl_info.rank <= 0:
+        video_grids = torch.cat(video_grids, dim=0)
+
+        def get_file_name():
+            return os.path.join(args.save_img_path,
+                                f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
+
+        if args.num_frames == 1:
+            save_image(video_grids / 255.0, get_file_name(),
+                    nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+        else:
+            video_grids = save_video_grid(video_grids)
+            imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6)
+
+        print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--max_sequence_length", type=int, default=512)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--model_type', type=str, default="dit", choices=['dit', 'udit', 'latte'])
+    parser.add_argument('--compile', action='store_true')
+    parser.add_argument('--save_memory', action='store_true')
+    args = parser.parse_args()
+
+    if torch_npu is not None:
+        npu_config.print_msg(args)
+
+    # 初始化分布式环境
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print('world_size', world_size)
+    if torch_npu is not None and npu_config.on_npu:
+        torch_npu.npu.set_device(local_rank)
+    else:
+        torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+    initialize_sequence_parallel_state(world_size)
+    # torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+    # print(11111111111111111111, local_rank, device)
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    # text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", 
+    #                                                cache_dir=args.cache_dir, low_cpu_mem_usage=True, 
+    #                                                torch_dtype=weight_dtype).to(device)
+    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", 
+    #                                           cache_dir=args.cache_dir)
+    
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
+                                                  low_cpu_mem_usage=True, torch_dtype=weight_dtype).to(device)
+    tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+    # set eval mode
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler(clip_sample=False)
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler(clip_sample=False)
+    elif args.sample_method == 'DPMSolverMultistep':
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path, exist_ok=True)
+
+    if args.num_frames == 1:
+        video_length = 1
+        ext = 'jpg'
+    else:
+        ext = 'mp4'
+
+    latest_path = None
+    save_img_path = args.save_img_path
+    # while True:
+        # cur_path = get_latest_path()
+        # print(cur_path, latest_path)
+        # if cur_path == latest_path:
+        #     time.sleep(5)
+        #     continue
+        # time.sleep(1)
+        # latest_path = cur_path
+
+        # if npu_config is not None:
+        #     npu_config.print_msg(f"The latest_path is {latest_path}")
+        # else:
+        #     print(f"The latest_path is {latest_path}")
+
+
+    if latest_path is None:
+        latest_path = ''
+    # full_path = f"{args.model_path}/{latest_path}/model_ema"
+    full_path = f"{args.model_path}"
+    # full_path = f"{args.model_path}/{latest_path}/model"
+    pipeline = load_t2v_checkpoint(full_path)
+    print('load model')
+    if npu_config is not None and npu_config.on_npu and npu_config.profiling:
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+        )
+        profile_output_path = "/home/image_data/npu_profiling_t2v"
+        os.makedirs(profile_output_path, exist_ok=True)
+
+        with torch_npu.profiler.profile(
+                activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                with_stack=True,
+                record_shapes=True,
+                profile_memory=True,
+                experimental_config=experimental_config,
+                schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1,
+                                                        skip_first=0),
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+        ) as prof:
+            run_model_and_save_images(pipeline, latest_path)
+            prof.step()
+    else:
+        # print('gpu')
+        run_model_and_save_images(pipeline, latest_path)
diff --git a/opensora/sample/sample_t2v_sp_on_npu.py b/opensora/sample/sample_t2v_sp_on_npu.py
new file mode 100644
index 000000000..b9f960f3f
--- /dev/null
+++ b/opensora/sample/sample_t2v_sp_on_npu.py
@@ -0,0 +1,273 @@
+import math
+import os
+import torch
+import argparse
+import torchvision
+import torch.distributed as dist
+
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
+from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
+from omegaconf import OmegaConf
+from torchvision.utils import save_image
+from transformers import T5EncoderModel, MT5EncoderModel, T5Tokenizer, AutoTokenizer
+
+import os, sys
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.adaptor.modules import replace_with_fp32_forwards
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
+from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import save_video_grid
+
+from opensora.sample.pipeline_opensora_sp import OpenSoraPipeline
+
+import imageio
+
+try:
+    import torch_npu
+    from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, get_sequence_parallel_state, hccl_info
+except:
+    pass
+import time
+from opensora.npu_config import npu_config
+
+
+def load_t2v_checkpoint(model_path):
+    if args.model_3d:
+        # transformer_model = UDiTT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+        #                                                 low_cpu_mem_usage=False, device_map=None,
+        #                                                 torch_dtype=weight_dtype)
+        transformer_model = OpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
+                                                        low_cpu_mem_usage=False, device_map=None,
+                                                        torch_dtype=weight_dtype)
+    print(transformer_model.config)
+
+    # set eval mode
+    transformer_model.eval()
+    pipeline = OpenSoraPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler,
+                                transformer=transformer_model).to(device)
+
+    return pipeline
+
+
+def get_latest_path():
+    # Get the most recent checkpoint
+    dirs = os.listdir(args.model_path)
+    dirs = [d for d in dirs if d.startswith("checkpoint")]
+    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+    path = dirs[-1] if len(dirs) > 0 else None
+
+    return path
+
+
+def run_model_and_save_images(pipeline, model_path):
+    video_grids = []
+    if not isinstance(args.text_prompt, list):
+        args.text_prompt = [args.text_prompt]
+    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
+        text_prompt = open(args.text_prompt[0], 'r').readlines()
+        args.text_prompt = [i.strip() for i in text_prompt]
+
+    checkpoint_name = f"{os.path.basename(model_path)}"
+
+    positive_prompt = "(masterpiece), (best quality), (ultra-detailed), {}. emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, sharp focus, high budget, cinemascope, moody, epic, gorgeous"
+    negative_prompt = "nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
+    for index, prompt in enumerate(args.text_prompt):
+        print('Processing the ({}) prompt'.format(prompt))
+        videos = pipeline(positive_prompt.format(prompt),
+                          negative_prompt=negative_prompt,
+                          num_frames=args.num_frames,
+                          height=args.height,
+                          width=args.width,
+                          num_inference_steps=args.num_sampling_steps,
+                          guidance_scale=args.guidance_scale,
+                          num_images_per_prompt=1,
+                          mask_feature=True,
+                          device=args.device,
+                          max_sequence_length=512,
+                          ).images
+        print(videos.shape)
+        if hccl_info.rank <= 0:
+            try:
+                if args.num_frames == 1:
+                    videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
+                    save_image(videos / 255.0, os.path.join(args.save_img_path,
+                                                            f'{args.model_path.split(os.sep)[-1]}_{args.sample_method}_{index}_{checkpoint_name}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
+                               nrow=1, normalize=True, value_range=(0, 1))  # t c h w
+
+                else:
+                    imageio.mimwrite(
+                        os.path.join(
+                            args.save_img_path,
+                            f'{args.model_path.split(os.sep)[-1]}_{args.sample_method}_{index}_{checkpoint_name}__gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
+                        ), videos[0],
+                        fps=args.fps, quality=6, codec='libx264',
+                        output_params=['-threads', '20'])  # highest quality is 10, lowest is 0
+            except:
+                print('Error when saving {}'.format(prompt))
+            video_grids.append(videos)
+    if hccl_info.rank <= 0:
+        video_grids = torch.cat(video_grids, dim=0)
+        def get_file_name():
+            return os.path.join(args.save_img_path,
+                                f'{args.model_path.split(os.sep)[-1]}_{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
+
+        if args.num_frames == 1:
+            save_image(video_grids / 255.0, get_file_name(),
+                    nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
+        else:
+            video_grids = save_video_grid(video_grids)
+            imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6)
+
+        print('save path {}'.format(args.save_img_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
+    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
+    parser.add_argument("--num_frames", type=int, default=1)
+    parser.add_argument("--height", type=int, default=512)
+    parser.add_argument("--width", type=int, default=512)
+    parser.add_argument("--device", type=str, default='cuda:0')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
+    parser.add_argument("--guidance_scale", type=float, default=7.5)
+    parser.add_argument("--sample_method", type=str, default="PNDM")
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument("--fps", type=int, default=24)
+    parser.add_argument("--run_time", type=int, default=0)
+    parser.add_argument("--text_prompt", nargs='+')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument('--model_3d', action='store_true')
+    parser.add_argument('--udit', action='store_true')
+    parser.add_argument('--save_memory', action='store_true')
+    args = parser.parse_args()
+
+    npu_config.print_msg(args)
+    npu_config.conv_dtype = torch.bfloat16
+    replace_with_fp32_forwards()
+
+    # 初始化分布式环境
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    if npu_config.on_npu:
+        torch_npu.npu.set_device(local_rank)
+    dist.init_process_group(backend='hccl', init_method='env://', world_size=8, rank=local_rank)
+    initialize_sequence_parallel_state(world_size)
+
+    # torch.manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    device = torch.cuda.current_device()
+
+    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
+    vae = CausalVAEModelWrapper(args.ae_path)
+    vae.vae = vae.vae.to(device=device, dtype=torch.float32)
+    if args.enable_tiling:
+        vae.vae.enable_tiling()
+        vae.vae.tile_overlap_factor = args.tile_overlap_factor
+        vae.vae.tile_sample_min_size = 512
+        vae.vae.tile_latent_min_size = 64
+        vae.vae.tile_sample_min_size_t = 29
+        vae.vae.tile_latent_min_size_t = 8
+        if args.save_memory:
+            vae.vae.tile_sample_min_size = 256
+            vae.vae.tile_latent_min_size = 32
+            vae.vae.tile_sample_min_size_t = 29
+            vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
+                                                  low_cpu_mem_usage=True, torch_dtype=torch.float16).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+
+    # set eval mode
+    vae.eval()
+    text_encoder.eval()
+
+    if args.sample_method == 'DDIM':  #########
+        scheduler = DDIMScheduler()
+    elif args.sample_method == 'EulerDiscrete':
+        scheduler = EulerDiscreteScheduler()
+    elif args.sample_method == 'DDPM':  #############
+        scheduler = DDPMScheduler()
+    elif args.sample_method == 'DPMSolverMultistep':
+        scheduler = DPMSolverMultistepScheduler()
+    elif args.sample_method == 'DPMSolverSinglestep':
+        scheduler = DPMSolverSinglestepScheduler()
+    elif args.sample_method == 'PNDM':
+        scheduler = PNDMScheduler()
+    elif args.sample_method == 'HeunDiscrete':  ########
+        scheduler = HeunDiscreteScheduler()
+    elif args.sample_method == 'EulerAncestralDiscrete':
+        scheduler = EulerAncestralDiscreteScheduler()
+    elif args.sample_method == 'DEISMultistep':
+        scheduler = DEISMultistepScheduler()
+    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
+        scheduler = KDPM2AncestralDiscreteScheduler()
+
+    if not os.path.exists(args.save_img_path):
+        os.makedirs(args.save_img_path, exist_ok=True)
+
+    if args.num_frames == 1:
+        video_length = 1
+        ext = 'jpg'
+    else:
+        ext = 'mp4'
+
+    latest_path = None
+    save_img_path = args.save_img_path
+    first_in = False
+    while True:
+        cur_path = get_latest_path()
+        if cur_path == latest_path:
+            time.sleep(60)
+            continue
+
+        if not first_in:
+            first_in = True
+        else:
+            time.sleep(1200)
+
+        latest_path = cur_path
+
+        npu_config.print_msg(f"The latest_path is {latest_path}")
+        full_path = f"{args.model_path}/{latest_path}/model_ema"
+        pipeline = load_t2v_checkpoint(full_path)
+
+        if npu_config.on_npu and npu_config.profiling:
+            experimental_config = torch_npu.profiler._ExperimentalConfig(
+                profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+                aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+            )
+            profile_output_path = "/home/image_data/npu_profiling_t2v"
+            os.makedirs(profile_output_path, exist_ok=True)
+
+            with torch_npu.profiler.profile(
+                    activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                    with_stack=True,
+                    record_shapes=True,
+                    profile_memory=True,
+                    experimental_config=experimental_config,
+                    schedule=torch_npu.profiler.schedule(wait=10000, warmup=0, active=1, repeat=1,
+                                                         skip_first=0),
+                    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+            ) as prof:
+                run_model_and_save_images(pipeline, latest_path)
+                prof.step()
+        else:
+            run_model_and_save_images(pipeline, latest_path)
diff --git a/opensora/serve/gradio_utils.py b/opensora/serve/gradio_utils.py
new file mode 100644
index 000000000..743aa0994
--- /dev/null
+++ b/opensora/serve/gradio_utils.py
@@ -0,0 +1,88 @@
+import random
+
+import imageio
+import uuid
+import torch
+
+import numpy as np
+
+MAX_SEED = np.iinfo(np.int32).max
+
+DESCRIPTION = """
+        <img src="https://s21.ax1x.com/2024/07/14/pk5pLBF.jpg" alt="Description" width="256" height="256">
+        
+        # Open-Sora Plan 93x720p
+        #### [Open-Sora Plan 93x720p](https://github.com/PKU-YuanGroup/Open-Sora-Plan) is a transformer-based text-to-video diffusion system trained on text embeddings from mT5. This demo uses the []() checkpoint.
+        #### Multilingual prompts SUPPORT; 支持多语言
+        #### Welcome to Star🌟 our [GitHub](https://github.com/PKU-YuanGroup/Open-Sora-Plan)
+        """
+NEG_PROMPT = """
+    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
+    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry. 
+    """
+
+
+style_list = [
+    {
+        "name": "(Default)",
+        "prompt": "(masterpiece), (best quality), (ultra-detailed), (unwatermarked), {prompt}",
+        "negative_prompt": NEG_PROMPT,
+    },
+    {
+        "name": "Cinematic",
+        "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured. ",
+    },
+    {
+        "name": "Photographic",
+        "prompt": "cinematic photo, a close-up of  {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly. ",
+    },
+    {
+        "name": "Anime",
+        "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed",
+        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast. ",
+    },
+    {
+        "name": "Manga",
+        "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style. ",
+    },
+    {
+        "name": "Digital Art",
+        "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
+        "negative_prompt": "photo, photorealistic, realism, ugly. ",
+    },
+    {
+        "name": "Pixel art",
+        "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
+        "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic. ",
+    },
+    {
+        "name": "Fantasy art",
+        "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
+        "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white. ",
+    },
+    {
+        "name": "Neonpunk",
+        "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
+        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured. ",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting. ",
+    },
+]
+
+
+
+def save_video(video):
+    unique_name = str(uuid.uuid4()) + ".mp4"
+    imageio.mimwrite(unique_name, video, fps=23, quality=6)
+    return unique_name
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
diff --git a/opensora/serve/gradio_web_server.py b/opensora/serve/gradio_web_server.py
new file mode 100644
index 000000000..f7af2d615
--- /dev/null
+++ b/opensora/serve/gradio_web_server.py
@@ -0,0 +1,310 @@
+#!/usr/bin/env python
+from __future__ import annotations
+import argparse
+import os
+import sys
+import gradio as gr
+from diffusers import ConsistencyDecoderVAE, DPMSolverMultistepScheduler, Transformer2DModel, AutoencoderKL, SASolverScheduler
+
+import torch
+from typing import Tuple
+from datetime import datetime
+# import spaces
+from opensora.models import CausalVAEModelWrapper
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config
+from opensora.models.causalvideovae import ae_norm, ae_denorm
+from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderModel
+from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
+                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
+                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
+                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
+from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
+from opensora.sample.pipeline_opensora import OpenSoraPipeline
+from opensora.serve.gradio_utils import DESCRIPTION, MAX_SEED, style_list, randomize_seed_fn, save_video
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
+CACHE_EXAMPLES = False
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "6000"))
+MAX_VIDEO_FRAME = int(os.getenv("MAX_IMAGE_SIZE", "93"))
+SPEED_UP_T5 = os.getenv("USE_TORCH_COMPILE", "0") == "1"
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
+ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+PORT = int(os.getenv("DEMO_PORT", "15432"))
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+STYLE_NAMES = list(styles.keys())
+DEFAULT_STYLE_NAME = "(Default)"
+SCHEDULE_NAME = [
+    "PNDM-Solver", "EulerA-Solver", "DPM-Solver", "SA-Solver", 
+    "DDIM-Solver", "Euler-Solver", "DDPM-Solver", "DEISM-Solver"]
+DEFAULT_SCHEDULE_NAME = "EulerA-Solver"
+NUM_IMAGES_PER_PROMPT = 1
+
+
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_path", type=str, required=True)
+parser.add_argument("--ae_path", type=str, required=True)
+parser.add_argument("--ae", type=str, default="CausalVAEModel_D4_4x8x8")
+parser.add_argument("--text_encoder_name", type=str, default="google/mt5-xxl")
+parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+args = parser.parse_args()
+
+def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    if not negative:
+        negative = ""
+    return p.replace("{prompt}", positive), n + negative
+
+if torch.cuda.is_available():
+    weight_dtype = torch.bfloat16
+    T5_token_max_length = 512
+
+    vae = CausalVAEModelWrapper(args.ae_path, args.cache_dir).eval()
+    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
+    vae.vae.enable_tiling()
+    vae.vae.tile_overlap_factor = 0.125
+    vae.vae.tile_sample_min_size = 256
+    vae.vae.tile_latent_min_size = 32
+    vae.vae.tile_sample_min_size_t = 29
+    vae.vae.tile_latent_min_size_t = 8
+    vae.vae_scale_factor = ae_stride_config[args.ae]
+
+    # text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir, 
+    #                                                low_cpu_mem_usage=True, torch_dtype=weight_dtype)
+    # tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+    text_encoder = MT5EncoderModel.from_pretrained("google/mt5-xxl", cache_dir=args.cache_dir, 
+                                                   low_cpu_mem_usage=True, torch_dtype=weight_dtype)
+    tokenizer = AutoTokenizer.from_pretrained("google/mt5-xxl", cache_dir=args.cache_dir)
+    transformer = OpenSoraT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, low_cpu_mem_usage=False, 
+                                              device_map=None, torch_dtype=weight_dtype)
+    scheduler = EulerAncestralDiscreteScheduler()
+    pipe = OpenSoraPipeline(vae=vae,
+                                text_encoder=text_encoder,
+                                tokenizer=tokenizer,
+                                scheduler=scheduler, 
+                                transformer=transformer)
+    pipe.to(device)
+    print("Loaded on Device!")
+
+    # speed-up T5
+    if SPEED_UP_T5:
+        pipe.text_encoder.to_bettertransformer()
+
+    if USE_TORCH_COMPILE:
+        pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=True)
+        print("Model Compiled!")
+
+# @spaces.GPU(duration=120)
+@torch.no_grad()
+@torch.inference_mode()
+def generate(
+        prompt: str,
+        negative_prompt: str = "",
+        style: str = DEFAULT_STYLE_NAME,
+        use_negative_prompt: bool = False,
+        seed: int = 0,
+        # frame: int = 29,
+        schedule: str = 'DPM-Solver',
+        guidance_scale: float = 4.5,
+        num_inference_steps: int = 25,
+        randomize_seed: bool = False,
+        progress=gr.Progress(track_tqdm=True),
+):
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator().manual_seed(seed)
+
+    if schedule == 'DPM-Solver':
+        if not isinstance(pipe.scheduler, DPMSolverMultistepScheduler):
+            pipe.scheduler = DPMSolverMultistepScheduler()
+    elif schedule == "PNDM-Solver":
+        if not isinstance(pipe.scheduler, PNDMScheduler):
+            pipe.scheduler = PNDMScheduler()
+    elif schedule == "DDIM-Solver":
+        if not isinstance(pipe.scheduler, DDIMScheduler):
+            pipe.scheduler = DDIMScheduler()
+    elif schedule == "Euler-Solver":
+        if not isinstance(pipe.scheduler, EulerDiscreteScheduler):
+            pipe.scheduler = EulerDiscreteScheduler()
+    elif schedule == "DDPM-Solver":
+        if not isinstance(pipe.scheduler, DDPMScheduler):
+            pipe.scheduler = DDPMScheduler()
+    elif schedule == "EulerA-Solver":
+        if not isinstance(pipe.scheduler, EulerAncestralDiscreteScheduler):
+            pipe.scheduler = EulerAncestralDiscreteScheduler()
+    elif schedule == "DEISM-Solver":
+        if not isinstance(pipe.scheduler, DEISMultistepScheduler):
+            pipe.scheduler = DEISMultistepScheduler()
+    elif schedule == "SA-Solver":
+        if not isinstance(pipe.scheduler, SASolverScheduler):
+            pipe.scheduler = SASolverScheduler.from_config(pipe.scheduler.config, algorithm_type='data_prediction', tau_func=lambda t: 1 if 200 <= t <= 800 else 0, predictor_order=2, corrector_order=2)
+    else:
+        raise ValueError(f"Unknown schedule: {schedule}")
+
+    if not use_negative_prompt:
+        negative_prompt = None  # type: ignore
+    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
+    print(prompt, negative_prompt)
+    videos = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_frames=93,
+        width=1280,
+        height=720,
+        # width=640,
+        # height=480,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        num_images_per_prompt=1,  # num_imgs
+        max_sequence_length=T5_token_max_length,
+    ).images
+
+    video_paths = [save_video(vid) for vid in videos]
+    print(video_paths)
+    return video_paths[0], seed
+
+
+examples = [
+    "A young man at his 20s is sitting on a piece of cloud in the sky, reading a book.", 
+    "an gray-haired man with a beard in his 60s, he is deep in thought pondering the history of the universe as he sits at a cafe in Paris, his eyes focus on people offscreen as they walk as he sits mostly motionless, he is dressed in a wool coat suit coat with a button-down shirt, he wears a brown beret and glasses and has a very professorial appearance, and the end he offers a subtle closed-mouth smile as if he found the answer to the mystery of life, the lighting is very cinematic with the golden light and the Parisian streets and city in the background, depth of field, cinematic 35mm film.", 
+    "a woman’s face, illuminated by the soft light of dawn, her expression serene and content as she wakes up in a cozy bedroom.", 
+    "a detective’s face, lit by a single desk lamp, his eyes scanning a wall covered in photos and notes, deep in thought.", 
+    "Audience members in a theater are captured in a series of medium shots, with a young man and woman in formal attire centrally positioned and illuminated by a spotlight effect.", 
+    "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.", 
+    "a realistic 3d rendering of a female character with curly blonde hair and blue eyes. she is wearing a black tank top and has a neutral expression while facing the camera directly. the background is a plain blue sky, and the scene is devoid of any other objects or text. the character is detailed, with realistic textures and lighting, suitable for a video game or high-quality animation. there is no movement or additional action in the video. the focus is entirely on the character's appearance and realistic rendering.", 
+    "A panda strumming a guitar under a bamboo grove, its paws gently plucking the strings as a group of mesmerized rabbits watch, the music blending with the rustle of bamboo leaves. HD.", 
+    "a woman with a vintage hairstyle and bright red lipstick, gazing seductively into the camera, the background blurred to keep the focus solely on her.", 
+    "In the jungle, a hidden temple stands guarded by statues of lions, their eyes glowing with emerald light, protecting secrets untold for millennia. 8K.", 
+    "an old man’s weathered face, with deep wrinkles and a thick white mustache, looking out to sea, the wind gently blowing through his hair.", 
+    "a soldier’s face, covered in dirt and sweat, his eyes filled with determination as he surveys the battlefield.", 
+    "A river that flows uphill, defying gravity as it returns lost treasures from the sea to the mountain top, each item telling a story of a voyage gone by. HD.", 
+    "a man’s face, lit only by the glow of his computer screen, his eyes wide and unblinking as he discovers something shocking online.", 
+    "On a deserted island, palm trees sway to summon a rainstorm, their leaves conducting the wind like maestros, orchestrating a symphony of thunder and lightning. High Resolution.", 
+    "a middle-aged man’s face, with a five o’clock shadow, staring pensively into the distance as rain softly taps against the window beside him, his thoughts deep and contemplative.", 
+    "a man’s face, his expression one of deep concentration as he works on a complex task.", 
+    "Drone view of waves crashing against the rugged cliffs along Big Sur's garay point beach.The crashing blue waters create white-tipped waves,while the golden light of the setting sun illuminates the rocky shore. A small island with a lighthouse sits in the distance, and green", 
+    "shrubbery covers the cliffs edge. The steep drop from the road down to the beach is adramatic feat, with the cliff's edges jutting out over the sea. This is a view that captures the raw beauty of the coast and the rugged landscape of the Pacific Coast Highway.", 
+    "a woman standing in a dimly lit room. she is wearing a traditional chinese outfit, which includes a red and gold dress with intricate designs and a matching headpiece. the woman has her hair styled in an updo, adorned with a gold accessory. her makeup is done in a way that accentuates her features, with red lipstick and dark eyeshadow. she is looking directly at the camera with a neutral expression. the room has a rustic feel, with wooden beams and a stone wall visible in the background. the lighting in the room is soft and warm, creating a contrast with the woman's vibrant attire. there are no texts or other objects in the video. the style of the video is a portrait, focusing on the woman and her attire.",
+]
+
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+    with gr.Row(equal_height=False):
+        with gr.Group():
+            with gr.Row():
+                use_negative_prompt = gr.Checkbox(label="Use additional negative prompt", value=False, visible=True)
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a additional negative prompt",
+                visible=True,
+            )
+            with gr.Row(visible=True):
+                schedule = gr.Radio(
+                    show_label=True,
+                    container=True,
+                    interactive=True,
+                    choices=SCHEDULE_NAME,
+                    value=DEFAULT_SCHEDULE_NAME,
+                    label="Sampler Schedule",
+                    visible=True,
+                )
+            style_selection = gr.Radio(
+                show_label=True,
+                container=True,
+                interactive=True,
+                choices=STYLE_NAMES,
+                value=DEFAULT_STYLE_NAME,
+                label="Video Style",
+            )
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=1,
+                    maximum=20,
+                    step=0.1,
+                    value=10.0,
+                )
+                inference_steps = gr.Slider(
+                    label="inference steps",
+                    minimum=10,
+                    maximum=200,
+                    step=1,
+                    value=50,
+                )
+        with gr.Group():
+            with gr.Row():
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt",
+                    container=False,
+                )
+                run_button = gr.Button("Run", scale=0)
+            result = gr.Video(label="Result")
+
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=[result, seed],
+        fn=generate,
+        cache_examples=CACHE_EXAMPLES,
+    )
+
+    use_negative_prompt.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt,
+        outputs=negative_prompt,
+        api_name=False,
+    )
+
+    gr.on(
+        triggers=[
+            prompt.submit,
+            negative_prompt.submit,
+            run_button.click,
+        ],
+        fn=generate,
+        inputs=[
+            prompt,
+            negative_prompt,
+            style_selection,
+            use_negative_prompt,
+            seed,
+            # frame,
+            schedule,
+            guidance_scale,
+            inference_steps,
+            randomize_seed,
+        ],
+        outputs=[result, seed],
+        api_name="run",
+    )
+
+if __name__ == "__main__":
+    
+
+    # demo.queue(max_size=20).launch(server_name='0.0.0.0', share=True)
+    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=11900, debug=True)
\ No newline at end of file
diff --git a/opensora/serve/style.css b/opensora/serve/style.css
new file mode 100644
index 000000000..73054b14b
--- /dev/null
+++ b/opensora/serve/style.css
@@ -0,0 +1 @@
+.gradio-container{width:1280px!important}
\ No newline at end of file
diff --git a/opensora/train/train_causalvae.py b/opensora/train/train_causalvae.py
new file mode 100644
index 000000000..37897bf17
--- /dev/null
+++ b/opensora/train/train_causalvae.py
@@ -0,0 +1,583 @@
+import os
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler, Subset
+import argparse
+import logging
+from colorlog import ColoredFormatter
+import tqdm
+from itertools import chain
+import wandb
+import random
+import numpy as np
+from pathlib import Path
+from einops import rearrange
+from causalvideovae.model import CausalVAEModel, EMA
+from causalvideovae.utils.utils import RealVideoDataset
+from causalvideovae.model.dataset_videobase import VideoDataset
+from causalvideovae.model.utils.module_utils import resolve_str_to_obj
+from causalvideovae.model.utils.video_utils import tensor_to_video
+import time
+try:
+    import lpips
+except:
+    raise Exception("Need lpips to valid.")
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def ddp_setup():
+    dist.init_process_group(backend="nccl")
+    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
+
+
+def setup_logger(rank):
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    formatter = ColoredFormatter(
+        f"[rank{rank}] %(log_color)s%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        log_colors={
+            "DEBUG": "cyan",
+            "INFO": "green",
+            "WARNING": "yellow",
+            "ERROR": "red",
+            "CRITICAL": "bold_red",
+        },
+        reset=True,
+        style="%",
+    )
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.DEBUG)
+    stream_handler.setFormatter(formatter)
+
+    if not logger.handlers:
+        logger.addHandler(stream_handler)
+
+    return logger
+
+
+def check_unused_params(model):
+    unused_params = []
+    for name, param in model.named_parameters():
+        if param.grad is None:
+            unused_params.append(name)
+    return unused_params
+
+
+def set_requires_grad_optimizer(optimizer, requires_grad):
+    for param_group in optimizer.param_groups:
+        for param in param_group["params"]:
+            param.requires_grad = requires_grad
+
+
+def total_params(model):
+    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params_in_millions = total_params / 1e6
+    return int(total_params_in_millions)
+
+
+def get_exp_name(args):
+    return f"{args.exp_name}-lr{args.lr:.2e}-bs{args.batch_size}-rs{args.resolution}-sr{args.sample_rate}-fr{args.num_frames}"
+
+def set_train(modules):
+    for module in modules:
+        module.train()
+            
+def set_eval(modules):
+    for module in modules:
+        module.eval()
+
+def set_modules_requires_grad(modules, requires_grad):
+    for module in modules:
+        module.requires_grad_(requires_grad)
+
+def save_checkpoint(
+    epoch,
+    batch_idx,
+    optimizer_state,
+    state_dict,
+    scaler_state,
+    checkpoint_dir,
+    filename="checkpoint.ckpt",
+    ema_state_dict={}
+):
+    filepath = checkpoint_dir / Path(filename)
+    torch.save(
+        {
+            "epoch": epoch,
+            "batch_idx": batch_idx,
+            "optimizer_state": optimizer_state,
+            "state_dict": state_dict,
+            "ema_state_dict": ema_state_dict,
+            "scaler_state": scaler_state,
+        },
+        filepath,
+    )
+    return filepath
+
+
+def valid(rank, model, val_dataloader, precision, args):
+    if args.eval_lpips:
+        lpips_model = lpips.LPIPS(net='alex', spatial=True)
+        lpips_model.to(rank)
+        lpips_model = DDP(lpips_model, device_ids=[rank])
+        lpips_model.requires_grad_(False)
+        lpips_model.eval()
+        
+    bar = None
+    if rank == 0:
+        bar = tqdm.tqdm(total=len(val_dataloader), desc="Validation...")
+        
+    psnr_list = []
+    lpips_list = []
+    video_log = []
+    num_video_log = args.eval_num_video_log
+    
+    with torch.no_grad():
+        for batch_idx, batch in enumerate(val_dataloader):
+            inputs = batch['video'].to(rank)
+            with torch.cuda.amp.autocast(dtype=precision):
+                video_recon, _ = model(inputs)
+            
+            # Upload videos
+            if rank == 0:
+                for i in range(len(video_recon)):
+                    if num_video_log <= 0:
+                        break
+                    video = tensor_to_video(video_recon[i])
+                    video_log.append(video)
+                    num_video_log -= 1
+                    
+            inputs = rearrange(inputs, "b c t h w -> (b t) c h w").contiguous()
+            video_recon = rearrange(video_recon, "b c t h w -> (b t) c h w").contiguous()
+            
+            # Calculate PSNR
+            mse = torch.mean(torch.square(inputs - video_recon), dim=(1,2,3))
+            psnr = 20 * torch.log10(1 / torch.sqrt(mse))
+            psnr = psnr.mean().detach().cpu().item()
+            
+            # Calculate LPIPS
+            if args.eval_lpips:
+                lpips_score = lpips_model.forward(inputs, video_recon).mean().detach().cpu().item()
+                lpips_list.append(lpips_score)
+                
+            psnr_list.append(psnr)
+            if rank == 0:
+                bar.update()
+    # Release gpus memory
+    torch.cuda.empty_cache()
+    return psnr_list, lpips_list, video_log
+
+def gather_valid_result(psnr_list, lpips_list, video_log_list, rank, world_size):
+    gathered_psnr_list = [None for _ in range(world_size)]
+    gathered_lpips_list = [None for _ in range(world_size)]
+    gathered_video_logs = [None for _ in range(world_size)]
+    
+    dist.all_gather_object(gathered_psnr_list, psnr_list)
+    dist.all_gather_object(gathered_lpips_list, lpips_list)
+    dist.all_gather_object(gathered_video_logs, video_log_list)
+    return np.array(gathered_psnr_list).mean(), np.array(gathered_lpips_list).mean(), list(chain(*gathered_video_logs))
+
+def train(args):
+    # Setup logger
+    ddp_setup()
+    rank = int(os.environ["LOCAL_RANK"])
+    logger = setup_logger(rank)
+
+    # Init
+    ckpt_dir = Path(args.ckpt_dir) / Path(get_exp_name(args))
+    if rank == 0:
+        try:
+            ckpt_dir.mkdir(exist_ok=False, parents=True)
+        except:
+            logger.warning(f"`{ckpt_dir}` exists!")
+            time.sleep(5)
+            
+        logger.warning("Connecting to WANDB...")
+        wandb.init(
+            project=os.environ.get("WANDB_PROJECT", "causalvideovae"),
+            config=args,
+            name=get_exp_name(args)
+        )
+    dist.barrier()
+    
+    # Load generator model
+    if args.pretrained_model_name_or_path is not None:
+        if rank == 0:
+            logger.warning(
+                f"You are loading a checkpoint from `{args.pretrained_model_name_or_path}`."
+            )
+        model = CausalVAEModel.from_pretrained(
+            args.pretrained_model_name_or_path, ignore_mismatched_sizes=False
+        )
+    elif args.model_config is not None:
+        if rank == 0:
+            logger.warning(f"Model will be inited randomly.")
+        model = CausalVAEModel.from_config(args.model_config)
+    else:
+        raise Exception(
+            "You should set either `--pretrained_model_name_or_path` or `--model_config`"
+        )
+
+    # Load discriminator model
+    disc_cls = resolve_str_to_obj(args.disc_cls, append=False)
+    logger.warning(f"disc_class: {args.disc_cls} perceptual_weight: {args.perceptual_weight}  loss_type: {args.loss_type}")
+    disc = disc_cls(
+        disc_start=args.disc_start,
+        disc_weight=args.disc_weight,
+        kl_weight=args.kl_weight,
+        logvar_init=args.logvar_init,
+        perceptual_weight=args.perceptual_weight,
+        loss_type=args.loss_type
+    )
+
+    # DDP
+    model = model.to(rank)
+    model = DDP(
+        model, device_ids=[rank], find_unused_parameters=args.find_unused_parameters
+    )
+    disc = disc.to(rank)
+    disc = DDP(
+        disc, device_ids=[rank], find_unused_parameters=args.find_unused_parameters
+    )
+    
+    # Load dataset
+    dataset = VideoDataset(
+        args.video_path,
+        sequence_length=args.num_frames,
+        resolution=args.resolution,
+        sample_rate=args.sample_rate,
+        dynamic_sample=args.dynamic_sample,
+    )
+    ddp_sampler = DistributedSampler(dataset)
+    dataloader = DataLoader(
+        dataset, batch_size=args.batch_size, sampler=ddp_sampler, pin_memory=True, num_workers=args.dataset_num_worker
+    )
+
+    val_dataset = RealVideoDataset(
+        real_video_dir=args.eval_video_path,
+        num_frames=args.eval_num_frames,
+        sample_rate=args.eval_sample_rate,
+        crop_size=args.eval_resolution,
+        resolution=args.eval_resolution,
+    )
+    indices = range(args.eval_subset_size)
+    val_dataset = Subset(val_dataset, indices=indices)
+    val_sampler = DistributedSampler(val_dataset)
+    val_dataloader = DataLoader(val_dataset, batch_size=args.eval_batch_size, sampler=val_sampler, pin_memory=True)
+    
+    
+    # Optimizer
+    modules_to_train = [module for module in model.module.get_decoder()]
+    if not args.freeze_encoder:
+        modules_to_train += [module for module in model.module.get_encoder()]
+    else:
+        for module in model.module.get_encoder():
+            module.eval()
+            module.requires_grad_(False)
+        logger.warning("Encoder is freezed!")
+        
+    parameters_to_train = []
+    for module in modules_to_train:
+        parameters_to_train += module.parameters()
+        
+    gen_optimizer = torch.optim.Adam(parameters_to_train, lr=args.lr)
+    disc_optimizer = torch.optim.Adam(
+        disc.module.discriminator.parameters(), lr=args.lr
+    )
+
+    # AMP scaler
+    scaler = torch.cuda.amp.GradScaler()
+    precision = torch.bfloat16
+    if args.mix_precision == "fp16":
+        precision = torch.float16
+    elif args.mix_precision == "fp32":
+        precision = torch.float32
+
+    # Load from checkpoint
+    start_epoch = 0
+    start_batch_idx = 0
+    if args.resume_from_checkpoint:
+        if not os.path.isfile(args.resume_from_checkpoint):
+            raise Exception(
+                f"Make sure `{args.resume_from_checkpoint}` is a ckpt file."
+            )
+        checkpoint = torch.load(args.resume_from_checkpoint, map_location="cpu")
+        if "ema_state_dict" in checkpoint and len(checkpoint['ema_state_dict']) > 0 and os.environ.get("NOT_USE_EMA_MODEL", 0) == 0:
+            logger.info("Load from EMA state dict! If you want to load from original state dict, you should set NOT_USE_EMA_MODEL=1.")
+            # model.module.load_state_dict(checkpoint["ema_state_dict"])
+            ckpt = {key.replace("module.", ""): value for key, value in checkpoint["ema_state_dict"].items()}
+            missing_keys, unexpected_keys = model.module.load_state_dict(ckpt, strict=True)
+        else:
+            missing_keys, unexpected_keys = model.module.load_state_dict(checkpoint["state_dict"]["gen_model"], strict=True)
+        print(f"missing_keys: {len(missing_keys)}, unexpected_keys: {len(unexpected_keys)}")
+        disc.module.load_state_dict(checkpoint["state_dict"]["dics_model"], strict=True)
+        scaler.load_state_dict(checkpoint["scaler_state"])
+        gen_optimizer.load_state_dict(checkpoint["optimizer_state"]["gen_optimizer"])
+        disc_optimizer.load_state_dict(checkpoint["optimizer_state"]["disc_optimizer"])
+        if not args.not_resume_training_process:
+            start_epoch = checkpoint["epoch"]
+            start_batch_idx = checkpoint.get("batch_idx", 0)
+            logger.info(
+                f"Checkpoint loaded from {args.resume_from_checkpoint}, starting from epoch {start_epoch} batch {start_batch_idx}"
+            )
+        else:
+            logger.warning(
+                f"Checkpoint loaded from {args.resume_from_checkpoint}, starting from epoch {start_epoch} batch {start_batch_idx}. But training process is not resumed."
+            )
+    
+    if args.ema:
+        logger.warning(f"Start with EMA. EMA decay = {args.ema_decay}.")
+        ema = EMA(model, args.ema_decay)
+        ema.register()
+            
+    # Training loop
+    logger.info("Prepared!")
+    dist.barrier()
+    if rank == 0:
+        logger.info(f"=== Model Params ===")
+        logger.info(f"Generator:\t\t{total_params(model.module)}M")
+        logger.info(f"\t- Encoder:\t{total_params(model.module.encoder):d}M")
+        logger.info(f"\t- Decoder:\t{total_params(model.module.decoder):d}M")
+        logger.info(f"Discriminator:\t{total_params(disc.module):d}M")
+        logger.info(f"===========")
+        logger.info(f"Precision is set to: {args.mix_precision}!")
+        logger.info("Start training!")
+
+    # Training Bar
+    bar_desc = ""
+    bar = None
+    if rank == 0:
+        max_steps = (
+            args.epochs * len(dataloader) if args.max_steps is None else args.max_steps
+        )
+        bar = tqdm.tqdm(total=max_steps, desc=bar_desc.format(current_epoch=0, loss=0))
+        bar_desc = "Epoch: {current_epoch}, Loss: {loss}"
+        logger.warning("Training Details: ")
+        logger.warning(f" Max steps: {max_steps}")
+        logger.warning(f" Dataset Samples: {len(dataloader)}")
+        logger.warning(
+            f" Total Batch Size: {args.batch_size} * {os.environ['WORLD_SIZE']}"
+        )
+    dist.barrier()
+
+    # Training Loop
+    num_epochs = args.epochs
+    current_step = 1
+
+    def update_bar(bar):
+        if rank == 0:
+            bar.desc = bar_desc.format(current_epoch=epoch, loss=f"-")
+            bar.update()
+    
+    for epoch in range(num_epochs):
+        set_train(modules_to_train)
+        ddp_sampler.set_epoch(epoch)  # Shuffle data at every epoch
+        for batch_idx, batch in enumerate(dataloader):
+
+            if epoch <= start_epoch and batch_idx < start_batch_idx:
+                update_bar(bar)
+                current_step += 1
+                continue
+
+            inputs = batch["video"].to(rank)
+            if (
+                current_step % 2 == 1
+                and current_step >= disc.module.discriminator_iter_start
+            ):
+                set_modules_requires_grad(modules_to_train, False)
+                step_gen = False
+                step_dis = True
+            else:
+                set_modules_requires_grad(modules_to_train, True)
+                step_gen = True
+                step_dis = False
+
+            assert (
+                step_gen or step_dis
+            ), "You should backward either Gen or Dis in a step."
+
+            with torch.cuda.amp.autocast(dtype=precision):
+                outputs, posterior = model(inputs)
+
+            # Generator Step
+            if step_gen:
+                with torch.cuda.amp.autocast(dtype=precision):
+                    g_loss, g_log = disc(
+                        inputs,
+                        outputs,
+                        posterior,
+                        optimizer_idx=0,
+                        global_step=current_step,
+                        last_layer=model.module.get_last_layer(),
+                        split="train",
+                    )
+                gen_optimizer.zero_grad()
+                scaler.scale(g_loss).backward()
+                scaler.step(gen_optimizer)
+                scaler.update()
+                if args.ema:
+                    ema.update()
+                if rank == 0 and current_step % args.log_steps == 0:
+                    wandb.log({"train/generator_loss": g_loss.item()}, step=current_step)
+
+            # Discriminator Step
+            if step_dis:
+                with torch.cuda.amp.autocast(dtype=precision):
+                    d_loss, d_log = disc(
+                        inputs,
+                        outputs,
+                        posterior,
+                        optimizer_idx=1,
+                        global_step=current_step,
+                        last_layer=None,
+                        split="train",
+                    )
+                disc_optimizer.zero_grad()
+                scaler.scale(d_loss).backward()
+                scaler.step(disc_optimizer)
+                scaler.update()
+                if rank == 0 and current_step % args.log_steps == 0:
+                    wandb.log({"train/discriminator_loss": d_loss.item()}, step=current_step)
+
+            def valid_model(model, name=""):
+                set_eval(modules_to_train)
+                psnr_list, lpips_list, video_log = valid(rank, model, val_dataloader, precision, args)
+                valid_psnr, valid_lpips, valid_video_log = gather_valid_result(psnr_list, lpips_list, video_log, rank, dist.get_world_size())
+                if rank == 0:
+                    name = "_" + name if name != "" else name
+                    print(f"psnr: {valid_psnr}   lpips: {valid_lpips}")
+                    wandb.log({f"val{name}/recon": wandb.Video(np.array(valid_video_log), fps=10)}, step=current_step)
+                    wandb.log({f"val{name}/psnr": valid_psnr}, step=current_step)
+                    wandb.log({f"val{name}/lpips": valid_lpips}, step=current_step)
+                    logger.info(f"{name} Validation done.")
+            
+            if current_step % args.eval_steps == 0 or current_step == 1:
+                if rank == 0:
+                    logger.info("Starting validation...")
+                valid_model(model)
+                if args.ema:
+                    ema.apply_shadow()
+                    valid_model(model, "ema")
+                    ema.restore()
+                        
+            # Checkpoint
+            if current_step % args.save_ckpt_step == 0 and rank == 0:
+                file_path = save_checkpoint(
+                    epoch,
+                    batch_idx,
+                    {
+                        "gen_optimizer": gen_optimizer.state_dict(),
+                        "disc_optimizer": disc_optimizer.state_dict(),
+                    },
+                    {
+                        "gen_model": model.module.state_dict(),
+                        "dics_model": disc.module.state_dict(),
+                    },
+                    scaler.state_dict(),
+                    ckpt_dir,
+                    f"checkpoint-{current_step}.ckpt",
+                    ema_state_dict=ema.shadow if args.ema else {}
+                )
+                logger.info(f"Checkpoint has been saved to `{file_path}`.")
+
+            # Update step
+            update_bar(bar)
+            current_step += 1
+
+    dist.destroy_process_group()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Distributed Training")
+    # Exp setting
+    parser.add_argument(
+        "--exp_name", type=str, default="test", help="number of epochs to train"
+    )
+    parser.add_argument("--seed", type=int, default=1234, help="seed")
+    # Training setting
+    parser.add_argument(
+        "--epochs", type=int, default=100, help="number of epochs to train"
+    )
+    parser.add_argument(
+        "--max_steps", type=int, default=None, help="number of epochs to train"
+    )
+    parser.add_argument("--save_ckpt_step", type=int, default=1000, help="")
+    parser.add_argument("--ckpt_dir", type=str, default="./results/", help="")
+    parser.add_argument(
+        "--batch_size", type=int, default=1, help="batch size for training"
+    )
+    parser.add_argument("--lr", type=float, default=1e-5, help="learning rate")
+    parser.add_argument("--log_steps", type=int, default=5, help="log steps")
+    parser.add_argument("--freeze_encoder", action="store_true", help="")
+    
+    # Data
+    parser.add_argument("--video_path", type=str, default=None, help="")
+    parser.add_argument("--num_frames", type=int, default=17, help="")
+    parser.add_argument("--resolution", type=int, default=256, help="")
+    parser.add_argument("--sample_rate", type=int, default=2, help="")
+    parser.add_argument("--dynamic_sample", type=bool, default=False, help="")
+    # Generator model
+    parser.add_argument("--find_unused_parameters", action="store_true", help="")
+    parser.add_argument(
+        "--pretrained_model_name_or_path", type=str, default=None, help=""
+    )
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="")
+    parser.add_argument("--not_resume_training_process", action="store_true", help="")
+    parser.add_argument("--model_config", type=str, default=None, help="")
+    parser.add_argument(
+        "--mix_precision",
+        type=str,
+        default="bf16",
+        choices=["fp16", "bf16", "fp32"],
+        help="precision for training",
+    )
+
+    # Discriminator Model
+    parser.add_argument("--load_disc_from_checkpoint", type=str, default=None, help="")
+    parser.add_argument(
+        "--disc_cls",
+        type=str,
+        default="causalvideovae.model.losses.LPIPSWithDiscriminator3D",
+        help="",
+    )
+    parser.add_argument("--disc_start", type=int, default=5, help="")
+    parser.add_argument("--disc_weight", type=float, default=0.5, help="")
+    parser.add_argument("--kl_weight", type=float, default=1e-06, help="")
+    parser.add_argument("--perceptual_weight", type=float, default=1.0, help="")
+    parser.add_argument("--loss_type", type=str, default="l1", help="")
+    parser.add_argument("--logvar_init", type=float, default=0.0, help="")
+
+    # Validation
+    parser.add_argument("--eval_steps", type=int, default=1000, help="")
+    parser.add_argument("--eval_video_path", type=str, default=None, help="")
+    parser.add_argument("--eval_num_frames", type=int, default=17, help="")
+    parser.add_argument("--eval_resolution", type=int, default=256, help="")
+    parser.add_argument("--eval_sample_rate", type=int, default=1, help="")
+    parser.add_argument("--eval_batch_size", type=int, default=8, help="")
+    parser.add_argument("--eval_subset_size", type=int, default=100, help="")
+    parser.add_argument("--eval_num_video_log", type=int, default=2, help="")
+    parser.add_argument("--eval_lpips", action="store_true", help="")
+    
+    # Dataset
+    parser.add_argument("--dataset_num_worker", type=int, default=16, help="")
+    
+    # EMA
+    parser.add_argument("--ema", action="store_true", help="")
+    parser.add_argument("--ema_decay", type=float, default=0.999, help="")
+
+    args = parser.parse_args()
+
+    set_random_seed(args.seed)
+    train(args)
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/opensora/train/train_inpaint.py b/opensora/train/train_inpaint.py
new file mode 100644
index 000000000..eb581a232
--- /dev/null
+++ b/opensora/train/train_inpaint.py
@@ -0,0 +1,876 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+A minimal training script for DiT using PyTorch DDP.
+"""
+import argparse
+import logging
+import math
+import os
+import shutil
+from pathlib import Path
+from typing import Optional
+import gc
+import numpy as np
+from einops import rearrange
+from tqdm import tqdm
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+    from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, \
+        destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
+    from opensora.acceleration.communications import prepare_parallel_data, broadcast
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import initialize_sequence_parallel_state, \
+        destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
+    from opensora.utils.communications import prepare_parallel_data, broadcast
+    pass
+import time
+from dataclasses import field, dataclass
+from torch.utils.data import DataLoader
+from copy import deepcopy
+import accelerate
+import torch
+from torch.nn import functional as F
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
+from packaging import version
+from tqdm.auto import tqdm
+
+import diffusers
+from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import compute_snr
+from diffusers.utils import check_min_version, is_wandb_available
+
+
+from opensora.utils.ema import EMAModel
+from opensora.dataset import getdataset
+from opensora.models import CausalVAEModelWrapper
+from opensora.models.text_encoder import get_text_enc, get_text_warpper
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config
+from opensora.models.causalvideovae import ae_norm, ae_denorm
+from opensora.models.diffusion import Diffusion_models, Diffusion_models_class
+from opensora.utils.dataset_utils import Collate, LengthGroupedSampler
+
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.24.0")
+logger = get_logger(__name__)
+
+class ProgressInfo:
+    def __init__(self, global_step, train_loss=0.0):
+        self.global_step = global_step
+        self.train_loss = train_loss
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    # use LayerNorm, GeLu, SiLu always as fp32 mode
+    if args.enable_stable_fp32:
+        replace_with_fp32_forwards()
+    if torch_npu is not None and npu_config is not None:
+        npu_config.print_msg(args)
+        npu_config.seed_everything(args.seed)
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if args.num_frames != 1 and args.use_image_num == 0:
+        initialize_sequence_parallel_state(args.sp_size)
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now.
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training we cast all non-trainable weigths to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # Create model:
+    kwargs = {}
+    ae = CausalVAEModelWrapper(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval()
+    if args.enable_tiling:
+        ae.vae.enable_tiling()
+        ae.vae.tile_overlap_factor = args.tile_overlap_factor
+
+    kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True}
+    text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval()
+
+    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
+    ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w)
+    assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})"
+    args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w
+    args.ae_stride = args.ae_stride_h
+    patch_size = args.model[-3:]
+    patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2])
+    args.patch_size = patch_size_h
+    args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w
+    assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})"
+    # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})."
+    assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})."
+    assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})."
+
+    args.stride_t = ae_stride_t * patch_size_t
+    args.stride = ae_stride_h * patch_size_h
+    latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w)
+    ae.latent_size = latent_size
+
+    if args.num_frames % 2 == 1:
+        args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1
+    else:
+        latent_size_t = args.num_frames // ae_stride_t
+
+    model_kwargs = {'vae_scale_factor_t': ae_stride_t}
+
+    model = Diffusion_models[args.model](
+        in_channels=ae_channel_config[args.ae],
+        out_channels=ae_channel_config[args.ae],
+        # caption_channels=4096,
+        # cross_attention_dim=1152,
+        attention_bias=True,
+        sample_size=latent_size,
+        sample_size_t=latent_size_t,
+        num_vector_embeds=None,
+        activation_fn="gelu-approximate",
+        num_embeds_ada_norm=1000,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        double_self_attention=False,
+        upcast_attention=False,
+        # norm_type="ada_norm_single",
+        norm_elementwise_affine=False,
+        norm_eps=1e-6,
+        attention_type='default',
+        attention_mode=args.attention_mode,
+        interpolation_scale_h=args.interpolation_scale_h,
+        interpolation_scale_w=args.interpolation_scale_w,
+        interpolation_scale_t=args.interpolation_scale_t,
+        downsampler=args.downsampler,
+        # compress_kv_factor=args.compress_kv_factor,
+        use_rope=args.use_rope,
+        # model_max_length=args.model_max_length,
+        use_stable_fp32=args.enable_stable_fp32,
+        **model_kwargs,
+    )
+    model.gradient_checkpointing = args.gradient_checkpointing
+
+    pretrained_transformer_model_path = args.pretrained_transformer_model_path
+    pretrained_model_path = dict(transformer_model=pretrained_transformer_model_path)
+    if pretrained_transformer_model_path is not None:
+        model.custom_load_state_dict(pretrained_model_path)
+
+    noise_scheduler = DDPMScheduler()
+
+    # Freeze main models
+    ae.vae.requires_grad_(False)
+    text_enc.requires_grad_(False)
+
+    model.train()
+
+    # Move unet, vae and text_encoder to device and cast to weight_dtype
+    # The VAE is in float32 to avoid NaN losses.
+    ae.vae.to(accelerator.device, dtype=torch.float32)
+    # ae.vae.to(accelerator.device, dtype=weight_dtype)
+    text_enc.to(accelerator.device, dtype=weight_dtype)
+
+    # Create EMA for the unet.
+    if args.use_ema:
+        ema_model = deepcopy(model)
+        ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step,
+                             model_cls=Diffusion_models_class[args.model], model_config=ema_model.config)
+    
+     # `accelerate` 0.16.0 will have better support for customized saving
+    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
+        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+        def save_model_hook(models, weights, output_dir):
+            if accelerator.is_main_process:
+                if args.use_ema:
+                    ema_model.save_pretrained(os.path.join(output_dir, "model_ema"))
+
+                for i, model in enumerate(models):
+                    model.save_pretrained(os.path.join(output_dir, "model"))
+                    if weights:  # Don't pop if empty
+                        # make sure to pop weight so that corresponding model is not saved again
+                        weights.pop()
+
+
+        def load_model_hook(models, input_dir):
+            # loading ema with customed 'from_pretrained' function
+            if args.use_ema:
+                load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model])
+                ema_model.load_state_dict(load_model.state_dict())
+                ema_model.to(accelerator.device)
+                del load_model
+
+            for i in range(len(models)):
+                # pop models so that they are not loaded again
+                model = models.pop()
+
+                # load diffusers style into model
+                load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model")
+                model.register_to_config(**load_model.config)
+
+                model.load_state_dict(load_model.state_dict())
+                del load_model
+
+        accelerator.register_save_state_pre_hook(save_model_hook)
+        accelerator.register_load_state_pre_hook(load_model_hook)
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    if args.scale_lr:
+        args.learning_rate = (
+                args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+        )
+
+    params_to_optimize = list(filter(lambda p: p.requires_grad, model.parameters()))
+    # Optimizer creation
+    if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"):
+        logger.warning(
+            f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]."
+            "Defaulting to adamW"
+        )
+        args.optimizer = "adamw"
+
+    if args.use_8bit_adam and not args.optimizer.lower() == "adamw":
+        logger.warning(
+            f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was "
+            f"set to {args.optimizer.lower()}"
+        )
+
+    if args.optimizer.lower() == "adamw":
+        if args.use_8bit_adam:
+            try:
+                import bitsandbytes as bnb
+            except ImportError:
+                raise ImportError(
+                    "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+                )
+
+            optimizer_class = bnb.optim.AdamW8bit
+        else:
+            optimizer_class = torch.optim.AdamW
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+        )
+
+    if args.optimizer.lower() == "prodigy":
+        try:
+            import prodigyopt
+        except ImportError:
+            raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+
+        optimizer_class = prodigyopt.Prodigy
+
+        if args.learning_rate <= 0.1:
+            logger.warning(
+                "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
+            )
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            beta3=args.prodigy_beta3,
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            decouple=args.prodigy_decouple,
+            use_bias_correction=args.prodigy_use_bias_correction,
+            safeguard_warmup=args.prodigy_safeguard_warmup,
+        )
+    logger.info(f"optimizer: {optimizer}")
+    
+    # Setup data:
+    train_dataset = getdataset(args)
+    logger.info(f"train_dataset: {train_dataset.__class__.__name__}")
+    sampler = LengthGroupedSampler(
+                args.train_batch_size,
+                world_size=accelerator.num_processes,
+                lengths=train_dataset.lengths, 
+                group_frame=args.group_frame, 
+                group_resolution=args.group_resolution, 
+            ) if args.group_frame or args.group_resolution else None
+    train_dataloader = DataLoader(
+        train_dataset,
+        shuffle=sampler is None,
+        # pin_memory=True,
+        collate_fn=Collate(args),
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        sampler=sampler if args.group_frame or args.group_resolution else None, 
+        drop_last=True, 
+        # prefetch_factor=4
+    )
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+
+    # Prepare everything with our `accelerator`.
+    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, lr_scheduler
+    )
+    if args.use_ema:
+        ema_model.to(accelerator.device)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    # NOTE wandb
+    if accelerator.is_main_process:
+        logger.info("init trackers...")
+        project_name = os.getenv('PROJECT', os.path.basename(args.output_dir))
+        entity = os.getenv('ENTITY', None)
+        run_name = os.getenv('WANDB_NAME', None)
+        init_kwargs = {
+            "entity": entity,
+            "run_name": run_name,
+        }
+        accelerator.init_trackers(project_name=project_name, config=vars(args), init_kwargs=init_kwargs)
+
+    # Train!
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size
+    logger.info("***** Running training *****")
+    logger.info(f"  Model = {model}")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Total trainable parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B")
+    global_step = 0
+    first_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+            if npu_config is not None:
+                train_dataset.n_used_elements = global_step * args.train_batch_size
+
+    else:
+        initial_global_step = 0
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=initial_global_step,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=not accelerator.is_local_main_process,
+    )
+    progress_info = ProgressInfo(global_step, train_loss=0.0)
+
+    def sync_gradients_info(loss):
+        # Checks if the accelerator has performed an optimization step behind the scenes
+        if args.use_ema:
+            ema_model.step(model.parameters())
+        progress_bar.update(1)
+        progress_info.global_step += 1
+        end_time = time.time()
+        one_step_duration = end_time - start_time
+        accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step)
+        if torch_npu is not None and npu_config is not None:
+            npu_config.print_msg(f"Step: [{progress_info.global_step}], local_loss={loss.detach().item()}, "
+                                 f"train_loss={progress_info.train_loss}, time_cost={one_step_duration}",
+                                 rank=0)
+        progress_info.train_loss = 0.0
+
+        # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
+        if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
+            if progress_info.global_step % args.checkpointing_steps == 0:
+                # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+                if accelerator.is_main_process and args.checkpoints_total_limit is not None:
+                    checkpoints = os.listdir(args.output_dir)
+                    checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+                    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+                    # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+                    if len(checkpoints) >= args.checkpoints_total_limit:
+                        num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+                        removing_checkpoints = checkpoints[0:num_to_remove]
+
+                        logger.info(
+                            f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+                        )
+                        logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+                        for removing_checkpoint in removing_checkpoints:
+                            removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+                            shutil.rmtree(removing_checkpoint)
+
+                save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}")
+                accelerator.save_state(save_path)
+                logger.info(f"Saved state to {save_path}")
+
+        logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+        progress_bar.set_postfix(**logs)
+
+    def run(model_input, model_kwargs, prof):
+        global start_time
+        start_time = time.time()
+
+        try:
+            in_channels = ae_channel_config[args.ae]
+            model_input, masked_x, video_mask = model_input[:, 0:in_channels], model_input[:, in_channels:2 * in_channels], model_input[:, 2 * in_channels:]
+        except:
+            raise ValueError("masked_x and video_mask is None!")
+
+        noise = torch.randn_like(model_input)
+        if args.noise_offset:
+            # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+            noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1),
+                                                     device=model_input.device)
+
+        bsz = model_input.shape[0]
+        current_step_frame = model_input.shape[2]
+        # Sample a random timestep for each image without bias.
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device)
+        if current_step_frame != 1 and get_sequence_parallel_state():  # image do not need sp
+            broadcast(timesteps)
+
+        # Add noise to the model input according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+
+        noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)
+
+        model_pred = model(
+            torch.cat([noisy_model_input, masked_x, video_mask], dim=1),
+            timesteps,
+            **model_kwargs,
+        )[0]
+       
+        # Get the target for loss depending on the prediction type
+        if args.prediction_type is not None:
+            # set prediction_type of scheduler if defined
+            noise_scheduler.register_to_config(prediction_type=args.prediction_type)
+
+        if noise_scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif noise_scheduler.config.prediction_type == "v_prediction":
+            target = noise_scheduler.get_velocity(model_input, noise, timesteps)
+        elif noise_scheduler.config.prediction_type == "sample":
+            # We set the target to latents here, but the model_pred will return the noise sample prediction.
+            target = model_input
+            # We will have to subtract the noise residual from the prediction to get the target sample.
+            model_pred = model_pred - noise
+        else:
+            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+
+        mask = model_kwargs.get('attention_mask', None)
+        if torch.all(mask.bool()):
+            mask = None
+        if get_sequence_parallel_state():
+            assert mask is None
+        b, c, _, _, _ = model_pred.shape
+        if mask is not None:
+            mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float()  # b t h w -> b c t h w
+            mask = mask.reshape(b, -1)
+        if args.snr_gamma is None:
+            # model_pred: b c t h w, attention_mask: b t h w
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+            loss = loss.reshape(b, -1)
+            if mask is not None:
+                loss = (loss * mask).sum() / mask.sum()  # mean loss on unpad patches
+            else:
+                loss = loss.mean()
+        else:
+            # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
+            # Since we predict the noise instead of x_0, the original formulation is slightly changed.
+            # This is discussed in Section 4.2 of the same paper.
+            snr = compute_snr(noise_scheduler, timesteps)
+            mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(
+                dim=1
+            )[0]
+            if noise_scheduler.config.prediction_type == "epsilon":
+                mse_loss_weights = mse_loss_weights / snr
+            elif noise_scheduler.config.prediction_type == "v_prediction":
+                mse_loss_weights = mse_loss_weights / (snr + 1)
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+            loss = loss.reshape(b, -1)
+            mse_loss_weights = mse_loss_weights.reshape(b, 1)
+            if mask is not None:
+                loss = (loss * mask * mse_loss_weights).sum() / mask.sum()  # mean loss on unpad patches
+            else:
+                loss = (loss * mse_loss_weights).mean()
+
+        # Gather the losses across all processes for logging (if we use distributed training).
+        avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+        progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps
+
+        # Backpropagate
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+            params_to_clip = params_to_optimize
+            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad()
+
+        if accelerator.sync_gradients:
+            sync_gradients_info(loss)
+
+        if prof is not None:
+            prof.step()
+
+
+        return loss
+
+    def train_one_step(step_, data_item_, prof_=None):
+        train_loss = 0.0
+        x, attn_mask, input_ids, cond_mask = data_item_
+        # assert torch.all(attn_mask.bool()), 'must all visible'
+        # Sample noise that we'll add to the latents
+        # import ipdb;ipdb.set_trace()
+        if args.group_frame or args.group_resolution:
+            if not args.group_frame:
+                each_latent_frame = torch.any(attn_mask.flatten(-2), dim=-1).int().sum(-1).tolist()
+                # logger.info(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask '
+                #             f'each_latent_frame: {each_latent_frame}')
+                logger.info(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask '
+                            f'each_latent_frame: {each_latent_frame}')
+        assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))'
+        x = x.to(accelerator.device, dtype=ae.vae.dtype)  # B 3*C T H W, 16
+
+        attn_mask = attn_mask.to(accelerator.device)  # B T H W
+        input_ids = input_ids.to(accelerator.device)  # B 1 L
+        cond_mask = cond_mask.to(accelerator.device)  # B 1 L
+
+        with torch.no_grad():
+            B, N, L = input_ids.shape  # B 1 L
+            # use batch inference
+            input_ids_ = input_ids.reshape(-1, L)
+            cond_mask_ = cond_mask.reshape(-1, L)
+            cond = text_enc(input_ids_, cond_mask_)  # B 1 L D
+            cond = cond.reshape(B, N, L, -1)
+
+            def preprocess_x_for_inpaint(x):
+                # NOTE vae-styled mask, deprecated
+                if args.use_vae_preprocessed_mask:
+                    x, masked_x, mask = x[:, :3], x[:, 3:6], x[:, 6:9]
+                    x, masked_x, mask = ae.encode(x), ae.encode(masked_x), ae.encode(mask)
+                else:
+                    x, masked_x, mask = x[:, :3], x[:, 3:6], x[:, 6:7]
+                    x, masked_x = ae.encode(x), ae.encode(masked_x)
+                    batch_size, channels, frame, height, width = mask.shape
+                    mask = rearrange(mask, 'b c t h w -> (b c t) 1 h w')
+                    mask = F.interpolate(mask, size=latent_size, mode='bilinear')
+                    mask = rearrange(mask, '(b c t) 1 h w -> b c t h w', t=frame, b=batch_size)
+                    mask_first_frame = mask[:, :, 0:1].repeat(1, 1, ae_stride_t, 1, 1).contiguous()
+                    mask = torch.cat([mask_first_frame, mask[:, :, 1:]], dim=2)
+                    mask = mask.view(batch_size, ae_stride_t, latent_size_t, latent_size[0], latent_size[1]).contiguous()
+
+                return x, masked_x, mask
+
+            # Map input images to latent space + normalize latents
+            x, masked_x, mask = preprocess_x_for_inpaint(x) # B 3*C T H W -> (B C T H W) * 3 
+            x = torch.cat([x, masked_x, mask], dim=1) # (B C T H W) * 3 -> B 3*C T H W
+
+        current_step_frame = x.shape[2]
+        current_step_sp_state = get_sequence_parallel_state()
+        if args.sp_size != 1:  # enable sp
+            if current_step_frame == 1:  # but image do not need sp
+                set_sequence_parallel_state(False)
+            else:
+                set_sequence_parallel_state(True)
+        if get_sequence_parallel_state():
+            x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask,
+                                                                                 args.use_image_num)
+            for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size):
+                with accelerator.accumulate(model):
+                    st_idx = iter * args.train_sp_batch_size
+                    ed_idx = (iter + 1) * args.train_sp_batch_size
+                    model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx],
+                                        attention_mask=attn_mask[st_idx: ed_idx],
+                                        encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num)
+                    run(x[st_idx: ed_idx], model_kwargs, prof_)
+
+        else:
+            with accelerator.accumulate(model):
+                assert not torch.any(torch.isnan(x)), 'after vae'
+                x = x.to(weight_dtype)
+                model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask,
+                                    encoder_attention_mask=cond_mask, use_image_num=args.use_image_num,)
+                run(x, model_kwargs, prof_)
+
+        set_sequence_parallel_state(current_step_sp_state)  # in case the next step use sp, which need broadcast(timesteps)
+
+        if progress_info.global_step >= args.max_train_steps:
+            return True
+
+        return False
+
+    def train_all_epoch(prof_=None):
+        for epoch in range(first_epoch, args.num_train_epochs):
+            progress_info.train_loss = 0.0
+            if progress_info.global_step >= args.max_train_steps:
+                return True
+
+            for step, data_item in enumerate(train_dataloader):
+                if train_one_step(step, data_item, prof_):
+                    break
+
+                if step >= 2 and torch_npu is not None and npu_config is not None:
+                    npu_config.free_mm()
+
+    if npu_config is not None and npu_config.on_npu and npu_config.profiling:
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+        )
+        profile_output_path = f"/home/image_data/npu_profiling_t2v/{os.getenv('PROJECT_NAME', 'local')}"
+        os.makedirs(profile_output_path, exist_ok=True)
+
+        with torch_npu.profiler.profile(
+                activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                with_stack=True,
+                record_shapes=True,
+                profile_memory=True,
+                experimental_config=experimental_config,
+                schedule=torch_npu.profiler.schedule(wait=npu_config.profiling_step, warmup=0, active=1, repeat=1,
+                                                     skip_first=0),
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+        ) as prof:
+            train_all_epoch(prof)
+    else:
+        train_all_epoch()
+    accelerator.wait_for_everyone()
+    accelerator.end_training()
+    if npu_config is not None and get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+     # dataset & dataloader
+    parser.add_argument("--dataset", type=str, required=True)
+    parser.add_argument("--data", type=str, required='')
+    parser.add_argument("--sample_rate", type=int, default=1)
+    parser.add_argument("--train_fps", type=int, default=24)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--num_frames", type=int, default=65)
+    parser.add_argument("--max_height", type=int, default=320)
+    parser.add_argument("--max_width", type=int, default=240)
+    parser.add_argument("--use_img_from_vid", action="store_true")
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--model_max_length", type=int, default=512)
+    parser.add_argument('--cfg', type=float, default=0.1)
+    parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.")
+    parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.")
+    parser.add_argument("--group_frame", action="store_true")
+    parser.add_argument("--group_resolution", action="store_true")
+
+    # text encoder & vae & diffusion model
+    parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122")
+    parser.add_argument('--enable_8bit_t5', action='store_true')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument("--compress_kv", action="store_true")
+    parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers")
+    parser.add_argument('--use_rope', action='store_true')
+    parser.add_argument('--compress_kv_factor', type=int, default=1)
+    parser.add_argument('--interpolation_scale_h', type=float, default=1.0)
+    parser.add_argument('--interpolation_scale_w', type=float, default=1.0)
+    parser.add_argument('--interpolation_scale_t', type=float, default=1.0)
+    parser.add_argument("--downsampler", type=str, default=None)
+    parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse")
+    parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse")
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument('--enable_stable_fp32', action='store_true')
+    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
+
+    # diffusion setting
+    parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.")
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.999)
+    parser.add_argument("--ema_start_step", type=int, default=0)
+    parser.add_argument("--noise_offset", type=float, default=0.02, help="The scale of noise offset.")
+    parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
+
+    # validation & logs
+    parser.add_argument("--num_sampling_steps", type=int, default=50)
+    parser.add_argument('--guidance_scale', type=float, default=2.5)
+    parser.add_argument("--enable_tracker", action="store_true")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store."))
+    parser.add_argument("--checkpointing_steps", type=int, default=500,
+                        help=(
+                            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+                            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+                            " training using `--resume_from_checkpoint`."
+                        ),
+                        )
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None,
+                        help=(
+                            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+                            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+                        ),
+                        )
+    parser.add_argument("--logging_dir", type=str, default="logs",
+                        help=(
+                            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+                            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+                        ),
+                        )
+    parser.add_argument("--report_to", type=str, default="tensorboard",
+                        help=(
+                            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+                            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+                        ),
+                        )
+    # optimizer & scheduler
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform.  If provided, overrides num_train_epochs.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")
+    parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]')
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")
+    parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
+    parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.")
+    parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW")
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.")
+    parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params")
+    parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.")
+    parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW")
+    parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--prodigy_beta3", type=float, default=None,
+                        help="coefficients for computing the Prodidy stepsize using running averages. If set to None, "
+                             "uses the value of square root of beta2. Ignored if optimizer is adamW",
+                        )
+    parser.add_argument("--lr_scheduler", type=str, default="constant",
+                        help=(
+                            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+                            ' "constant", "constant_with_warmup"]'
+                        ),
+                        )
+    parser.add_argument("--allow_tf32", action="store_true",
+                        help=(
+                            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+                            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+                        ),
+                        )
+    parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"],
+                        help=(
+                            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+                            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+                            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+                        ),
+                        )
+
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
+    parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
+
+    # inpaint
+    parser.add_argument("--i2v_ratio", type=float, default=0.5) # for inpainting mode
+    parser.add_argument("--transition_ratio", type=float, default=0.4) # for inpainting mode
+    parser.add_argument("--v2v_ratio", type=float, default=0.1) # for inpainting mode
+    parser.add_argument("--clear_video_ratio", type=float, default=0.0)
+    parser.add_argument("--default_text_ratio", type=float, default=0.1)
+    parser.add_argument("--pretrained_transformer_model_path", type=str, default=None)
+    parser.add_argument("--use_vae_preprocessed_mask", action="store_true")
+
+    args = parser.parse_args()
+    main(args)
\ No newline at end of file
diff --git a/opensora/train/train_t2v_diffusers.py b/opensora/train/train_t2v_diffusers.py
new file mode 100644
index 000000000..dffedb5e5
--- /dev/null
+++ b/opensora/train/train_t2v_diffusers.py
@@ -0,0 +1,957 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+A minimal training script for DiT using PyTorch DDP.
+"""
+import argparse
+import logging
+import math
+import os
+import shutil
+from pathlib import Path
+from typing import Optional
+import gc
+import numpy as np
+from einops import rearrange
+from tqdm import tqdm
+
+from opensora.adaptor.modules import replace_with_fp32_forwards
+
+try:
+    import torch_npu
+    from opensora.npu_config import npu_config
+    from opensora.acceleration.parallel_states import initialize_sequence_parallel_state, \
+        destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
+    from opensora.acceleration.communications import prepare_parallel_data, broadcast
+except:
+    torch_npu = None
+    npu_config = None
+    from opensora.utils.parallel_states import initialize_sequence_parallel_state, \
+        destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
+    from opensora.utils.communications import prepare_parallel_data, broadcast
+    pass
+import time
+from dataclasses import field, dataclass
+from torch.utils.data import DataLoader
+from copy import deepcopy
+import accelerate
+import torch
+from torch.nn import functional as F
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
+from packaging import version
+from tqdm.auto import tqdm
+
+import diffusers
+from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import compute_snr
+from diffusers.utils import check_min_version, is_wandb_available
+
+from opensora.utils.ema import EMAModel
+from opensora.dataset import getdataset
+from opensora.models import CausalVAEModelWrapper
+from opensora.models.text_encoder import get_text_enc, get_text_warpper
+from opensora.models.causalvideovae import ae_stride_config, ae_channel_config
+from opensora.models.causalvideovae import ae_norm, ae_denorm
+from opensora.models.diffusion import Diffusion_models, Diffusion_models_class
+from opensora.utils.dataset_utils import Collate, LengthGroupedSampler
+from opensora.sample.pipeline_opensora import OpenSoraPipeline
+
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.24.0")
+logger = get_logger(__name__)
+
+
+@torch.inference_mode()
+def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, ema=False):
+    positive_prompt = "(masterpiece), (best quality), (ultra-detailed), {}. emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, sharp focus, high budget, cinemascope, moody, epic, gorgeous"
+    negative_prompt = """nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, 
+                        """
+    validation_prompt = [
+        "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
+        "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."
+        ]
+    if 'mt5' in args.text_encoder_name:
+        validation_prompt_cn = [
+            "一只戴着墨镜在泳池当救生员的猫咪。",
+            "这是一个宁静的水下场景，一只海龟游过珊瑚礁。海龟带着绿褐色的龟壳，优雅地游向画面右侧，成为视频的焦点。背景中的珊瑚礁生机盎然，为海龟的旅程提供了生动多彩的背景。几条小鱼在海龟周围穿梭，为画面增添了动感和活力。"
+            ]
+        validation_prompt += validation_prompt_cn
+    logger.info(f"Running validation....\n")
+    model = accelerator.unwrap_model(model)
+    scheduler = DPMSolverMultistepScheduler()
+    opensora_pipeline = OpenSoraPipeline(vae=vae,
+                                         text_encoder=text_encoder,
+                                         tokenizer=tokenizer,
+                                         scheduler=scheduler,
+                                         transformer=model).to(device=accelerator.device)
+    videos = []
+    for prompt in validation_prompt:
+        logger.info('Processing the ({}) prompt'.format(prompt))
+        video = opensora_pipeline(
+                                positive_prompt.format(prompt),
+                                negative_prompt=negative_prompt, 
+                                num_frames=args.num_frames,
+                                height=args.max_height,
+                                width=args.max_width,
+                                num_inference_steps=args.num_sampling_steps,
+                                guidance_scale=args.guidance_scale,
+                                enable_temporal_attentions=True,
+                                num_images_per_prompt=1,
+                                mask_feature=True,
+                                max_sequence_length=args.model_max_length,
+                                ).images
+        videos.append(video[0])
+    # import ipdb;ipdb.set_trace()
+    gc.collect()
+    torch.cuda.empty_cache()
+    videos = torch.stack(videos).numpy()
+    videos = rearrange(videos, 'b t h w c -> b t c h w')
+    for tracker in accelerator.trackers:
+        if tracker.name == "tensorboard":
+            if videos.shape[1] == 1:
+                assert args.num_frames == 1
+                images = rearrange(videos, 'b 1 c h w -> (b 1) h w c')
+                np_images = np.stack([np.asarray(img) for img in images])
+                tracker.writer.add_images(f"{'ema_' if ema else ''}validation", np_images, global_step, dataformats="NHWC")
+            else:
+                np_videos = np.stack([np.asarray(vid) for vid in videos])
+                tracker.writer.add_video(f"{'ema_' if ema else ''}validation", np_videos, global_step, fps=24)
+        if tracker.name == "wandb":
+            import wandb
+            if videos.shape[1] == 1:
+                images = rearrange(videos, 'b 1 c h w -> (b 1) h w c')
+                logs = {
+                    f"{'ema_' if ema else ''}validation": [
+                        wandb.Image(image, caption=f"{i}: {prompt}")
+                        for i, (image, prompt) in enumerate(zip(images, validation_prompt))
+                    ]
+                }
+            else:
+                logs = {
+                    f"{'ema_' if ema else ''}validation": [
+                        wandb.Video(video, caption=f"{i}: {prompt}", fps=24)
+                        for i, (video, prompt) in enumerate(zip(videos, validation_prompt))
+                    ]
+                }
+            tracker.log(logs, step=global_step)
+
+    del opensora_pipeline
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+class ProgressInfo:
+    def __init__(self, global_step, train_loss=0.0):
+        self.global_step = global_step
+        self.train_loss = train_loss
+
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    # use LayerNorm, GeLu, SiLu always as fp32 mode
+    if args.enable_stable_fp32:
+        replace_with_fp32_forwards()
+    if torch_npu is not None and npu_config is not None:
+        npu_config.print_msg(args)
+        npu_config.seed_everything(args.seed)
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with=args.report_to,
+        project_config=accelerator_project_config,
+    )
+
+    if args.num_frames != 1 and args.use_image_num == 0:
+        initialize_sequence_parallel_state(args.sp_size)
+
+    if args.report_to == "wandb":
+        if not is_wandb_available():
+            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now. On GPU...
+    if torch_npu is None and npu_config is None and args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training we cast all non-trainable weigths to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # Create model:
+    kwargs = {}
+    ae = CausalVAEModelWrapper(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval()
+    if args.enable_tiling:
+        ae.vae.enable_tiling()
+        ae.vae.tile_overlap_factor = args.tile_overlap_factor
+
+    kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True}
+    text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval()
+
+    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
+    ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w)
+    assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})"
+    args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w
+    args.ae_stride = args.ae_stride_h
+    patch_size = args.model[-3:]
+    patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2])
+    args.patch_size = patch_size_h
+    args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w
+    assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})"
+    # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})."
+    assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})."
+    assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})."
+
+    args.stride_t = ae_stride_t * patch_size_t
+    args.stride = ae_stride_h * patch_size_h
+    latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w)
+    ae.latent_size = latent_size
+
+    if args.num_frames % 2 == 1:
+        args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1
+    else:
+        latent_size_t = args.num_frames // ae_stride_t
+    model = Diffusion_models[args.model](
+        in_channels=ae_channel_config[args.ae],
+        out_channels=ae_channel_config[args.ae],
+        # caption_channels=4096,
+        # cross_attention_dim=1152,
+        attention_bias=True,
+        sample_size=latent_size,
+        sample_size_t=latent_size_t,
+        num_vector_embeds=None,
+        activation_fn="gelu-approximate",
+        num_embeds_ada_norm=1000,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        double_self_attention=False,
+        upcast_attention=False,
+        # norm_type="ada_norm_single",
+        norm_elementwise_affine=False,
+        norm_eps=1e-6,
+        attention_type='default',
+        attention_mode=args.attention_mode,
+        interpolation_scale_h=args.interpolation_scale_h,
+        interpolation_scale_w=args.interpolation_scale_w,
+        interpolation_scale_t=args.interpolation_scale_t,
+        downsampler=args.downsampler,
+        # compress_kv_factor=args.compress_kv_factor,
+        use_rope=args.use_rope,
+        # model_max_length=args.model_max_length,
+        use_stable_fp32=args.enable_stable_fp32, 
+    )
+    model.gradient_checkpointing = args.gradient_checkpointing
+
+    # # use pretrained model?
+    if args.pretrained:
+        model_state_dict = model.state_dict()
+        if 'safetensors' in args.pretrained:  # pixart series
+            from safetensors.torch import load_file as safe_load
+            # import ipdb;ipdb.set_trace()
+            pretrained_checkpoint = safe_load(args.pretrained, device="cpu")
+            pretrained_keys = set(list(pretrained_checkpoint.keys()))
+            model_keys = set(list(model_state_dict.keys()))
+            common_keys = list(pretrained_keys & model_keys)
+            checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()}
+            # if checkpoint['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and checkpoint['pos_embed.proj.weight'].ndim == 4:
+            #     logger.info(f"Resize pos_embed, {checkpoint['pos_embed.proj.weight'].shape} -> {model.pos_embed.proj.weight.shape}")
+            #     repeat = model.pos_embed.proj.weight.shape[2]
+            #     checkpoint['pos_embed.proj.weight'] = checkpoint['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat)
+                # del checkpoint['proj_out.weight'], checkpoint['proj_out.bias']
+        else:  # latest stage training weight
+            checkpoint = torch.load(args.pretrained, map_location='cpu')
+            if 'model' in checkpoint:
+                checkpoint = checkpoint['model']
+        missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+        logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}')
+        logger.info(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!')
+
+    # Freeze vae and text encoders.
+    ae.vae.requires_grad_(False)
+    text_enc.requires_grad_(False)
+    # Set model as trainable.
+    model.train()
+
+    noise_scheduler = DDPMScheduler()
+    # Move unet, vae and text_encoder to device and cast to weight_dtype
+    # The VAE is in float32 to avoid NaN losses.
+    ae.vae.to(accelerator.device, dtype=torch.float32)
+    # ae.vae.to(accelerator.device, dtype=weight_dtype)
+    text_enc.to(accelerator.device, dtype=weight_dtype)
+
+    # Create EMA for the unet.
+    if args.use_ema:
+        ema_model = deepcopy(model)
+        ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step,
+                             model_cls=Diffusion_models_class[args.model], model_config=ema_model.config)
+
+    # `accelerate` 0.16.0 will have better support for customized saving
+    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
+        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+        def save_model_hook(models, weights, output_dir):
+            if accelerator.is_main_process:
+                if args.use_ema:
+                    ema_model.save_pretrained(os.path.join(output_dir, "model_ema"))
+
+                for i, model in enumerate(models):
+                    model.save_pretrained(os.path.join(output_dir, "model"))
+                    if weights:  # Don't pop if empty
+                        # make sure to pop weight so that corresponding model is not saved again
+                        weights.pop()
+
+        def load_model_hook(models, input_dir):
+            if args.use_ema:
+                load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model])
+                ema_model.load_state_dict(load_model.state_dict())
+                ema_model.to(accelerator.device)
+                del load_model
+
+            for i in range(len(models)):
+                # pop models so that they are not loaded again
+                model = models.pop()
+
+                # load diffusers style into model
+                load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model")
+                model.register_to_config(**load_model.config)
+
+                model.load_state_dict(load_model.state_dict())
+                del load_model
+
+        accelerator.register_save_state_pre_hook(save_model_hook)
+        accelerator.register_load_state_pre_hook(load_model_hook)
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32:
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    if args.scale_lr:
+        args.learning_rate = (
+                args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+        )
+
+    params_to_optimize = model.parameters()
+    # Optimizer creation
+    if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"):
+        logger.warning(
+            f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]."
+            "Defaulting to adamW"
+        )
+        args.optimizer = "adamw"
+
+    if args.use_8bit_adam and not args.optimizer.lower() == "adamw":
+        logger.warning(
+            f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was "
+            f"set to {args.optimizer.lower()}"
+        )
+
+    if args.optimizer.lower() == "adamw":
+        if args.use_8bit_adam:
+            try:
+                import bitsandbytes as bnb
+            except ImportError:
+                raise ImportError(
+                    "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+                )
+
+            optimizer_class = bnb.optim.AdamW8bit
+        else:
+            optimizer_class = torch.optim.AdamW
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+        )
+
+    if args.optimizer.lower() == "prodigy":
+        try:
+            import prodigyopt
+        except ImportError:
+            raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+
+        optimizer_class = prodigyopt.Prodigy
+
+        if args.learning_rate <= 0.1:
+            logger.warning(
+                "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
+            )
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            beta3=args.prodigy_beta3,
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            decouple=args.prodigy_decouple,
+            use_bias_correction=args.prodigy_use_bias_correction,
+            safeguard_warmup=args.prodigy_safeguard_warmup,
+        )
+    logger.info(f"optimizer: {optimizer}")
+    
+    # Setup data:
+    train_dataset = getdataset(args)
+    sampler = LengthGroupedSampler(
+                args.train_batch_size,
+                world_size=accelerator.num_processes,
+                lengths=train_dataset.lengths, 
+                group_frame=args.group_frame, 
+                group_resolution=args.group_resolution, 
+            ) if (args.group_frame or args.group_resolution) else None
+    train_dataloader = DataLoader(
+        train_dataset,
+        shuffle=sampler is None,
+        # pin_memory=True,
+        collate_fn=Collate(args),
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        sampler=sampler if (args.group_frame or args.group_resolution) else None, 
+        drop_last=True, 
+        # prefetch_factor=4
+    )
+    logger.info(f'after train_dataloader')
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+
+    # Prepare everything with our `accelerator`.
+    # model.requires_grad_(False)
+    # model.pos_embed.requires_grad_(True)
+    logger.info(f'before accelerator.prepare')
+    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, lr_scheduler
+    )
+    logger.info(f'after accelerator.prepare')
+    if args.use_ema:
+        ema_model.to(accelerator.device)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args))
+
+    # Train!
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size
+    logger.info("***** Running training *****")
+    logger.info(f"  Model = {model}")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Total training parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B")
+    global_step = 0
+    first_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+
+            if npu_config is not None:
+                train_dataset.n_used_elements = global_step * args.train_batch_size
+
+    else:
+        initial_global_step = 0
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=initial_global_step,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=not accelerator.is_local_main_process,
+    )
+    progress_info = ProgressInfo(global_step, train_loss=0.0)
+
+    def sync_gradients_info(loss):
+        # Checks if the accelerator has performed an optimization step behind the scenes
+        if args.use_ema:
+            ema_model.step(model.parameters())
+        progress_bar.update(1)
+        progress_info.global_step += 1
+        end_time = time.time()
+        one_step_duration = end_time - start_time
+        accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step)
+        if torch_npu is not None and npu_config is not None:
+            npu_config.print_msg(f"Step: [{progress_info.global_step}], local_loss={loss.detach().item()}, "
+                                 f"train_loss={progress_info.train_loss}, time_cost={one_step_duration}",
+                                 rank=0)
+        progress_info.train_loss = 0.0
+
+        # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
+        if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
+            if progress_info.global_step % args.checkpointing_steps == 0:
+                # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+                if accelerator.is_main_process and args.checkpoints_total_limit is not None:
+                    checkpoints = os.listdir(args.output_dir)
+                    checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+                    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+                    # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+                    if len(checkpoints) >= args.checkpoints_total_limit:
+                        num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+                        removing_checkpoints = checkpoints[0:num_to_remove]
+
+                        logger.info(
+                            f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+                        )
+                        logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+                        for removing_checkpoint in removing_checkpoints:
+                            removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+                            shutil.rmtree(removing_checkpoint)
+
+                save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}")
+                accelerator.save_state(save_path)
+                logger.info(f"Saved state to {save_path}")
+
+        logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+        progress_bar.set_postfix(**logs)
+
+    def run(model_input, model_kwargs, prof):
+        global start_time
+        start_time = time.time()
+
+        noise = torch.randn_like(model_input)
+        if args.noise_offset:
+            # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+            noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1),
+                                                     device=model_input.device)
+
+        bsz = model_input.shape[0]
+        current_step_frame = model_input.shape[2]
+        # Sample a random timestep for each image without bias.
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device)
+        if current_step_frame != 1 and get_sequence_parallel_state():  # image do not need sp
+            broadcast(timesteps)
+
+        # Add noise to the model input according to the noise magnitude at each timestep
+        # (this is the forward diffusion process)
+
+        noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)
+
+        model_pred = model(
+            noisy_model_input,
+            timesteps,
+            **model_kwargs
+        )[0]
+        # Get the target for loss depending on the prediction type
+        if args.prediction_type is not None:
+            # set prediction_type of scheduler if defined
+            noise_scheduler.register_to_config(prediction_type=args.prediction_type)
+
+        if noise_scheduler.config.prediction_type == "epsilon":
+            target = noise
+        elif noise_scheduler.config.prediction_type == "v_prediction":
+            target = noise_scheduler.get_velocity(model_input, noise, timesteps)
+        elif noise_scheduler.config.prediction_type == "sample":
+            # We set the target to latents here, but the model_pred will return the noise sample prediction.
+            target = model_input
+            # We will have to subtract the noise residual from the prediction to get the target sample.
+            model_pred = model_pred - noise
+        else:
+            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+        mask = model_kwargs.get('attention_mask', None)
+        if torch.all(mask.bool()):
+            mask = None
+        if get_sequence_parallel_state():
+            assert mask is None
+        b, c, _, _, _ = model_pred.shape
+        if mask is not None:
+            mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float()  # b t h w -> b c t h w
+            mask = mask.reshape(b, -1)
+        if args.snr_gamma is None:
+            # model_pred: b c t h w, attention_mask: b t h w
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+            loss = loss.reshape(b, -1)
+            if mask is not None:
+                loss = (loss * mask).sum() / mask.sum()  # mean loss on unpad patches
+            else:
+                loss = loss.mean()
+        else:
+            # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
+            # Since we predict the noise instead of x_0, the original formulation is slightly changed.
+            # This is discussed in Section 4.2 of the same paper.
+            snr = compute_snr(noise_scheduler, timesteps)
+            mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(
+                dim=1
+            )[0]
+            if noise_scheduler.config.prediction_type == "epsilon":
+                mse_loss_weights = mse_loss_weights / snr
+            elif noise_scheduler.config.prediction_type == "v_prediction":
+                mse_loss_weights = mse_loss_weights / (snr + 1)
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+            loss = loss.reshape(b, -1)
+            mse_loss_weights = mse_loss_weights.reshape(b, 1)
+            if mask is not None:
+                loss = (loss * mask * mse_loss_weights).sum() / mask.sum()  # mean loss on unpad patches
+            else:
+                loss = (loss * mse_loss_weights).mean()
+
+        # Gather the losses across all processes for logging (if we use distributed training).
+        avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
+        progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps
+
+        # Backpropagate
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+            params_to_clip = model.parameters()
+            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad()
+
+        if accelerator.sync_gradients:
+            sync_gradients_info(loss)
+
+        if accelerator.is_main_process:
+
+            if progress_info.global_step % args.checkpointing_steps == 0:
+
+                if args.enable_tracker:
+                    log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator,
+                                   weight_dtype, progress_info.global_step)
+
+                    if args.use_ema and npu_config is None:
+                        # Store the UNet parameters temporarily and load the EMA parameters to perform inference.
+                        ema_model.store(model.parameters())
+                        ema_model.copy_to(model.parameters())
+                        log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator,
+                                    weight_dtype, progress_info.global_step, ema=True)
+                        # Switch back to the original UNet parameters.
+                        ema_model.restore(model.parameters())
+
+        if prof is not None:
+            prof.step()
+
+
+        return loss
+
+    def train_one_step(step_, data_item_, prof_=None):
+        train_loss = 0.0
+        x, attn_mask, input_ids, cond_mask = data_item_
+        if args.group_frame or args.group_resolution:
+            if not args.group_frame:
+                each_latent_frame = torch.any(attn_mask.flatten(-2), dim=-1).int().sum(-1).tolist()
+                # logger.info(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask '
+                #             f'each_latent_frame: {each_latent_frame}')
+                print(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask '
+                            f'each_latent_frame: {each_latent_frame}')
+        assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))'
+        x = x.to(accelerator.device, dtype=ae.vae.dtype)  # B C T+num_images H W, 16 + 4
+
+        attn_mask = attn_mask.to(accelerator.device)  # B T+num_images H W
+        input_ids = input_ids.to(accelerator.device)  # B 1+num_images L
+        cond_mask = cond_mask.to(accelerator.device)  # B 1+num_images L
+        # if accelerator.process_index == 0:
+        #     logger.info(f'rank: {accelerator.process_index}, x: {x.shape}, attn_mask: {attn_mask.shape}')
+
+        with torch.no_grad():
+            B, N, L = input_ids.shape  # B 1+num_images L
+            # use batch inference
+            input_ids_ = input_ids.reshape(-1, L)
+            cond_mask_ = cond_mask.reshape(-1, L)
+            cond = text_enc(input_ids_, cond_mask_)  # B 1+num_images L D
+            cond = cond.reshape(B, N, L, -1)
+            # Map input images to latent space + normalize latents
+            x = ae.encode(x)  # B C T H W
+
+            # def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None:
+            #     from examples.rec_video import array_to_video
+            #     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
+            # videos = ae.decode(x)[0]
+            # videos = videos.transpose(0, 1)
+            # custom_to_video(videos.to(torch.float32), fps=24, output_file='tmp.mp4')
+            # import sys;sys.exit()
+        current_step_frame = x.shape[2]
+        current_step_sp_state = get_sequence_parallel_state()
+        if args.sp_size != 1:  # enable sp
+            if current_step_frame == 1:  # but image do not need sp
+                set_sequence_parallel_state(False)
+            else:
+                set_sequence_parallel_state(True)
+        if get_sequence_parallel_state():
+            x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask,
+                                                                                 args.use_image_num)
+            for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size):
+                with accelerator.accumulate(model):
+                    st_idx = iter * args.train_sp_batch_size
+                    ed_idx = (iter + 1) * args.train_sp_batch_size
+                    model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx],
+                                        attention_mask=attn_mask[st_idx: ed_idx],
+                                        encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num)
+                    run(x[st_idx: ed_idx], model_kwargs, prof_)
+
+        else:
+            with accelerator.accumulate(model):
+                assert not torch.any(torch.isnan(x)), 'after vae'
+                x = x.to(weight_dtype)
+                model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask,
+                                    encoder_attention_mask=cond_mask, use_image_num=args.use_image_num)
+                run(x, model_kwargs, prof_)
+
+        set_sequence_parallel_state(current_step_sp_state)  # in case the next step use sp, which need broadcast(timesteps)
+
+        if progress_info.global_step >= args.max_train_steps:
+            return True
+
+        return False
+
+    def train_all_epoch(prof_=None):
+        for epoch in range(first_epoch, args.num_train_epochs):
+            progress_info.train_loss = 0.0
+            if progress_info.global_step >= args.max_train_steps:
+                return True
+
+            for step, data_item in enumerate(train_dataloader):
+                if train_one_step(step, data_item, prof_):
+                    break
+
+                if step >= 2 and torch_npu is not None and npu_config is not None:
+                    npu_config.free_mm()
+
+    if npu_config is not None and npu_config.on_npu and npu_config.profiling:
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization
+        )
+        profile_output_path = f"/home/image_data/npu_profiling_t2v/{os.getenv('PROJECT_NAME', 'local')}"
+        os.makedirs(profile_output_path, exist_ok=True)
+
+        with torch_npu.profiler.profile(
+                activities=[torch_npu.profiler.ProfilerActivity.NPU, torch_npu.profiler.ProfilerActivity.CPU],
+                with_stack=True,
+                record_shapes=True,
+                profile_memory=True,
+                experimental_config=experimental_config,
+                schedule=torch_npu.profiler.schedule(wait=npu_config.profiling_step, warmup=0, active=1, repeat=1,
+                                                     skip_first=0),
+                on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f"{profile_output_path}/")
+        ) as prof:
+            train_all_epoch(prof)
+    else:
+        train_all_epoch()
+    accelerator.wait_for_everyone()
+    accelerator.end_training()
+    if get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # dataset & dataloader
+    parser.add_argument("--dataset", type=str, required=True)
+    parser.add_argument("--data", type=str, required='')
+    parser.add_argument("--sample_rate", type=int, default=1)
+    parser.add_argument("--train_fps", type=int, default=24)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--num_frames", type=int, default=65)
+    parser.add_argument("--max_height", type=int, default=320)
+    parser.add_argument("--max_width", type=int, default=240)
+    parser.add_argument("--use_img_from_vid", action="store_true")
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--model_max_length", type=int, default=512)
+    parser.add_argument('--cfg', type=float, default=0.1)
+    parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.")
+    parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.")
+    parser.add_argument("--group_frame", action="store_true")
+    parser.add_argument("--group_resolution", action="store_true")
+
+    # text encoder & vae & diffusion model
+    parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122")
+    parser.add_argument('--enable_8bit_t5', action='store_true')
+    parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument("--compress_kv", action="store_true")
+    parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers")
+    parser.add_argument('--use_rope', action='store_true')
+    parser.add_argument('--compress_kv_factor', type=int, default=1)
+    parser.add_argument('--interpolation_scale_h', type=float, default=1.0)
+    parser.add_argument('--interpolation_scale_w', type=float, default=1.0)
+    parser.add_argument('--interpolation_scale_t', type=float, default=1.0)
+    parser.add_argument("--downsampler", type=str, default=None)
+    parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse")
+    parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse")
+    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--pretrained", type=str, default=None)
+    parser.add_argument('--enable_stable_fp32', action='store_true')
+    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
+
+    # diffusion setting
+    parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.")
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.999)
+    parser.add_argument("--ema_start_step", type=int, default=0)
+    parser.add_argument("--noise_offset", type=float, default=0.02, help="The scale of noise offset.")
+    parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
+
+    # validation & logs
+    parser.add_argument("--num_sampling_steps", type=int, default=20)
+    parser.add_argument('--guidance_scale', type=float, default=4.5)
+    parser.add_argument("--enable_tracker", action="store_true")
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store."))
+    parser.add_argument("--checkpointing_steps", type=int, default=500,
+                        help=(
+                            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+                            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+                            " training using `--resume_from_checkpoint`."
+                        ),
+                        )
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None,
+                        help=(
+                            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+                            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+                        ),
+                        )
+    parser.add_argument("--logging_dir", type=str, default="logs",
+                        help=(
+                            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+                            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+                        ),
+                        )
+    parser.add_argument("--report_to", type=str, default="tensorboard",
+                        help=(
+                            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+                            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+                        ),
+                        )
+    # optimizer & scheduler
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform.  If provided, overrides num_train_epochs.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")
+    parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]')
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")
+    parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
+    parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.")
+    parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW")
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.")
+    parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params")
+    parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.")
+    parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW")
+    parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--prodigy_beta3", type=float, default=None,
+                        help="coefficients for computing the Prodidy stepsize using running averages. If set to None, "
+                             "uses the value of square root of beta2. Ignored if optimizer is adamW",
+                        )
+    parser.add_argument("--lr_scheduler", type=str, default="constant",
+                        help=(
+                            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+                            ' "constant", "constant_with_warmup"]'
+                        ),
+                        )
+    parser.add_argument("--allow_tf32", action="store_true",
+                        help=(
+                            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+                            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+                        ),
+                        )
+    parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"],
+                        help=(
+                            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+                            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+                            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+                        ),
+                        )
+
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
+    parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
+
+    args = parser.parse_args()
+    main(args)
diff --git a/opensora/utils/communications.py b/opensora/utils/communications.py
new file mode 100644
index 000000000..c7a79e3a6
--- /dev/null
+++ b/opensora/utils/communications.py
@@ -0,0 +1,122 @@
+import torch
+import torch.distributed as dist
+from einops import rearrange
+from opensora.utils.parallel_states import nccl_info
+
+def broadcast(input_: torch.Tensor):
+    sp_size = nccl_info.world_size
+    src = nccl_info.rank // sp_size * sp_size
+    dist.broadcast(input_, src=src, group=nccl_info.group)
+
+_COUNT = 0
+def _all_to_all(
+    input_: torch.Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+):
+    group = nccl_info.group
+    sp_size = nccl_info.world_size
+    input_list = [t.contiguous() for t in torch.tensor_split(input_, sp_size, scatter_dim)]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+def _single_all_to_all(
+    input_: torch.Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    enable_HCCL=False,
+):
+
+    sp_size = nccl_info.world_size
+    inp_shape = list(input_.shape)
+    inp_shape[scatter_dim] = inp_shape[scatter_dim] // sp_size
+    if scatter_dim < 1:
+        input_t = input_.reshape(
+            [sp_size, inp_shape[scatter_dim]] + \
+            inp_shape[scatter_dim + 1:]
+        )
+    else:
+        # transpose groups of heads with the seq-len parallel dimension, so that we can scatter them!
+        input_t = input_.reshape(
+            [-1, sp_size, inp_shape[scatter_dim]] + \
+            inp_shape[scatter_dim + 1:]
+        ).transpose(0, 1).contiguous()
+
+    output = torch.empty_like(input_t)
+    dist.all_to_all_single(output, input_t, group=nccl_info.group)
+    # if scattering the seq-dim, transpose the heads back to the original dimension
+    if scatter_dim < 1:
+        output = output.transpose(0, 1).contiguous()
+
+    return output.reshape(
+        inp_shape[: gather_dim] + [inp_shape[gather_dim] * sp_size, ] + inp_shape[gather_dim + 1:])
+
+
+class _AllToAll(torch.autograd.Function):
+    """All-to-all communication.
+
+    Args:
+        input_: input matrix
+        process_group: communication group
+        scatter_dim: scatter dimension
+        gather_dim: gather dimension
+    """
+
+    @staticmethod
+    def forward(ctx, input_, scatter_dim, gather_dim, all_to_all_func):
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.all_to_all = all_to_all_func
+        output = ctx.all_to_all(input_, scatter_dim, gather_dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output = ctx.all_to_all(
+            grad_output,
+            ctx.gather_dim,
+            ctx.scatter_dim,
+        )
+        return (
+            grad_output,
+            None,
+            None,
+            None,
+        )
+
+def all_to_all_SBH(
+    input_: torch.Tensor,
+    scatter_dim: int = 1,
+    gather_dim: int = 0,
+):
+    return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all)
+
+def all_to_all_BSND(
+    input_: torch.Tensor,
+    scatter_dim: int = 2,
+    gather_dim: int = 1,
+):
+    return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all)
+
+
+def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num):
+    def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
+        hidden_states = _single_all_to_all(hidden_states, scatter_dim=2, gather_dim=0, enable_HCCL=True)
+        encoder_hidden_states = _single_all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0, enable_HCCL=True)
+        attention_mask = _single_all_to_all(attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True)
+        encoder_attention_mask = _single_all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0, enable_HCCL=True)
+        return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
+
+    sp_size = nccl_info.world_size
+    frame = hidden_states.shape[2]
+    assert frame % sp_size == 0, "frame should be a multiple of sp_size"
+
+    encoder_hidden_states = rearrange(encoder_hidden_states, 'b 1 (n x) h -> b n x h',
+                                     n=sp_size, x=encoder_hidden_states.shape[2]//sp_size).contiguous()
+    hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask = all_to_all(hidden_states,
+                                                                                            encoder_hidden_states,
+                                                                                            attention_mask.repeat(1, sp_size, 1, 1),
+                                                                                            encoder_attention_mask.repeat(1, sp_size, 1))
+
+    return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num
\ No newline at end of file
diff --git a/opensora/utils/dataset_utils.py b/opensora/utils/dataset_utils.py
new file mode 100644
index 000000000..449050afa
--- /dev/null
+++ b/opensora/utils/dataset_utils.py
@@ -0,0 +1,329 @@
+import math
+from einops import rearrange
+import decord
+from torch.nn import functional as F
+import torch
+from typing import Optional
+import torch.utils
+import torch.utils.data
+import torch
+from torch.utils.data import Sampler
+from typing import List
+from collections import Counter
+import random
+
+
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+class DecordInit(object):
+    """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader."""
+
+    def __init__(self, num_threads=1):
+        self.num_threads = num_threads
+        self.ctx = decord.cpu(0)
+
+    def __call__(self, filename):
+        """Perform the Decord initialization.
+        Args:
+            results (dict): The resulting dict to be modified and passed
+                to the next transform in pipeline.
+        """
+        reader = decord.VideoReader(filename,
+                                    ctx=self.ctx,
+                                    num_threads=self.num_threads)
+        return reader
+
+    def __repr__(self):
+        repr_str = (f'{self.__class__.__name__}('
+                    f'sr={self.sr},'
+                    f'num_threads={self.num_threads})')
+        return repr_str
+
+def pad_to_multiple(number, ds_stride):
+    remainder = number % ds_stride
+    if remainder == 0:
+        return number
+    else:
+        padding = ds_stride - remainder
+        return number + padding
+
+class Collate:
+    def __init__(self, args):
+        self.batch_size = args.train_batch_size
+        self.group_frame = args.group_frame
+        self.group_resolution = args.group_resolution
+
+        self.max_height = args.max_height
+        self.max_width = args.max_width
+        self.ae_stride = args.ae_stride
+
+        self.ae_stride_t = args.ae_stride_t
+        self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride)
+
+        self.patch_size = args.patch_size
+        self.patch_size_t = args.patch_size_t
+
+        self.num_frames = args.num_frames
+        self.use_image_num = args.use_image_num
+        self.max_thw = (self.num_frames, self.max_height, self.max_width)
+
+    def package(self, batch):
+        batch_tubes = [i['pixel_values'] for i in batch]  # b [c t h w]
+        input_ids = [i['input_ids'] for i in batch]  # b [1 l]
+        cond_mask = [i['cond_mask'] for i in batch]  # b [1 l]
+        return batch_tubes, input_ids, cond_mask
+
+    def __call__(self, batch):
+        batch_tubes, input_ids, cond_mask = self.package(batch)
+
+        ds_stride = self.ae_stride * self.patch_size
+        t_ds_stride = self.ae_stride_t * self.patch_size_t
+        
+        pad_batch_tubes, attention_mask, input_ids, cond_mask = self.process(batch_tubes, input_ids, cond_mask, t_ds_stride, ds_stride, self.max_thw, self.ae_stride_thw)
+        assert not torch.any(torch.isnan(pad_batch_tubes)), 'after pad_batch_tubes'
+        return pad_batch_tubes, attention_mask, input_ids, cond_mask
+
+
+    def process(self, batch_tubes, input_ids, cond_mask, t_ds_stride, ds_stride, max_thw, ae_stride_thw):
+        # pad to max multiple of ds_stride
+        batch_input_size = [i.shape for i in batch_tubes]  # [(c t h w), (c t h w)]
+        assert len(batch_input_size) == self.batch_size
+        if self.group_frame or self.group_resolution or self.batch_size == 1:  #
+            len_each_batch = batch_input_size
+            idx_length_dict = dict([*zip(list(range(self.batch_size)), len_each_batch)])
+            count_dict = Counter(len_each_batch)
+            if len(count_dict) != 1:
+                sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
+                # import ipdb;ipdb.set_trace()
+                # print(batch, idx_length_dict, count_dict, sorted_by_value)
+                pick_length = sorted_by_value[-1][0]  # the highest frequency
+                candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length]
+                random_select_batch = [random.choice(candidate_batch) for _ in range(len(len_each_batch) - len(candidate_batch))]
+                print(batch_input_size, idx_length_dict, count_dict, sorted_by_value, pick_length, candidate_batch, random_select_batch)
+                pick_idx = candidate_batch + random_select_batch
+
+                batch_tubes = [batch_tubes[i] for i in pick_idx]
+                batch_input_size = [i.shape for i in batch_tubes]  # [(c t h w), (c t h w)]
+                input_ids = [input_ids[i] for i in pick_idx]  # b [1, l]
+                cond_mask = [cond_mask[i] for i in pick_idx]  # b [1, l]
+
+            for i in range(1, self.batch_size):
+                assert batch_input_size[0] == batch_input_size[i]
+            max_t = max([i[1] for i in batch_input_size])
+            max_h = max([i[2] for i in batch_input_size])
+            max_w = max([i[3] for i in batch_input_size])
+        else:
+            max_t, max_h, max_w = max_thw
+        pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1+self.ae_stride_t, t_ds_stride), \
+                                          pad_to_multiple(max_h, ds_stride), \
+                                          pad_to_multiple(max_w, ds_stride)
+        pad_max_t = pad_max_t + 1 - self.ae_stride_t
+        each_pad_t_h_w = [
+            [
+                pad_max_t - i.shape[1],
+                pad_max_h - i.shape[2],
+                pad_max_w - i.shape[3]
+                ] for i in batch_tubes
+                ]
+        pad_batch_tubes = [
+            F.pad(im, (0, pad_w, 0, pad_h, 0, pad_t), value=0) 
+            for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)
+            ]
+        pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0)
+
+
+        max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
+        max_latent_size = [
+            ((max_tube_size[0]-1) // ae_stride_thw[0] + 1),
+            max_tube_size[1] // ae_stride_thw[1],
+            max_tube_size[2] // ae_stride_thw[2]
+            ]
+        valid_latent_size = [
+            [
+                int(math.ceil((i[1]-1) / ae_stride_thw[0])) + 1,
+                int(math.ceil(i[2] / ae_stride_thw[1])),
+                int(math.ceil(i[3] / ae_stride_thw[2]))
+                ] for i in batch_input_size]
+        attention_mask = [
+            F.pad(torch.ones(i, dtype=pad_batch_tubes.dtype), (0, max_latent_size[2] - i[2], 
+                                                               0, max_latent_size[1] - i[1],
+                                                               0, max_latent_size[0] - i[0]), value=0) for i in valid_latent_size]
+        attention_mask = torch.stack(attention_mask)  # b t h w
+        if self.batch_size == 1 or self.group_frame or self.group_resolution:
+            assert torch.all(attention_mask.bool())
+            
+        input_ids = torch.stack(input_ids)  # b 1 l
+        cond_mask = torch.stack(cond_mask)  # b 1 l
+
+        return pad_batch_tubes, attention_mask, input_ids, cond_mask
+
+ 
+def split_to_even_chunks(indices, lengths, num_chunks, batch_size):
+    """
+    Split a list of indices into `chunks` chunks of roughly equal lengths.
+    """
+
+    if len(indices) % num_chunks != 0:
+        chunks = [indices[i::num_chunks] for i in range(num_chunks)]
+    else:
+        num_indices_per_chunk = len(indices) // num_chunks
+
+        chunks = [[] for _ in range(num_chunks)]
+        chunks_lengths = [0 for _ in range(num_chunks)]
+        for index in indices:
+            shortest_chunk = chunks_lengths.index(min(chunks_lengths))
+            chunks[shortest_chunk].append(index)
+            chunks_lengths[shortest_chunk] += lengths[index]
+            if len(chunks[shortest_chunk]) == num_indices_per_chunk:
+                chunks_lengths[shortest_chunk] = float("inf")
+    # return chunks
+
+    pad_chunks = []
+    for idx, chunk in enumerate(chunks):
+        if batch_size != len(chunk):
+            assert batch_size > len(chunk)
+            if len(chunk) != 0:
+                chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))]
+            else:
+                chunk = random.choice(pad_chunks)
+                print(chunks[idx], '->', chunk)
+        pad_chunks.append(chunk)
+    return pad_chunks
+
+def group_frame_fun(indices, lengths):
+    # sort by num_frames
+    indices.sort(key=lambda i: lengths[i], reverse=True)
+    return indices
+
+def group_resolution_fun(indices):
+    raise NotImplementedError
+    return indices
+
+def group_frame_and_resolution_fun(indices):
+    raise NotImplementedError
+    return indices
+
+def last_group_frame_fun(shuffled_megabatches, lengths):
+    re_shuffled_megabatches = []
+    # print('shuffled_megabatches', len(shuffled_megabatches))
+    for i_megabatch, megabatch in enumerate(shuffled_megabatches):
+        re_megabatch = []
+        for i_batch, batch in enumerate(megabatch):
+            assert len(batch) != 0
+            len_each_batch = [lengths[i] for i in batch]
+            idx_length_dict = dict([*zip(batch, len_each_batch)])
+            count_dict = Counter(len_each_batch)
+            if len(count_dict) != 1:
+                sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
+                # print(batch, idx_length_dict, count_dict, sorted_by_value)
+                pick_length = sorted_by_value[-1][0]  # the highest frequency
+                candidate_batch = [idx for idx, length in idx_length_dict.items() if length == pick_length]
+                random_select_batch = [random.choice(candidate_batch) for i in range(len(len_each_batch) - len(candidate_batch))]
+                # print(batch, idx_length_dict, count_dict, sorted_by_value, pick_length, candidate_batch, random_select_batch)
+                batch = candidate_batch + random_select_batch
+                # print(batch)
+            re_megabatch.append(batch)
+        re_shuffled_megabatches.append(re_megabatch)
+    
+    
+    # for megabatch, re_megabatch in zip(shuffled_megabatches, re_shuffled_megabatches):
+    #     for batch, re_batch in zip(megabatch, re_megabatch):
+    #         for i, re_i in zip(batch, re_batch):
+    #             if i != re_i:
+    #                 print(i, re_i)
+    return re_shuffled_megabatches
+                
+    
+def last_group_resolution_fun(indices):
+    raise NotImplementedError
+    return indices
+
+def last_group_frame_and_resolution_fun(indices):
+    raise NotImplementedError
+    return indices
+
+def get_length_grouped_indices(lengths, batch_size, world_size, generator=None, group_frame=False, group_resolution=False, seed=42):
+    # We need to use torch for the random part as a distributed sampler will set the random seed for torch.
+    if generator is None:
+        generator = torch.Generator().manual_seed(seed)  # every rank will generate a fixed order but random index
+    # print('lengths', lengths)
+    
+    indices = torch.randperm(len(lengths), generator=generator).tolist()
+    # print('indices', len(indices))
+
+    if group_frame and not group_resolution:
+        indices = group_frame_fun(indices, lengths)
+    elif not group_frame and group_resolution:
+        indices = group_resolution_fun(indices)
+    elif group_frame and group_resolution:
+        indices = group_frame_and_resolution_fun(indices)
+    # print('sort indices', len(indices))
+    # print('sort indices', indices)
+    # print('sort lengths', [lengths[i] for i in indices])
+    
+    megabatch_size = world_size * batch_size
+    megabatches = [indices[i: i + megabatch_size] for i in range(0, len(lengths), megabatch_size)]
+    # print('megabatches', len(megabatches))
+    # print('\nmegabatches', megabatches)
+    megabatches = [sorted(megabatch, key=lambda i: lengths[i], reverse=True) for megabatch in megabatches]
+    # print('sort megabatches', len(megabatches))
+    # megabatches_len = [[lengths[i] for i in megabatch] for megabatch in megabatches]
+    # print('\nsorted megabatches', megabatches)
+    # print('\nsorted megabatches_len', megabatches_len)
+    megabatches = [split_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches]
+    # print('nsplit_to_even_chunks megabatches', len(megabatches))
+    # print('\nsplit_to_even_chunks megabatches', megabatches)
+    # print('\nsplit_to_even_chunks len', [lengths[i] for megabatch in megabatches for batch in megabatch for i in batch])
+    # return [i for megabatch in megabatches for batch in megabatch for i in batch]
+
+    indices = torch.randperm(len(megabatches), generator=generator).tolist()
+    shuffled_megabatches = [megabatches[i] for i in indices]
+    # print('shuffled_megabatches', len(shuffled_megabatches))
+    if group_frame and not group_resolution:
+        shuffled_megabatches = last_group_frame_fun(shuffled_megabatches, lengths)
+    elif not group_frame and group_resolution:
+        shuffled_megabatches = last_group_resolution_fun(shuffled_megabatches, indices)
+    elif group_frame and group_resolution:
+        shuffled_megabatches = last_group_frame_and_resolution_fun(shuffled_megabatches, indices)
+    # print('\nshuffled_megabatches', shuffled_megabatches)
+    # import ipdb;ipdb.set_trace()
+    # print('\nshuffled_megabatches len', [lengths[i] for megabatch in shuffled_megabatches for batch in megabatch for i in batch])
+
+    return [i for megabatch in shuffled_megabatches for batch in megabatch for i in batch]
+
+
+class LengthGroupedSampler(Sampler):
+    r"""
+    Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while
+    keeping a bit of randomness.
+    """
+
+    def __init__(
+        self,
+        batch_size: int,
+        world_size: int,
+        lengths: Optional[List[int]] = None, 
+        group_frame=False, 
+        group_resolution=False, 
+        generator=None,
+    ):
+        if lengths is None:
+            raise ValueError("Lengths must be provided.")
+
+        self.batch_size = batch_size
+        self.world_size = world_size
+        self.lengths = lengths
+        self.group_frame = group_frame
+        self.group_resolution = group_resolution
+        self.generator = generator
+
+    def __len__(self):
+        return len(self.lengths)
+
+    def __iter__(self):
+        indices = get_length_grouped_indices(self.lengths, self.batch_size, self.world_size, group_frame=self.group_frame, 
+                                             group_resolution=self.group_resolution, generator=self.generator)
+        return iter(indices)
diff --git a/opensora/utils/downloader.py b/opensora/utils/downloader.py
new file mode 100644
index 000000000..f2ac4017b
--- /dev/null
+++ b/opensora/utils/downloader.py
@@ -0,0 +1,18 @@
+import gdown
+import os
+
+opensora_cache_home = os.path.expanduser(
+    os.getenv("OPENSORA_HOME", os.path.join("~/.cache", "opensora"))
+)
+
+
+def gdown_download(id, fname, cache_dir=None):
+    cache_dir = opensora_cache_home if not cache_dir else cache_dir
+
+    os.makedirs(cache_dir, exist_ok=True)
+    destination = os.path.join(cache_dir, fname)
+    if os.path.exists(destination):
+        return destination
+
+    gdown.download(id=id, output=destination, quiet=False)
+    return destination
diff --git a/opensora/utils/ema.py b/opensora/utils/ema.py
new file mode 100644
index 000000000..3ac045381
--- /dev/null
+++ b/opensora/utils/ema.py
@@ -0,0 +1,328 @@
+import contextlib
+import copy
+import random
+from typing import Any, Dict, Iterable, List, Optional, Union
+
+from diffusers.utils import (
+    deprecate,
+    is_torchvision_available,
+    is_transformers_available,
+)
+
+if is_transformers_available():
+    import transformers
+
+if is_torchvision_available():
+    from torchvision import transforms
+
+import numpy as np
+import torch
+
+
+# Adapted from diffusers-style ema https://github.com/huggingface/diffusers/blob/main/src/diffusers/training_utils.py#L263
+class EMAModel:
+    """
+    Exponential Moving Average of models weights
+    """
+
+    def __init__(
+        self,
+        parameters: Iterable[torch.nn.Parameter],
+        decay: float = 0.9999,
+        min_decay: float = 0.0,
+        update_after_step: int = 0,
+        use_ema_warmup: bool = False,
+        inv_gamma: Union[float, int] = 1.0,
+        power: Union[float, int] = 2 / 3,
+        model_cls: Optional[Any] = None,
+        model_config: Dict[str, Any] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            parameters (Iterable[torch.nn.Parameter]): The parameters to track.
+            decay (float): The decay factor for the exponential moving average.
+            min_decay (float): The minimum decay factor for the exponential moving average.
+            update_after_step (int): The number of steps to wait before starting to update the EMA weights.
+            use_ema_warmup (bool): Whether to use EMA warmup.
+            inv_gamma (float):
+                Inverse multiplicative factor of EMA warmup. Default: 1. Only used if `use_ema_warmup` is True.
+            power (float): Exponential factor of EMA warmup. Default: 2/3. Only used if `use_ema_warmup` is True.
+            device (Optional[Union[str, torch.device]]): The device to store the EMA weights on. If None, the EMA
+                        weights will be stored on CPU.
+
+        @crowsonkb's notes on EMA Warmup:
+            If gamma=1 and power=1, implements a simple average. gamma=1, power=2/3 are good values for models you plan
+            to train for a million or more steps (reaches decay factor 0.999 at 31.6K steps, 0.9999 at 1M steps),
+            gamma=1, power=3/4 for models you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999
+            at 215.4k steps).
+        """
+
+        if isinstance(parameters, torch.nn.Module):
+            deprecation_message = (
+                "Passing a `torch.nn.Module` to `ExponentialMovingAverage` is deprecated. "
+                "Please pass the parameters of the module instead."
+            )
+            deprecate(
+                "passing a `torch.nn.Module` to `ExponentialMovingAverage`",
+                "1.0.0",
+                deprecation_message,
+                standard_warn=False,
+            )
+            parameters = parameters.parameters()
+
+            # set use_ema_warmup to True if a torch.nn.Module is passed for backwards compatibility
+            use_ema_warmup = True
+
+        if kwargs.get("max_value", None) is not None:
+            deprecation_message = "The `max_value` argument is deprecated. Please use `decay` instead."
+            deprecate("max_value", "1.0.0", deprecation_message, standard_warn=False)
+            decay = kwargs["max_value"]
+
+        if kwargs.get("min_value", None) is not None:
+            deprecation_message = "The `min_value` argument is deprecated. Please use `min_decay` instead."
+            deprecate("min_value", "1.0.0", deprecation_message, standard_warn=False)
+            min_decay = kwargs["min_value"]
+
+        parameters = list(parameters)
+        self.shadow_params = [p.clone().detach() for p in parameters]
+
+        if kwargs.get("device", None) is not None:
+            deprecation_message = "The `device` argument is deprecated. Please use `to` instead."
+            deprecate("device", "1.0.0", deprecation_message, standard_warn=False)
+            self.to(device=kwargs["device"])
+
+        self.temp_stored_params = None
+
+        self.decay = decay
+        self.min_decay = min_decay
+        self.update_after_step = update_after_step
+        self.use_ema_warmup = use_ema_warmup
+        self.inv_gamma = inv_gamma
+        self.power = power
+        self.optimization_step = 0
+        self.cur_decay_value = None  # set in `step()`
+
+        self.model_cls = model_cls
+        self.model_config = model_config
+
+    @classmethod
+    def extract_ema_kwargs(cls, kwargs):
+        """
+        Extracts the EMA kwargs from the kwargs of a class method.
+        """
+        ema_kwargs = {}
+        for key in [
+            "decay",
+            "min_decay",
+            "optimization_step",
+            "update_after_step",
+            "use_ema_warmup",
+            "inv_gamma",
+            "power",
+        ]:
+            if kwargs.get(key, None) is not None:
+                ema_kwargs[key] = kwargs.pop(key)
+        return ema_kwargs
+
+    @classmethod
+    def from_pretrained(cls, path, model_cls) -> "EMAModel":
+        config = model_cls.load_config(path)
+        ema_kwargs = cls.extract_ema_kwargs(config)
+        model = model_cls.from_pretrained(path)
+
+        ema_model = cls(model.parameters(), model_cls=model_cls, model_config=config)
+
+        ema_model.load_state_dict(ema_kwargs)
+        return ema_model
+
+    def save_pretrained(self, path):
+        if self.model_cls is None:
+            raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")
+
+        if self.model_config is None:
+            raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.")
+
+        model = self.model_cls.from_config(self.model_config)
+        state_dict = self.state_dict()
+        state_dict.pop("shadow_params", None)
+
+        model.register_to_config(**state_dict)
+        self.copy_to(model.parameters())
+        model.save_pretrained(path)
+
+    def get_decay(self, optimization_step: int) -> float:
+        """
+        Compute the decay factor for the exponential moving average.
+        """
+        step = max(0, optimization_step - self.update_after_step - 1)
+
+        if step <= 0:
+            return 0.0
+
+        if self.use_ema_warmup:
+            cur_decay_value = 1 - (1 + step / self.inv_gamma) ** -self.power
+        else:
+            cur_decay_value = (1 + step) / (10 + step)
+
+        cur_decay_value = min(cur_decay_value, self.decay)
+        # make sure decay is not smaller than min_decay
+        cur_decay_value = max(cur_decay_value, self.min_decay)
+        return cur_decay_value
+
+    @torch.no_grad()
+    def step(self, parameters: Iterable[torch.nn.Parameter]):
+        if isinstance(parameters, torch.nn.Module):
+            deprecation_message = (
+                "Passing a `torch.nn.Module` to `ExponentialMovingAverage.step` is deprecated. "
+                "Please pass the parameters of the module instead."
+            )
+            deprecate(
+                "passing a `torch.nn.Module` to `ExponentialMovingAverage.step`",
+                "1.0.0",
+                deprecation_message,
+                standard_warn=False,
+            )
+            parameters = parameters.parameters()
+
+        parameters = list(parameters)
+
+        self.optimization_step += 1
+
+        # Compute the decay factor for the exponential moving average.
+        decay = self.get_decay(self.optimization_step)
+        self.cur_decay_value = decay
+        one_minus_decay = 1 - decay
+
+        context_manager = contextlib.nullcontext
+        if is_transformers_available() and transformers.deepspeed.is_deepspeed_zero3_enabled():
+            import deepspeed
+
+        for s_param, param in zip(self.shadow_params, parameters):
+            if is_transformers_available() and transformers.deepspeed.is_deepspeed_zero3_enabled():
+                context_manager = deepspeed.zero.GatheredParameters(param, modifier_rank=None)
+
+            with context_manager():
+                if param.requires_grad:
+                    s_param.sub_(one_minus_decay * (s_param - param))
+                else:
+                    s_param.copy_(param)
+
+    def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        """
+        Copy current averaged parameters into given collection of parameters.
+
+        Args:
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                updated with the stored moving averages. If `None`, the parameters with which this
+                `ExponentialMovingAverage` was initialized will be used.
+        """
+        parameters = list(parameters)
+        for s_param, param in zip(self.shadow_params, parameters):
+            param.data.copy_(s_param.to(param.device).data)
+
+
+    def to(self, device=None, dtype=None) -> None:
+        r"""Move internal buffers of the ExponentialMovingAverage to `device`.
+
+        Args:
+            device: like `device` argument to `torch.Tensor.to`
+        """
+        # .to() on the tensors handles None correctly
+        self.shadow_params = [
+            p.to(device=device, dtype=dtype) if p.is_floating_point() else p.to(device=device)
+            for p in self.shadow_params
+        ]
+
+    def state_dict(self) -> dict:
+        r"""
+        Returns the state of the ExponentialMovingAverage as a dict. This method is used by accelerate during
+        checkpointing to save the ema state dict.
+        """
+        # Following PyTorch conventions, references to tensors are returned:
+        # "returns a reference to the state and not its copy!" -
+        # https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict
+        return {
+            "decay": self.decay,
+            "min_decay": self.min_decay,
+            "optimization_step": self.optimization_step,
+            "update_after_step": self.update_after_step,
+            "use_ema_warmup": self.use_ema_warmup,
+            "inv_gamma": self.inv_gamma,
+            "power": self.power,
+            "shadow_params": self.shadow_params,
+        }
+
+    def store(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        r"""
+        Args:
+        Save the current parameters for restoring later.
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                temporarily stored.
+        """
+        self.temp_stored_params = [param.detach().cpu().clone() for param in parameters]
+
+    def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None:
+        r"""
+        Args:
+        Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without:
+        affecting the original optimization process. Store the parameters before the `copy_to()` method. After
+        validation (or model saving), use this to restore the former parameters.
+            parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+                updated with the stored parameters. If `None`, the parameters with which this
+                `ExponentialMovingAverage` was initialized will be used.
+        """
+        if self.temp_stored_params is None:
+            raise RuntimeError("This ExponentialMovingAverage has no `store()`ed weights " "to `restore()`")
+        for c_param, param in zip(self.temp_stored_params, parameters):
+            param.data.copy_(c_param.data)
+
+        # Better memory-wise.
+        self.temp_stored_params = None
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        r"""
+        Args:
+        Loads the ExponentialMovingAverage state. This method is used by accelerate during checkpointing to save the
+        ema state dict.
+            state_dict (dict): EMA state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        # deepcopy, to be consistent with module API
+        state_dict = copy.deepcopy(state_dict)
+
+        self.decay = state_dict.get("decay", self.decay)
+        if self.decay < 0.0 or self.decay > 1.0:
+            raise ValueError("Decay must be between 0 and 1")
+
+        self.min_decay = state_dict.get("min_decay", self.min_decay)
+        if not isinstance(self.min_decay, float):
+            raise ValueError("Invalid min_decay")
+
+        self.optimization_step = state_dict.get("optimization_step", self.optimization_step)
+        if not isinstance(self.optimization_step, int):
+            raise ValueError("Invalid optimization_step")
+
+        self.update_after_step = state_dict.get("update_after_step", self.update_after_step)
+        if not isinstance(self.update_after_step, int):
+            raise ValueError("Invalid update_after_step")
+
+        self.use_ema_warmup = state_dict.get("use_ema_warmup", self.use_ema_warmup)
+        if not isinstance(self.use_ema_warmup, bool):
+            raise ValueError("Invalid use_ema_warmup")
+
+        self.inv_gamma = state_dict.get("inv_gamma", self.inv_gamma)
+        if not isinstance(self.inv_gamma, (float, int)):
+            raise ValueError("Invalid inv_gamma")
+
+        self.power = state_dict.get("power", self.power)
+        if not isinstance(self.power, (float, int)):
+            raise ValueError("Invalid power")
+
+        shadow_params = state_dict.get("shadow_params", None)
+        if shadow_params is not None:
+            self.shadow_params = shadow_params
+            if not isinstance(self.shadow_params, list):
+                raise ValueError("shadow_params must be a list")
+            if not all(isinstance(p, torch.Tensor) for p in self.shadow_params):
+                raise ValueError("shadow_params must all be Tensors")
\ No newline at end of file
diff --git a/opensora/utils/lora_utils.py b/opensora/utils/lora_utils.py
new file mode 100644
index 000000000..ce1cec13c
--- /dev/null
+++ b/opensora/utils/lora_utils.py
@@ -0,0 +1,82 @@
+
+from peft import get_peft_model, PeftModel
+import os
+from copy import deepcopy
+import torch
+from opensora.utils.ema import EMAModel
+
+class EMAModel_LoRA(EMAModel):
+    def __init__(self, lora_config, **kwargs):
+        super().__init__(**kwargs)
+        self.lora_config = lora_config
+    
+    @classmethod
+    def from_pretrained(cls, path, model_cls, lora_config, origin_model_path) -> "EMAModel":
+        # 1. load origin model
+        origin_model = model_cls.from_pretrained(origin_model_path)  # origin_model
+        # 2. convert to lora model automatically and load lora weight
+        lora_model = PeftModel.from_pretrained(origin_model, path)  # lora_origin_model
+        # 3. ema the whole lora_model
+        ema_model = cls(lora_config, parameters=lora_model.parameters(), model_cls=model_cls, model_config=origin_model.config)
+        # 4. load ema_kwargs, e.g decay...
+        ema_kwargs = torch.load(os.path.join(path, 'ema_kwargs.pt'))
+        ema_model.load_state_dict(ema_kwargs)
+        return ema_model
+
+    def save_pretrained(self, path):
+        if self.model_cls is None:
+            raise ValueError("`save_pretrained` can only be used if `model_cls` was defined at __init__.")
+
+        if self.model_config is None:
+            raise ValueError("`save_pretrained` can only be used if `model_config` was defined at __init__.")
+        # 1. init a base model randomly
+        model = self.model_cls.from_config(self.model_config)
+        # 2. convert lora_model
+        lora_model = get_peft_model(model, self.lora_config)
+        # 3. ema_lora_model weight to lora_model
+        self.copy_to(lora_model.parameters())
+        # 4. save lora weight
+        lora_model.save_pretrained(path)  # only lora weight
+        # 5. save ema_kwargs, e.g decay...
+        state_dict = self.state_dict()  # lora_model weight
+        state_dict.pop("shadow_params", None)
+        torch.save(state_dict, os.path.join(path, 'ema_kwargs.pt'))
+
+
+def maybe_zero_3(param, ignore_status=False, name=None):
+    from deepspeed import zero
+    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+    if hasattr(param, "ds_id"):
+        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
+            if not ignore_status:
+                logging.warning(f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}")
+        with zero.GatheredParameters([param]):
+            param = param.data.detach().cpu().clone()
+    else:
+        param = param.detach().cpu().clone()
+    return param
+
+# Borrowed from peft.utils.get_peft_model_state_dict
+def get_peft_state_maybe_zero_3(named_params, bias):
+    if bias == "none":
+        to_return = {k: t for k, t in named_params if "lora_" in k}
+    elif bias == "all":
+        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
+    elif bias == "lora_only":
+        to_return = {}
+        maybe_lora_bias = {}
+        lora_bias_names = set()
+        for k, t in named_params:
+            if "lora_" in k:
+                to_return[k] = t
+                bias_name = k.split("lora_")[0] + "bias"
+                lora_bias_names.add(bias_name)
+            elif "bias" in k:
+                maybe_lora_bias[k] = t
+        for k, t in maybe_lora_bias:
+            if bias_name in lora_bias_names:
+                to_return[bias_name] = t
+    else:
+        raise NotImplementedError
+    to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
+    return to_return
diff --git a/opensora/utils/parallel_states.py b/opensora/utils/parallel_states.py
new file mode 100644
index 000000000..def564f71
--- /dev/null
+++ b/opensora/utils/parallel_states.py
@@ -0,0 +1,44 @@
+import torch
+import torch.distributed as dist
+import os
+
+class COMM_INFO:
+    def __init__(self):
+        self.group = None
+        self.world_size = 0
+        self.rank = -1
+
+nccl_info = COMM_INFO()
+_SEQUENCE_PARALLEL_STATE = False
+def initialize_sequence_parallel_state(sequence_parallel_size):
+    global _SEQUENCE_PARALLEL_STATE
+    if sequence_parallel_size > 1:
+        _SEQUENCE_PARALLEL_STATE = True
+        initialize_sequence_parallel_group(sequence_parallel_size)
+
+def set_sequence_parallel_state(state):
+    global _SEQUENCE_PARALLEL_STATE
+    _SEQUENCE_PARALLEL_STATE = state
+
+def get_sequence_parallel_state():
+    return _SEQUENCE_PARALLEL_STATE
+
+def initialize_sequence_parallel_group(sequence_parallel_size):
+    """Initialize the sequence parallel group."""
+    rank = int(os.getenv('RANK', '0'))
+    world_size = int(os.getenv("WORLD_SIZE", '1'))
+    assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size"
+    # hccl
+    nccl_info.world_size = sequence_parallel_size
+    nccl_info.rank = rank
+    num_sequence_parallel_groups: int = world_size // sequence_parallel_size
+    for i in range(num_sequence_parallel_groups):
+        ranks = range(i * sequence_parallel_size, (i + 1) * sequence_parallel_size)
+        group = dist.new_group(ranks)
+        if rank in ranks:
+            nccl_info.group = group
+
+
+def destroy_sequence_parallel_group():
+    """Destroy the sequence parallel group."""
+    dist.destroy_process_group()
diff --git a/opensora/utils/taming_download.py b/opensora/utils/taming_download.py
new file mode 100644
index 000000000..5a62be778
--- /dev/null
+++ b/opensora/utils/taming_download.py
@@ -0,0 +1,145 @@
+"""Modified from https://github.com/CompVis/taming-transformers.git"""
+
+import os, hashlib
+import requests
+from tqdm import tqdm
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class KeyNotFoundError(Exception):
+    def __init__(self, cause, keys=None, visited=None):
+        self.cause = cause
+        self.keys = keys
+        self.visited = visited
+        messages = list()
+        if keys is not None:
+            messages.append("Key not found: {}".format(keys))
+        if visited is not None:
+            messages.append("Visited: {}".format(visited))
+        messages.append("Cause:\n{}".format(cause))
+        message = "\n".join(messages)
+        super().__init__(message)
+
+
+def retrieve(
+    list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False
+):
+    """Given a nested list or dict return the desired value at key expanding
+    callable nodes if necessary and :attr:`expand` is ``True``. The expansion
+    is done in-place.
+
+    Parameters
+    ----------
+        list_or_dict : list or dict
+            Possibly nested list or dictionary.
+        key : str
+            key/to/value, path like string describing all keys necessary to
+            consider to get to the desired value. List indices can also be
+            passed here.
+        splitval : str
+            String that defines the delimiter between keys of the
+            different depth levels in `key`.
+        default : obj
+            Value returned if :attr:`key` is not found.
+        expand : bool
+            Whether to expand callable nodes on the path or not.
+
+    Returns
+    -------
+        The desired value or if :attr:`default` is not ``None`` and the
+        :attr:`key` is not found returns ``default``.
+
+    Raises
+    ------
+        Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is
+        ``None``.
+    """
+
+    keys = key.split(splitval)
+
+    success = True
+    try:
+        visited = []
+        parent = None
+        last_key = None
+        for key in keys:
+            if callable(list_or_dict):
+                if not expand:
+                    raise KeyNotFoundError(
+                        ValueError(
+                            "Trying to get past callable node with expand=False."
+                        ),
+                        keys=keys,
+                        visited=visited,
+                    )
+                list_or_dict = list_or_dict()
+                parent[last_key] = list_or_dict
+
+            last_key = key
+            parent = list_or_dict
+
+            try:
+                if isinstance(list_or_dict, dict):
+                    list_or_dict = list_or_dict[key]
+                else:
+                    list_or_dict = list_or_dict[int(key)]
+            except (KeyError, IndexError, ValueError) as e:
+                raise KeyNotFoundError(e, keys=keys, visited=visited)
+
+            visited += [key]
+        # final expansion of retrieved value
+        if expand and callable(list_or_dict):
+            list_or_dict = list_or_dict()
+            parent[last_key] = list_or_dict
+    except KeyNotFoundError as e:
+        if default is None:
+            raise e
+        else:
+            list_or_dict = default
+            success = False
+
+    if not pass_success:
+        return list_or_dict
+    else:
+        return list_or_dict, success
+
diff --git a/opensora/utils/utils.py b/opensora/utils/utils.py
new file mode 100644
index 000000000..7cd353662
--- /dev/null
+++ b/opensora/utils/utils.py
@@ -0,0 +1,471 @@
+import os
+
+import torch
+
+import os
+import math
+import torch
+import logging
+import random
+import subprocess
+import numpy as np
+import torch.distributed as dist
+
+# from torch._six import inf
+from torch import inf
+from PIL import Image
+from typing import Union, Iterable
+import collections
+from collections import OrderedDict
+from torch.utils.tensorboard import SummaryWriter
+
+from diffusers.utils import is_bs4_available, is_ftfy_available
+
+import html
+import re
+import urllib.parse as ul
+
+if is_bs4_available():
+    from bs4 import BeautifulSoup
+
+if is_ftfy_available():
+    import ftfy
+
+_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]
+
+def to_2tuple(x):
+    if isinstance(x, collections.abc.Iterable):
+        return x
+    return (x, x)
+
+def find_model(model_name):
+    """
+    Finds a pre-trained Latte model, downloading it if necessary. Alternatively, loads a model from a local path.
+    """
+    assert os.path.isfile(model_name), f'Could not find Latte checkpoint at {model_name}'
+    checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
+
+    # if "ema" in checkpoint:  # supports checkpoints from train.py
+    #     print('Using Ema!')
+    #     checkpoint = checkpoint["ema"]
+    # else:
+    print('Using model!')
+    checkpoint = checkpoint['model']
+    return checkpoint
+
+#################################################################################
+#                             Training Clip Gradients                           #
+#################################################################################
+
+def get_grad_norm(
+        parameters: _tensor_or_tensors, norm_type: float = 2.0) -> torch.Tensor:
+    r"""
+    Copy from torch.nn.utils.clip_grad_norm_
+
+    Clips gradient norm of an iterable of parameters.
+
+    The norm is computed over all gradients together, as if they were
+    concatenated into a single vector. Gradients are modified in-place.
+
+    Args:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        max_norm (float or int): max norm of the gradients
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+        error_if_nonfinite (bool): if True, an error is thrown if the total
+            norm of the gradients from :attr:`parameters` is ``nan``,
+            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+
+    Returns:
+        Total norm of the parameter gradients (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    grads = [p.grad for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(grads) == 0:
+        return torch.tensor(0.)
+    device = grads[0].device
+    if norm_type == inf:
+        norms = [g.detach().abs().max().to(device) for g in grads]
+        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
+    return total_norm
+
+
+def clip_grad_norm_(
+        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
+        error_if_nonfinite: bool = False, clip_grad=True) -> torch.Tensor:
+    r"""
+    Copy from torch.nn.utils.clip_grad_norm_
+
+    Clips gradient norm of an iterable of parameters.
+
+    The norm is computed over all gradients together, as if they were
+    concatenated into a single vector. Gradients are modified in-place.
+
+    Args:
+        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+            single Tensor that will have gradients normalized
+        max_norm (float or int): max norm of the gradients
+        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+            infinity norm.
+        error_if_nonfinite (bool): if True, an error is thrown if the total
+            norm of the gradients from :attr:`parameters` is ``nan``,
+            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+
+    Returns:
+        Total norm of the parameter gradients (viewed as a single vector).
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    grads = [p.grad for p in parameters if p.grad is not None]
+    max_norm = float(max_norm)
+    norm_type = float(norm_type)
+    if len(grads) == 0:
+        return torch.tensor(0.)
+    device = grads[0].device
+    if norm_type == inf:
+        norms = [g.detach().abs().max().to(device) for g in grads]
+        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
+
+    if clip_grad:
+        if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
+            raise RuntimeError(
+                f'The total norm of order {norm_type} for gradients from '
+                '`parameters` is non-finite, so it cannot be clipped. To disable '
+                'this error and scale the gradients by the non-finite norm anyway, '
+                'set `error_if_nonfinite=False`')
+        clip_coef = max_norm / (total_norm + 1e-6)
+        # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
+        # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
+        # when the gradients do not reside in CPU memory.
+        clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
+        for g in grads:
+            g.detach().mul_(clip_coef_clamped.to(g.device))
+        # gradient_cliped = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
+        # print(gradient_cliped)
+    return total_norm
+
+
+def get_experiment_dir(root_dir, args):
+    # if args.pretrained is not None and 'Latte-XL-2-256x256.pt' not in args.pretrained:
+    #     root_dir += '-WOPRE'
+    if args.use_compile:
+        root_dir += '-Compile'  # speedup by torch compile
+    if args.attention_mode:
+        root_dir += f'-{args.attention_mode.upper()}'
+    # if args.enable_xformers_memory_efficient_attention:
+    #     root_dir += '-Xfor'
+    if args.gradient_checkpointing:
+        root_dir += '-Gc'
+    if args.mixed_precision:
+        root_dir += f'-{args.mixed_precision.upper()}'
+    root_dir += f'-{args.max_image_size}'
+    return root_dir
+
+def get_precision(args):
+    if args.mixed_precision == "bf16":
+        dtype = torch.bfloat16
+    elif args.mixed_precision == "fp16":
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+    return dtype
+
+#################################################################################
+#                             Training Logger                                   #
+#################################################################################
+
+def create_logger(logging_dir):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    if dist.get_rank() == 0:  # real logger
+        logging.basicConfig(
+            level=logging.INFO,
+            # format='[\033[34m%(asctime)s\033[0m] %(message)s',
+            format='[%(asctime)s] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S',
+            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+        )
+        logger = logging.getLogger(__name__)
+
+    else:  # dummy logger (does nothing)
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+def create_tensorboard(tensorboard_dir):
+    """
+    Create a tensorboard that saves losses.
+    """
+    if dist.get_rank() == 0:  # real tensorboard
+        # tensorboard
+        writer = SummaryWriter(tensorboard_dir)
+
+    return writer
+
+
+def write_tensorboard(writer, *args):
+    '''
+    write the loss information to a tensorboard file.
+    Only for pytorch DDP mode.
+    '''
+    if dist.get_rank() == 0:  # real tensorboard
+        writer.add_scalar(args[0], args[1], args[2])
+
+
+#################################################################################
+#                      EMA Update/ DDP Training Utils                           #
+#################################################################################
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def cleanup():
+    """
+    End DDP training.
+    """
+    dist.destroy_process_group()
+
+
+def setup_distributed(backend="nccl", port=None):
+    """Initialize distributed training environment.
+    support both slurm and torch.distributed.launch
+    see torch.distributed.init_process_group() for more details
+    """
+    num_gpus = torch.cuda.device_count()
+
+    if "SLURM_JOB_ID" in os.environ:
+        rank = int(os.environ["SLURM_PROCID"])
+        world_size = int(os.environ["SLURM_NTASKS"])
+        node_list = os.environ["SLURM_NODELIST"]
+        addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1")
+        # specify master port
+        if port is not None:
+            os.environ["MASTER_PORT"] = str(port)
+        elif "MASTER_PORT" not in os.environ:
+            # os.environ["MASTER_PORT"] = "29566"
+            os.environ["MASTER_PORT"] = str(29567 + num_gpus)
+        if "MASTER_ADDR" not in os.environ:
+            os.environ["MASTER_ADDR"] = addr
+        os.environ["WORLD_SIZE"] = str(world_size)
+        os.environ["LOCAL_RANK"] = str(rank % num_gpus)
+        os.environ["RANK"] = str(rank)
+    else:
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+
+    # torch.cuda.set_device(rank % num_gpus)
+
+    dist.init_process_group(
+        backend=backend,
+        world_size=world_size,
+        rank=rank,
+    )
+
+
+#################################################################################
+#                             Testing  Utils                                    #
+#################################################################################
+
+def save_video_grid(video, nrow=None):
+    b, t, h, w, c = video.shape
+
+    if nrow is None:
+        nrow = math.ceil(math.sqrt(b))
+    ncol = math.ceil(b / nrow)
+    padding = 1
+    video_grid = torch.zeros((t, (padding + h) * nrow + padding,
+                              (padding + w) * ncol + padding, c), dtype=torch.uint8)
+
+    print(video_grid.shape)
+    for i in range(b):
+        r = i // ncol
+        c = i % ncol
+        start_r = (padding + h) * r
+        start_c = (padding + w) * c
+        video_grid[:, start_r:start_r + h, start_c:start_c + w] = video[i]
+
+    return video_grid
+
+
+#################################################################################
+#                             MMCV  Utils                                    #
+#################################################################################
+
+
+def collect_env():
+    # Copyright (c) OpenMMLab. All rights reserved.
+    from mmcv.utils import collect_env as collect_base_env
+    from mmcv.utils import get_git_hash
+    """Collect the information of the running environments."""
+
+    env_info = collect_base_env()
+    env_info['MMClassification'] = get_git_hash()[:7]
+
+    for name, val in env_info.items():
+        print(f'{name}: {val}')
+
+    print(torch.cuda.get_arch_list())
+    print(torch.version.cuda)
+
+
+#################################################################################
+#                          Pixart-alpha  Utils                                  #
+#################################################################################
+
+
+bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
+
+def text_preprocessing(text, support_Chinese=True):
+    # The exact text cleaning as was in the training stage:
+    text = clean_caption(text, support_Chinese=support_Chinese)
+    text = clean_caption(text, support_Chinese=support_Chinese)
+    return text
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+def clean_caption(caption, support_Chinese=True):
+    caption = str(caption)
+    caption = ul.unquote_plus(caption)
+    caption = caption.strip().lower()
+    caption = re.sub('<person>', 'person', caption)
+    # urls:
+    caption = re.sub(
+        r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+        '', caption)  # regex for urls
+    caption = re.sub(
+        r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+        '', caption)  # regex for urls
+    # html:
+    caption = BeautifulSoup(caption, features='html.parser').text
+
+    # @<nickname>
+    caption = re.sub(r'@[\w\d]+\b', '', caption)
+
+    # 31C0—31EF CJK Strokes
+    # 31F0—31FF Katakana Phonetic Extensions
+    # 3200—32FF Enclosed CJK Letters and Months
+    # 3300—33FF CJK Compatibility
+    # 3400—4DBF CJK Unified Ideographs Extension A
+    # 4DC0—4DFF Yijing Hexagram Symbols
+    # 4E00—9FFF CJK Unified Ideographs
+    caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
+    caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
+    caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
+    caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
+    caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
+    caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
+    if not support_Chinese:
+        caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)  # Chinese
+    #######################################################
+
+    # все виды тире / all types of dash --> "-"
+    caption = re.sub(
+        r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
+        '-', caption)
+
+    # кавычки к одному стандарту
+    caption = re.sub(r'[`´«»“”¨]', '"', caption)
+    caption = re.sub(r'[‘’]', "'", caption)
+
+    # &quot;
+    caption = re.sub(r'&quot;?', '', caption)
+    # &amp
+    caption = re.sub(r'&amp', '', caption)
+
+    # ip adresses:
+    caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
+
+    # article ids:
+    caption = re.sub(r'\d:\d\d\s+$', '', caption)
+
+    # \n
+    caption = re.sub(r'\\n', ' ', caption)
+
+    # "#123"
+    caption = re.sub(r'#\d{1,3}\b', '', caption)
+    # "#12345.."
+    caption = re.sub(r'#\d{5,}\b', '', caption)
+    # "123456.."
+    caption = re.sub(r'\b\d{6,}\b', '', caption)
+    # filenames:
+    caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
+
+    #
+    caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
+    caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
+
+    caption = re.sub(bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+    caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
+
+    # this-is-my-cute-cat / this_is_my_cute_cat
+    regex2 = re.compile(r'(?:\-|\_)')
+    if len(re.findall(regex2, caption)) > 3:
+        caption = re.sub(regex2, ' ', caption)
+
+    caption = basic_clean(caption)
+
+    caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
+    caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
+    caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
+
+    caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
+    caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
+    caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
+    caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
+    caption = re.sub(r'\bpage\s+\d+\b', '', caption)
+
+    caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
+
+    caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
+
+    caption = re.sub(r'\b\s+\:\s+', r': ', caption)
+    caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
+    caption = re.sub(r'\s+', ' ', caption)
+
+    caption.strip()
+
+    caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
+    caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
+    caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
+    caption = re.sub(r'^\.\S+$', '', caption)
+
+    return caption.strip()
+
+
+if __name__ == '__main__':
+    
+    # caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
+    a = "امرأة مسنة بشعر أبيض ووجه مليء بالتجاعيد تجلس داخل سيارة قديمة الطراز، تنظر من خلال النافذة الجانبية بتعبير تأملي أو حزين قليلاً."
+    print(a)
+    print(text_preprocessing(a))
+
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 000000000..431e0c32e
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,49 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "opensora"
+version = "1.2.0"
+description = "Reproduce OpenAI's Sora."
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+]
+dependencies = [
+    "torch==2.1.0", "torchvision==0.16.0",
+    "transformers==4.40.1", "accelerate==0.29.3", "tokenizers==0.19.1", "diffusers==0.28.0", 
+    "albumentations==1.4.0", "av==11.0.0", "decord==0.6.0", "einops==0.7.0", "fastapi==0.110.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==2.2.1", "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", "test-tube==0.7.5",
+    "timm==0.9.16", "torchdiffeq==0.2.3", "torchmetrics==1.3.2", "tqdm==4.66.2", "urllib3==2.2.1", "uvicorn==0.27.1",
+    "scikit-video==1.1.11", "imageio-ffmpeg==0.4.9", "sentencepiece==0.1.99", "beautifulsoup4==4.12.3",
+    "ftfy==6.1.3", "moviepy==1.0.3", "wandb==0.16.3", "tensorboard==2.14.0", "pydantic==2.6.4", "gradio==4.0.0", "xformers==0.0.22.post7"
+]
+
+[project.optional-dependencies]
+train = ["deepspeed==0.12.6"]
+dev = ["mypy==1.8.0"]
+
+
+[project.urls]
+"Homepage" = "https://github.com/PKU-YuanGroup/Open-Sora-Plan"
+"Bug Tracker" = "https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues"
+
+[tool.setuptools.packages.find]
+exclude = ["assets*", "docker*", "docs", "scripts*"]
+
+[tool.wheel]
+exclude = ["assets*", "docker*", "docs", "scripts*"]
+
+[tool.mypy]
+warn_return_any = true
+warn_unused_configs = true
+ignore_missing_imports = true
+disallow_untyped_calls = true
+check_untyped_defs = true
+no_implicit_optional = true
diff --git a/scripts/accelerate_configs/ddp_config.yaml b/scripts/accelerate_configs/ddp_config.yaml
new file mode 100644
index 000000000..fa6cfbdd6
--- /dev/null
+++ b/scripts/accelerate_configs/ddp_config.yaml
@@ -0,0 +1,11 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: MULTI_GPU
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/deepspeed_zero2_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_config.yaml
new file mode 100644
index 000000000..03ce04191
--- /dev/null
+++ b/scripts/accelerate_configs/deepspeed_zero2_config.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero2.json
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29503
+main_training_function: main
+num_machines: 1
+num_processes: 8
+# gpu_ids: 0
+use_cpu: false
diff --git a/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
new file mode 100644
index 000000000..9c96bae67
--- /dev/null
+++ b/scripts/accelerate_configs/deepspeed_zero2_offload_config.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero2_offload.json
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/deepspeed_zero3_config.yaml b/scripts/accelerate_configs/deepspeed_zero3_config.yaml
new file mode 100644
index 000000000..b368b2be7
--- /dev/null
+++ b/scripts/accelerate_configs/deepspeed_zero3_config.yaml
@@ -0,0 +1,14 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero3.json
+ zero3_init_flag: true
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml b/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml
new file mode 100644
index 000000000..b01631caf
--- /dev/null
+++ b/scripts/accelerate_configs/deepspeed_zero3_offload_config.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero3_offload.json
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/default_config.yaml b/scripts/accelerate_configs/default_config.yaml
new file mode 100644
index 000000000..ac3bb99b5
--- /dev/null
+++ b/scripts/accelerate_configs/default_config.yaml
@@ -0,0 +1,12 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: MULTI_GPU
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: 29501
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+gpu_ids: 0,1,2,3,4,5,6,7
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/hostfile b/scripts/accelerate_configs/hostfile
new file mode 100644
index 000000000..becb02e09
--- /dev/null
+++ b/scripts/accelerate_configs/hostfile
@@ -0,0 +1,16 @@
+node30 slots=8
+node005 slots=8
+node006 slots=8
+node007 slots=8
+node008 slots=8
+node010 slots=8
+node012 slots=8
+node013 slots=8
+node019 slots=8
+node026 slots=8
+node027 slots=8
+node028 slots=8
+node029 slots=8
+node031 slots=8
+node059 slots=8
+node060 slots=8
\ No newline at end of file
diff --git a/scripts/accelerate_configs/multi_node_example.yaml b/scripts/accelerate_configs/multi_node_example.yaml
new file mode 100644
index 000000000..43a9f7f48
--- /dev/null
+++ b/scripts/accelerate_configs/multi_node_example.yaml
@@ -0,0 +1,18 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: DEEPSPEED
+deepspeed_config:
+ deepspeed_config_file: scripts/accelerate_configs/zero2.json
+ deepspeed_hostfile: scripts/accelerate_configs/hostfile
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: 100.64.24.30
+main_process_port: 29522
+main_training_function: main
+num_machines: 16
+num_processes: 128
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/scripts/accelerate_configs/multi_node_example_by_ddp.yaml b/scripts/accelerate_configs/multi_node_example_by_ddp.yaml
new file mode 100644
index 000000000..17fd420ac
--- /dev/null
+++ b/scripts/accelerate_configs/multi_node_example_by_ddp.yaml
@@ -0,0 +1,13 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: MULTI_GPU
+fsdp_config: {}
+main_process_port: 29501
+main_training_function: main
+num_machines: 2
+num_processes: 16
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero2.json b/scripts/accelerate_configs/zero2.json
new file mode 100644
index 000000000..8dcdd154d
--- /dev/null
+++ b/scripts/accelerate_configs/zero2.json
@@ -0,0 +1,25 @@
+{
+    "fp16": {
+        "enabled": false,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "communication_data_type": "fp32",
+    "gradient_clipping": 1.0,
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 5e8
+    }
+}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero2_npu.json b/scripts/accelerate_configs/zero2_npu.json
new file mode 100644
index 000000000..6f1ecdf2c
--- /dev/null
+++ b/scripts/accelerate_configs/zero2_npu.json
@@ -0,0 +1,25 @@
+{
+    "fp16": {
+        "enabled": false,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "communication_data_type": "fp32",
+    "gradient_clipping": 1.0,
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": true,
+        "allgather_bucket_size": 536870912,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": 536870912
+    }
+}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero2_offload.json b/scripts/accelerate_configs/zero2_offload.json
new file mode 100644
index 000000000..8546e6ee1
--- /dev/null
+++ b/scripts/accelerate_configs/zero2_offload.json
@@ -0,0 +1,26 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "offload_optimizer": {
+            "device": "cpu"
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 5e8
+    }
+}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero3.json b/scripts/accelerate_configs/zero3.json
new file mode 100644
index 000000000..def4450c2
--- /dev/null
+++ b/scripts/accelerate_configs/zero3.json
@@ -0,0 +1,28 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": 5e8,
+        "stage3_prefetch_bucket_size": 5e8,
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    }
+}
\ No newline at end of file
diff --git a/scripts/accelerate_configs/zero3_offload.json b/scripts/accelerate_configs/zero3_offload.json
new file mode 100644
index 000000000..c15bcf91b
--- /dev/null
+++ b/scripts/accelerate_configs/zero3_offload.json
@@ -0,0 +1,39 @@
+{
+  "fp16": {
+    "enabled": "auto",
+    "loss_scale": 0,
+    "loss_scale_window": 1000,
+    "initial_scale_power": 16,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+  "zero_optimization": {
+    "stage": 3,
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "sub_group_size": 1e9,
+    "reduce_bucket_size": 5e8,
+    "stage3_prefetch_bucket_size": "auto",
+    "stage3_param_persistence_threshold": "auto",
+    "stage3_max_live_parameters": 1e9,
+    "stage3_max_reuse_distance": 1e9,
+    "gather_16bit_weights_on_model_save": true
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "steps_per_print": 1e5,
+  "wall_clock_breakdown": false
+}
\ No newline at end of file
diff --git a/scripts/causalvae/eval.sh b/scripts/causalvae/eval.sh
new file mode 100644
index 000000000..0b7b909cc
--- /dev/null
+++ b/scripts/causalvae/eval.sh
@@ -0,0 +1,21 @@
+# REAL_DATASET_DIR=/remote-home1/dataset/OpenMMLab___Kinetics-400/raw/Kinetics-400/videos_val/
+REAL_DATASET_DIR=../dataset/webvid/videos
+EXP_NAME=decoder
+SAMPLE_RATE=3
+NUM_FRAMES=33
+RESOLUTION=256
+SUBSET_SIZE=50
+METRIC=ssim
+
+python opensora/models/causalvideovae/eval/eval_common_metric.py \
+    --batch_size 1 \
+    --real_video_dir ${REAL_DATASET_DIR} \
+    --generated_video_dir /remote-home1/lzj/dataset/gen/${EXP_NAME}_sr${SAMPLE_RATE}_nf${NUM_FRAMES}_res${RESOLUTION}_subset${SUBSET_SIZE} \
+    --device cuda:0 \
+    --sample_fps 3 \
+    --sample_rate ${SAMPLE_RATE} \
+    --num_frames ${NUM_FRAMES} \
+    --resolution ${RESOLUTION} \
+    --subset_size ${SUBSET_SIZE} \
+    --crop_size ${RESOLUTION} \
+    --metric ${METRIC}
\ No newline at end of file
diff --git a/scripts/causalvae/prepare_eval.sh b/scripts/causalvae/prepare_eval.sh
new file mode 100644
index 000000000..8ce0b5ba8
--- /dev/null
+++ b/scripts/causalvae/prepare_eval.sh
@@ -0,0 +1,24 @@
+export CUDA_VISIBLE_DEVICES=0
+REAL_DATASET_DIR=valid/
+EXP_NAME=decoder
+SAMPLE_RATE=1
+NUM_FRAMES=33
+RESOLUTION=512
+SUBSET_SIZE=1
+CKPT=/storage/lcm/Causal-Video-VAE/results/488dim8
+
+python opensora/models/causalvideovae/sample/rec_video_vae.py \
+    --batch_size 1 \
+    --real_video_dir ${REAL_DATASET_DIR} \
+    --generated_video_dir valid_gen/${EXP_NAME}_sr${SAMPLE_RATE}_nf${NUM_FRAMES}_res${RESOLUTION}_subset${SUBSET_SIZE} \
+    --device cuda \
+    --sample_fps 24 \
+    --sample_rate ${SAMPLE_RATE} \
+    --num_frames ${NUM_FRAMES} \
+    --resolution ${RESOLUTION} \
+    --subset_size ${SUBSET_SIZE} \
+    --crop_size ${RESOLUTION} \
+    --num_workers 8 \
+    --ckpt ${CKPT} \
+    --output_origin \
+    --enable_tiling
\ No newline at end of file
diff --git a/scripts/causalvae/rec_image.sh b/scripts/causalvae/rec_image.sh
new file mode 100644
index 000000000..3c65c54b3
--- /dev/null
+++ b/scripts/causalvae/rec_image.sh
@@ -0,0 +1,8 @@
+CUDA_VISIBLE_DEVICES=0 python examples/rec_image.py \
+    --ae_path "/storage/dataset/test140k" \
+    --image_path /storage/dataset/image/anytext3m/ocr_data/Art/images/gt_5544.jpg \
+    --rec_path rec.jpg \
+    --device cuda \
+    --short_size 512 \
+    --ae CausalVAEModel_4x8x8 \
+    --enable_tiling
\ No newline at end of file
diff --git a/scripts/causalvae/rec_video.sh b/scripts/causalvae/rec_video.sh
new file mode 100644
index 000000000..711000a3f
--- /dev/null
+++ b/scripts/causalvae/rec_video.sh
@@ -0,0 +1,11 @@
+CUDA_VISIBLE_DEVICES=0 python examples/rec_video.py \
+    --ae_path "/storage/dataset/488dim4_plus" \
+    --video_path /storage/clh/gen/64dim_488/origin/134445.mp4 \
+    --rec_path rec.mp4 \
+    --device cuda \
+    --sample_rate 1 \
+    --num_frames 65 \
+    --height 480 \
+    --width 640 \
+    --ae CausalVAEModel_4x8x8 \
+    --enable_tiling --tile_overlap_factor 0.125 --save_memory
diff --git a/scripts/causalvae/train.sh b/scripts/causalvae/train.sh
new file mode 100644
index 000000000..d97ba141f
--- /dev/null
+++ b/scripts/causalvae/train.sh
@@ -0,0 +1,45 @@
+# export https_proxy=http://127.0.0.1:8998
+# export http_proxy=http://127.0.0.1:8998
+unset https_proxy
+unset http_proxy
+export WANDB_PROJECT=causalvideovae_2.0
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+export GLOO_SOCKET_IFNAME=bond0
+export NCCL_SOCKET_IFNAME=bond0
+export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1
+export NCCL_IB_GID_INDEX=3
+export NCCL_IB_TC=162
+export NCCL_IB_TIMEOUT=22
+export NCCL_PXN_DISABLE=0
+export NCCL_IB_QPS_PER_CONNECTION=4
+
+EXP_NAME=latent4_resume_2
+
+torchrun \
+    --nnodes=1 --nproc_per_node=8 \
+    --master_addr=localhost \
+    --master_port=29600 \
+    scripts/causalvae/train_causalvae.py \
+    --exp_name ${EXP_NAME} \
+    --model_config scripts/config.json \
+    --video_path /storage/dataset/pexels/ \
+    --eval_video_path /storage/dataset/pixabay_v2/folder_01 \
+    --resolution 320 \
+    --epochs 1000 \
+    --num_frames 25 \
+    --batch_size 1 \
+    --disc_start 2000 \
+    --save_ckpt_step 2000 \
+    --eval_steps 500 \
+    --eval_batch_size 1 \
+    --eval_num_frames 33 \
+    --eval_sample_rate 3 \
+    --eval_subset_size 50 \
+    --eval_lpips \
+    --ema \
+    --ema_decay 0.999 \
+    --perceptual_weight 1.0 \
+    --loss_type l1 \
+    --disc_cls causalvideovae.model.losses.LPIPSWithDiscriminator3D \
+    --not_resume_training_process \
+    --resume_from_checkpoint /storage/lcm/Causal-Video-VAE/results/latent8_3d-lr1.00e-05-bs1-rs320-sr2-fr25/checkpoint-14000.ckpt
\ No newline at end of file
diff --git a/scripts/slurm/placeholder b/scripts/slurm/placeholder
new file mode 100644
index 000000000..e69de29bb
diff --git a/scripts/text_condition/gpu/sample_inpaint.sh b/scripts/text_condition/gpu/sample_inpaint.sh
new file mode 100644
index 000000000..1a340aa3f
--- /dev/null
+++ b/scripts/text_condition/gpu/sample_inpaint.sh
@@ -0,0 +1,18 @@
+CUDA_VISIBLE_DEVICES=7 python -m opensora.sample.sample_inpaint \
+    --model_path /path/to/checkpoint-xxx/model \
+    --num_frames 93 \
+    --height 480 \
+    --width 640 \
+    --cache_dir "../cache_dir" \
+    --text_encoder_name google/mt5-xxl \
+    --text_prompt examples/prompt_list_0.txt \
+    --conditional_images_path examples/conditional_images_path.txt \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --save_img_path "./sample_video_93x480p_cfg7.5_step100" \
+    --fps 24 \
+    --guidance_scale 7.5 \
+    --num_sampling_steps 100 \
+    --enable_tiling \
+    --max_sequence_length 512 \
+    --sample_method PDNM \
diff --git a/scripts/text_condition/gpu/sample_inpaint_ddp.sh b/scripts/text_condition/gpu/sample_inpaint_ddp.sh
new file mode 100644
index 000000000..a2706e7b9
--- /dev/null
+++ b/scripts/text_condition/gpu/sample_inpaint_ddp.sh
@@ -0,0 +1,21 @@
+
+torchrun --nnodes=1 --nproc_per_node 8  --master_port 29503 \
+    -m opensora.sample.sample_inpaint_ddp \
+    --model_path /path/to/checkpoint-xxx/model \
+    --num_frames 93 \
+    --height 480 \
+    --width 640 \
+    --cache_dir "../cache_dir" \
+    --text_encoder_name google/mt5-xxl \
+    --text_prompt examples/prompt_list_0.txt \
+    --conditional_images_path examples/conditional_images_path.txt \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --save_img_path "./sample_video_93x480p_cfg7.5_step100" \
+    --fps 24 \
+    --guidance_scale 7.5 \
+    --num_sampling_steps 100 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --max_sequence_length 512 \
+    --sample_method PNDM \
\ No newline at end of file
diff --git a/scripts/text_condition/gpu/sample_inpaint_sp.sh b/scripts/text_condition/gpu/sample_inpaint_sp.sh
new file mode 100644
index 000000000..9aa9c1b4c
--- /dev/null
+++ b/scripts/text_condition/gpu/sample_inpaint_sp.sh
@@ -0,0 +1,21 @@
+
+torchrun --nnodes=1 --nproc_per_node 8  --master_port 29503 \
+    -m opensora.sample.sample_inpaint_sp \
+    --model_path /path/to/checkpoint-xxx/model \
+    --num_frames 93 \
+    --height 480 \
+    --width 640 \
+    --cache_dir "../cache_dir" \
+    --text_encoder_name google/mt5-xxl \
+    --text_prompt examples/prompt_list_0.txt \
+    --conditional_images_path examples/conditional_images_path.txt \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --save_img_path "./sample_video_93x480p_cfg7.5_step100" \
+    --fps 24 \
+    --guidance_scale 7.5 \
+    --num_sampling_steps 100 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --max_sequence_length 512 \
+    --sample_method PNDM \
\ No newline at end of file
diff --git a/scripts/text_condition/gpu/sample_t2v.sh b/scripts/text_condition/gpu/sample_t2v.sh
new file mode 100644
index 000000000..67dc33da6
--- /dev/null
+++ b/scripts/text_condition/gpu/sample_t2v.sh
@@ -0,0 +1,18 @@
+CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \
+    --model_path /path/to/checkpoint-xxx/model_ema \
+    --num_frames 29 \
+    --height 480 \
+    --width 640 \
+    --cache_dir "./cache_dir" \
+    --text_encoder_name google/mt5-xxl \
+    --text_prompt examples/prompt_list_0.txt \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --save_img_path "./sample_video_29x4800p_cfg7.5_step100" \
+    --fps 24 \
+    --guidance_scale 7.5 \
+    --num_sampling_steps 100 \
+    --enable_tiling \
+    --max_sequence_length 512 \
+    --sample_method EulerAncestralDiscrete \
+    --model_type "dit"
\ No newline at end of file
diff --git a/scripts/text_condition/gpu/sample_t2v_ddp.sh b/scripts/text_condition/gpu/sample_t2v_ddp.sh
new file mode 100644
index 000000000..6af17a6de
--- /dev/null
+++ b/scripts/text_condition/gpu/sample_t2v_ddp.sh
@@ -0,0 +1,21 @@
+
+torchrun --nnodes=1 --nproc_per_node 8  --master_port 29503 \
+    -m opensora.sample.sample_t2v_ddp \
+    --model_path /path/to/checkpoint-xxx/model_ema \
+    --num_frames 93 \
+    --height 720 \
+    --width 1280 \
+    --cache_dir "./cache_dir" \
+    --text_encoder_name google/mt5-xxl \
+    --text_prompt examples/prompt_list_0.txt \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --save_img_path "./sample_93x480p_cfg7.5_step100_eulera" \
+    --fps 24 \
+    --guidance_scale 7.5 \
+    --num_sampling_steps 100 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --max_sequence_length 512 \
+    --sample_method EulerAncestralDiscrete \
+    --model_type "dit"
\ No newline at end of file
diff --git a/scripts/text_condition/gpu/sample_t2v_sp.sh b/scripts/text_condition/gpu/sample_t2v_sp.sh
new file mode 100644
index 000000000..1cfb451ee
--- /dev/null
+++ b/scripts/text_condition/gpu/sample_t2v_sp.sh
@@ -0,0 +1,21 @@
+
+torchrun --nnodes=1 --nproc_per_node 8  --master_port 29503 \
+    -m opensora.sample.sample_t2v_sp \
+    --model_path /path/to/checkpoint-xxx/model_ema \
+    --num_frames 93 \
+    --height 720 \
+    --width 1280 \
+    --cache_dir "./cache_dir" \
+    --text_encoder_name google/mt5-xxl \
+    --text_prompt examples/prompt_list_0.txt \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --save_img_path "./sample_video_test" \
+    --fps 24 \
+    --guidance_scale 7.5 \
+    --num_sampling_steps 100 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --max_sequence_length 512 \
+    --sample_method EulerAncestralDiscrete \
+    --model_type "dit" 
\ No newline at end of file
diff --git a/scripts/text_condition/gpu/train_inpaint.sh b/scripts/text_condition/gpu/train_inpaint.sh
new file mode 100644
index 000000000..2276883d6
--- /dev/null
+++ b/scripts/text_condition/gpu/train_inpaint.sh
@@ -0,0 +1,74 @@
+PROJECT="your project name"
+export WANDB_API_KEY="your wanb key"
+export WANDB_MODE="online"
+export ENTITY="your name"
+export PROJECT=$PROJECT
+export HF_DATASETS_OFFLINE=1 
+export TRANSFORMERS_OFFLINE=1
+export TOKENIZERS_PARALLELISM=false
+# NCCL setting
+export GLOO_SOCKET_IFNAME=bond0
+export NCCL_SOCKET_IFNAME=bond0
+export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1
+export NCCL_IB_GID_INDEX=3
+export NCCL_IB_TC=162
+export NCCL_IB_TIMEOUT=22
+export NCCL_PXN_DISABLE=0
+export NCCL_IB_QPS_PER_CONNECTION=4
+export NCCL_ALGO=Ring
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+
+# export NCCL_ALGO=Tree
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
+    opensora/train/train_inpaint.py \
+    --model OpenSoraInpaint-ROPE-L/122 \
+    --text_encoder_name google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset inpaint \
+    --data "scripts/train_data/merge_data.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/path/to/causalvideovae" \
+    --sample_rate 1 \
+    --num_frames 93 \
+    --use_image_num 0 \
+    --max_height 480 \
+    --max_width 640 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.0 \
+    --interpolation_scale_w 1.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 10 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-5 \
+    --lr_scheduler="constant" \
+    --lr_warmup_steps=0 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=200 \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --noise_offset 0.02 \
+    --i2v_ratio 0.4 \
+    --transition_ratio 0.4 \
+    --v2v_ratio 0.1 \
+    --clear_video_ratio 0.0 \
+    --default_text_ratio 0.5 \
+    --use_rope \
+    --ema_decay 0.999 \
+    --speed_factor 1.0 \
+    --group_frame \
+    --output_dir runs/$PROJECT \
+    --pretrained_transformer_model_path "/path/to/pretrained_model" \
+    # --resume_from_checkpoint="latest" \
diff --git a/scripts/text_condition/gpu/train_t2v.sh b/scripts/text_condition/gpu/train_t2v.sh
new file mode 100644
index 000000000..5fcdf0020
--- /dev/null
+++ b/scripts/text_condition/gpu/train_t2v.sh
@@ -0,0 +1,66 @@
+# export WANDB_MODE="offline"
+export ENTITY="linbin"
+export HF_DATASETS_OFFLINE=1 
+export TRANSFORMERS_OFFLINE=1
+export PDSH_RCMD_TYPE=ssh
+# NCCL setting
+export GLOO_SOCKET_IFNAME=bond0
+export NCCL_SOCKET_IFNAME=bond0
+export NCCL_IB_HCA=mlx5_10:1,mlx5_11:1,mlx5_12:1,mlx5_13:1
+export NCCL_IB_GID_INDEX=3
+export NCCL_IB_TC=162
+export NCCL_IB_TIMEOUT=22
+export NCCL_PXN_DISABLE=0
+export NCCL_IB_QPS_PER_CONNECTION=4
+export NCCL_ALGO=Ring
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+# export NCCL_ALGO=Tree
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/deepspeed_zero2_config.yaml \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name google/mt5-xxl \
+    --cache_dir "./cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "/storage/dataset/test140k" \
+    --sample_rate 1 \
+    --num_frames 29 \
+    --max_height 480 \
+    --max_width 640 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.0 \
+    --interpolation_scale_w 1.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 10 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --lr_warmup_steps=0 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=250 \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --tile_overlap_factor 0.125 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --noise_offset 0.02 \
+    --use_rope \
+    --resume_from_checkpoint="latest" \
+    --ema_decay 0.999 \
+    --enable_tiling \
+    --speed_factor 1.0 \
+    --group_frame \
+    --sp_size 8 \
+    --train_sp_batch_size 2 \
+    --output_dir="bs16x8x1_29x480p_lr1e-4_snr5_ema999_opensora122_rope_mt5xxl_speed1.0" 
\ No newline at end of file
diff --git a/scripts/text_condition/npu/sample_inpaint.sh b/scripts/text_condition/npu/sample_inpaint.sh
new file mode 100644
index 000000000..005be4739
--- /dev/null
+++ b/scripts/text_condition/npu/sample_inpaint.sh
@@ -0,0 +1,34 @@
+WEIGHT_PATH="/home/opensora/pre_weights/"
+
+export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$(pwd)"
+export MASTER_PORT=12359
+
+if [ -z "$SAMPLE_SAVE_PATH" ]; then
+  export SAMPLE_SAVE_PATH="/home/image_data/sample_videos"
+fi
+
+if [ -z "$SAMPLE_HEIGHT" ]; then
+  echo "You should set both envs of SAMPLE_HEIGHT and SAMPLE_WIDTH"
+  return
+fi
+
+torchrun --nproc_per_node=8 opensora/sample/sample_inpaint_on_npu.py \
+    --model_path /home/image_data/checkpoints/${PROJECT_NAME} \
+    --num_frames ${NUM_FRAME} \
+    --height $SAMPLE_HEIGHT \
+    --width $SAMPLE_WIDTH \
+    --cache_dir "../cache_dir" \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --text_prompt "/home/image_data/checkpoints/i2v_validation_dir/prompt.txt" \
+    --ae CausalVAEModel_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --save_img_path "${SAMPLE_SAVE_PATH}/${PROJECT_NAME}" \
+    --fps 24 \
+    --max_sequence_length 512 \
+    --guidance_scale 5.0 \
+    --num_sampling_steps 50 \
+    --sample_method PNDM \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --model_3d \
+    --validation_dir "/home/image_data/checkpoints/i2v_validation_dir"
\ No newline at end of file
diff --git a/scripts/text_condition/npu/train_image3d_480p.sh b/scripts/text_condition/npu/train_image3d_480p.sh
new file mode 100644
index 000000000..c320b0bad
--- /dev/null
+++ b/scripts/text_condition/npu/train_image3d_480p.sh
@@ -0,0 +1,54 @@
+export PROJECT=$PROJECT_NAME
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+#export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data_on_npu.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --sample_rate 1 \
+    --num_frames 1 \
+    --max_height 480 \
+    --max_width 640 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.0 \
+    --interpolation_scale_w 1.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=8 \
+    --dataloader_num_workers 20 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --seed=10 \
+    --lr_warmup_steps=500 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=2000 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.0 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --resume_from_checkpoint="latest"
+
diff --git a/scripts/text_condition/npu/train_image3d_720p.sh b/scripts/text_condition/npu/train_image3d_720p.sh
new file mode 100644
index 000000000..2562dac26
--- /dev/null
+++ b/scripts/text_condition/npu/train_image3d_720p.sh
@@ -0,0 +1,54 @@
+export PROJECT=$PROJECT_NAME
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+#export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data_on_npu.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --sample_rate 1 \
+    --num_frames 1 \
+    --max_height 720 \
+    --max_width 1280 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.5 \
+    --interpolation_scale_w 2.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 8 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --seed=10 \
+    --lr_warmup_steps=500 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=1000 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.0 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --resume_from_checkpoint="latest"
+
diff --git a/scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh b/scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh
new file mode 100644
index 000000000..b7f751312
--- /dev/null
+++ b/scripts/text_condition/npu/train_inpaint_video3d_nx480p.sh
@@ -0,0 +1,61 @@
+export WANDB_KEY="720d886d8c437c2142c88056a1eab8ef78d64a1f"
+# # export WANDB_MODE="offline"
+export ENTITY="yunyangge"
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+# export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_inpaint.py \
+    --model OpenSoraInpaint-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset i2v \
+    --ae CausalVAEModel_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --video_data "./scripts/train_data/video_data_on_npu.txt" \
+    --image_data "./scripts/train_data/image_data_on_npu.txt" \
+    --sample_rate 1 \
+    --num_frames ${NUM_FRAME} \
+    --max_height 480 \
+    --max_width 640 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.0 \
+    --interpolation_scale_w 1.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 8 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-5 \
+    --lr_scheduler="constant" \
+    --seed=42 \
+    --lr_warmup_steps=0 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=200 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --pretrained "/home/image_data/checkpoints/gyy_pretrained_f125/model_ema/diffusion_pytorch_model.safetensors" \
+    --resume_from_checkpoint="latest" \
+    --i2v_ratio 0.5 \
+    --transition_ratio 0.4 \
+    --default_text_ratio 0.5 \
+    --group_frame \
+    --speed_factor 1.1 \
diff --git a/scripts/text_condition/npu/train_video3d_nx480p.sh b/scripts/text_condition/npu/train_video3d_nx480p.sh
new file mode 100644
index 000000000..d406a1ccb
--- /dev/null
+++ b/scripts/text_condition/npu/train_video3d_nx480p.sh
@@ -0,0 +1,57 @@
+export PROJECT=$PROJECT_NAME
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+#export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data_on_npu.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --sample_rate 1 \
+    --num_frames ${NUM_FRAME} \
+    --max_height 480 \
+    --max_width 640 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.0 \
+    --interpolation_scale_w 1.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 8 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --seed=10 \
+    --lr_warmup_steps=500 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=1000 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --resume_from_checkpoint="latest" \
+    --enable_stable_fp32 \
+    --ema_decay 0.999 \
+    --speed_factor 1.0 \
+    --drop_short_ratio 1.0
diff --git a/scripts/text_condition/npu/train_video3d_nx720p.sh b/scripts/text_condition/npu/train_video3d_nx720p.sh
new file mode 100644
index 000000000..415fb5728
--- /dev/null
+++ b/scripts/text_condition/npu/train_video3d_nx720p.sh
@@ -0,0 +1,57 @@
+export PROJECT=$PROJECT_NAME
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+#export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data_on_npu.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --sample_rate 1 \
+    --num_frames ${NUM_FRAME} \
+    --max_height 720 \
+    --max_width 1280 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.5 \
+    --interpolation_scale_w 2.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 8 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --seed=10 \
+    --lr_warmup_steps=500 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=1000 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --resume_from_checkpoint="latest" \
+    --enable_stable_fp32 \
+    --ema_decay 0.999 \
+    --speed_factor 1.0 \
+    --drop_short_ratio 1.0
diff --git a/scripts/text_condition/npu/train_video3d_sp_nx480p.sh b/scripts/text_condition/npu/train_video3d_sp_nx480p.sh
new file mode 100644
index 000000000..39268b88a
--- /dev/null
+++ b/scripts/text_condition/npu/train_video3d_sp_nx480p.sh
@@ -0,0 +1,59 @@
+export PROJECT=$PROJECT_NAME
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+#export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data_on_npu.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --sample_rate 1 \
+    --num_frames ${NUM_FRAME} \
+    --max_height 480 \
+    --max_width 640 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.0 \
+    --interpolation_scale_w 1.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 8 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --seed=10 \
+    --lr_warmup_steps=500 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=1000 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --resume_from_checkpoint="latest" \
+    --enable_stable_fp32 \
+    --ema_decay 0.999 \
+    --speed_factor 1.0 \
+    --drop_short_ratio 1.0 \
+    --sp_size 8 \
+    --train_sp_batch_size 1
diff --git a/scripts/text_condition/npu/train_video3d_sp_nx720p.sh b/scripts/text_condition/npu/train_video3d_sp_nx720p.sh
new file mode 100644
index 000000000..d5b08e540
--- /dev/null
+++ b/scripts/text_condition/npu/train_video3d_sp_nx720p.sh
@@ -0,0 +1,59 @@
+export PROJECT=$PROJECT_NAME
+WEIGHT_PATH="/home/opensora/pre_weights/"
+env
+#export WANDB_MODE='offline'
+export HCCL_OP_BASE_FFTS_MODE_ENABLE=TRUE
+export HCCL_ALGO="level0:NA;level1:H-D_R"
+export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
+
+accelerate launch \
+    --config_file scripts/accelerate_configs/multi_node_example_by_deepspeed.yaml \
+    --machine_rank=${MACHINE_RANK} \
+    --main_process_ip=${MAIN_PROCESS_IP_VALUE} \
+    opensora/train/train_t2v_diffusers.py \
+    --model OpenSoraT2V-ROPE-L/122 \
+    --text_encoder_name ${WEIGHT_PATH}/google/mt5-xxl \
+    --cache_dir "../cache_dir" \
+    --dataset t2v \
+    --data "scripts/train_data/merge_data_on_npu.txt" \
+    --ae CausalVAEModel_D4_4x8x8 \
+    --ae_path "${WEIGHT_PATH}/test140k/" \
+    --sample_rate 1 \
+    --num_frames ${NUM_FRAME} \
+    --max_height 720 \
+    --max_width 1280 \
+    --interpolation_scale_t 1.0 \
+    --interpolation_scale_h 1.5 \
+    --interpolation_scale_w 2.0 \
+    --attention_mode xformers \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --dataloader_num_workers 8 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=1000000 \
+    --learning_rate=1e-4 \
+    --lr_scheduler="constant" \
+    --seed=10 \
+    --lr_warmup_steps=500 \
+    --mixed_precision="bf16" \
+    --report_to="wandb" \
+    --checkpointing_steps=1000 \
+    --output_dir="/home/image_data/checkpoints/${PROJECT}/" \
+    --allow_tf32 \
+    --model_max_length 512 \
+    --use_image_num 0 \
+    --snr_gamma 5.0 \
+    --use_ema \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --enable_tiling \
+    --tile_overlap_factor 0.125 \
+    --use_rope \
+    --noise_offset 0.02 \
+    --resume_from_checkpoint="latest" \
+    --enable_stable_fp32 \
+    --ema_decay 0.999 \
+    --speed_factor 1.0 \
+    --drop_short_ratio 1.0 \
+    --sp_size 8 \
+    --train_sp_batch_size 1
diff --git a/scripts/train_data/merge_data.txt b/scripts/train_data/merge_data.txt
new file mode 100644
index 000000000..6b5792143
--- /dev/null
+++ b/scripts/train_data/merge_data.txt
@@ -0,0 +1,2 @@
+
+/storage/dataset/panda70m,/storage/dataset/filter_aes/filter_panda70m_json/panda70m_last_6268414_flowValue_940432_5.4plus.json
\ No newline at end of file
diff --git a/scripts/train_data/merge_data_on_npu.txt b/scripts/train_data/merge_data_on_npu.txt
new file mode 100644
index 000000000..d7e10fc05
--- /dev/null
+++ b/scripts/train_data/merge_data_on_npu.txt
@@ -0,0 +1 @@
+/home/local_dataset_6t/video_data_obs/movie_aes_motion,/home/opensora/captions/movie_aes_motion/bbc01_clips_aes_motion_final_250508.json
\ No newline at end of file
diff --git a/tools/get_img_info.py b/tools/get_img_info.py
new file mode 100644
index 000000000..0617f02b1
--- /dev/null
+++ b/tools/get_img_info.py
@@ -0,0 +1,51 @@
+
+import cv2
+from tqdm import tqdm
+from glob import glob
+import json
+import os
+from multiprocessing import Pool
+import argparse
+
+def get_image_size(image_path):
+    """
+    Given an image path, return its width and height.
+    """
+    try:
+        image = cv2.imread(image_path)
+        height, width = image.shape[:2]
+        return image_path, height, width
+    except Exception as e:
+        return image_path, None, None
+
+def process_image_paths(image_paths):
+    items = []
+    for image_path, height, width in image_paths:
+        path = image_path.replace(image_root if image_root.endswith('/') else image_root + '/', '')
+        item = dict(path=path, resolution=dict(height=height, width=width))
+        items.append(item)
+    return items
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_root", type=str, required=True)
+    parser.add_argument("--save_root", type=str, required=True)
+    parser.add_argument("--save_name", type=str, required=True)
+    args = parser.parse_args()
+    # image_root = '/storage/dataset/xx/xx_v1'
+    # save_root = '/storage/dataset/image'
+    image_root = args.image_root
+    save_root = args.save_root
+    os.makedirs(save_root, exist_ok=True)
+    save_name = '{}_{}_resolution.json'
+    all_paths = glob(os.path.join(image_root, '**', '*.jpg'), recursive=True)
+
+    num_processes = os.cpu_count()  # Use the number of CPU cores
+    num_processes = 128  # Use the number of CPU cores
+    with Pool(num_processes) as pool:
+        results = list(tqdm(pool.imap(get_image_size, all_paths), total=len(all_paths)))
+
+    items = process_image_paths(results)
+
+    with open(os.path.join(save_root, save_name.format(args.save_name, len(items))), 'w') as f:
+        json.dump(items, f, indent=2)
diff --git a/tools/merge_imginfo_to_anno.py b/tools/merge_imginfo_to_anno.py
new file mode 100644
index 000000000..32d3a1220
--- /dev/null
+++ b/tools/merge_imginfo_to_anno.py
@@ -0,0 +1,42 @@
+
+from tqdm import tqdm
+from glob import glob
+import json
+import os
+
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--anno_path", type=str, required=True)
+    parser.add_argument("--img_info_path", type=str, required=True)
+    parser.add_argument("--save_root", type=str, required=True)
+    parser.add_argument("--save_name", type=str, required=True)
+    args = parser.parse_args()
+
+    # anno_path = '/storage/anno_jsons/sam_1940032.json'
+    # img_info_path = '/storage/dataset/image/sam_2151074_resolution.json'
+    # save_root = '/storage/anno_jsons'
+    # save_name = 'sam_{}_resolution.json'
+    
+    anno_path = args.anno_path
+    img_info_path = args.img_info_path
+    save_root = args.save_root
+    save_name = '{}_{}_resolution.json'
+
+
+    with open(anno_path, 'r') as f:
+        anno = json.load(f)
+    with open(img_info_path, 'r') as f:
+        img_info = json.load(f)
+    img_info = {i['path']: i['resolution'] for i in img_info}
+
+    items = []
+    cnt = 0
+    for i in tqdm(anno):
+        resolution = img_info[i['path']]
+        i['resolution'] = resolution
+        items.append(i)
+    with open(os.path.join(save_root, save_name.format(args.save_name, len(items))), 'w') as f:
+        json.dump(items, f, indent=2)
\ No newline at end of file