diff --git a/README.md b/README.md
index 386dea7b1..8fa42ef8c 100644
--- a/README.md
+++ b/README.md
@@ -9,11 +9,9 @@
 <img src=assets/logo.png width="50%"/>
 </div>
 
-As scaling laws from language models are applied to diffusion transformers, the number of parameters in diffusion models has grown significantly. This trend is even more pronounced in video models, where people are scaling not just the number of parameters but also the sequence length. As a result, traditional post-training workflows for diffusion models, such as fine-tuning, distillation, and inference, are becoming increasingly difficult to manage with frameworks like HF Diffusers, which are primarily designed for simple data-parallel workloads.
+FastVideo is a scalable framework for post-training video diffusion models, addressing the growing challenges of fine-tuning, distillation, and inference as model sizes and sequence lengths increase. As a first step, it provides an efficient script for distilling and fine-tuning the 10B Mochi model, with plans to expand features and support for more models.
 
-That is why we launched this FastVideo project to try to build a scalable framework for post-training various video diffusion models. As the tiny first step, we now provide a simple and efficient script to distill and finetune the 10B Mochi model. We will continue to add more features and models to this project in the future.
-
-### Key Features
+### Features
 
 - FastMochi, a distilled Mochi model that can generate videos with merely 8 sampling steps.
 - Finetuning with FSDP (both master weight and ema weight), sequence parallelism, and selective gradient checkpointing.
@@ -56,7 +54,7 @@ Jump to a specific section:
 conda create -n fastmochi python=3.10.0 -y && conda activate fastmochi
 pip3 install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu121
 pip install packaging ninja && pip install flash-attn==2.7.0.post2 --no-build-isolation 
-"git+https://github.com/huggingface/diffusers.git@bf64b32652a63a1865a0528a73a13652b201698b"
+pip install "git+https://github.com/huggingface/diffusers.git@bf64b32652a63a1865a0528a73a13652b201698b"
 git clone https://github.com/hao-ai-lab/FastVideo.git
 cd FastVideo && pip install -e .
 ```
@@ -74,8 +72,9 @@ python scripts/download_hf.py --repo_id=FastVideo/FastMochi --local_dir=data/Fas
 
 Start the gradio UI with
 ```
-python3 ./demos/gradio_ui.py --model_dir weights/ --cpu_offload
+python fastvideo/demo/gradio_web_demo.py --model_path data/FastMochi
 ```
+
 We also provide CLI inference script featured with sequence parallelism.
 
 ```
@@ -83,16 +82,15 @@ export NUM_GPUS=4
 
 torchrun --nnodes=1 --nproc_per_node=$NUM_GPUS \
     fastvideo/sample/sample_t2v_mochi.py \
-    --model_path data/mochi \
+    --model_path data/FastMochi \
     --prompt_path assets/prompt.txt \
     --num_frames 163 \
     --height 480 \
     --width 848 \
     --num_inference_steps 8 \
-    --guidance_scale 4.5 \
+    --guidance_scale 1.5 \
     --output_path outputs_video/demo_video \
-    --shift 8 \
-    --seed 42 \
+    --seed 12345 \
     --scheduler_type "pcm_linear_quadratic" \
     --linear_threshold 0.1 \
     --linear_range 0.75
diff --git a/fastvideo/dataset/__init__.py b/fastvideo/dataset/__init__.py
index e33340b55..f28fc403f 100644
--- a/fastvideo/dataset/__init__.py
+++ b/fastvideo/dataset/__init__.py
@@ -4,31 +4,51 @@ from torchvision import transforms
 from torchvision.transforms import Lambda
 from fastvideo.dataset.t2v_datasets import T2V_dataset
 from fastvideo.dataset.latent_datasets import LatentDataset
-from fastvideo.dataset.transform import Normalize255, TemporalRandomCrop,CenterCropResizeVideo
+from fastvideo.dataset.transform import (
+    Normalize255,
+    TemporalRandomCrop,
+    CenterCropResizeVideo,
+)
+
 
 def getdataset(args):
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
-    norm_fun = Lambda(lambda x: 2. * x - 1.)
-    resize_topcrop = [CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True), ]
-    resize = [CenterCropResizeVideo((args.max_height, args.max_width)), ]
-    transform = transforms.Compose([
-        # Normalize255(),
-        *resize, 
-        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
-        # norm_fun
-    ])
-    transform_topcrop = transforms.Compose([
-        Normalize255(),
-        *resize_topcrop, 
-        # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
-        norm_fun
-    ])
+    norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
+    resize_topcrop = [
+        CenterCropResizeVideo((args.max_height, args.max_width), top_crop=True),
+    ]
+    resize = [
+        CenterCropResizeVideo((args.max_height, args.max_width)),
+    ]
+    transform = transforms.Compose(
+        [
+            # Normalize255(),
+            *resize,
+            # RandomHorizontalFlipVideo(p=0.5),  # in case their caption have position decription
+            # norm_fun
+        ]
+    )
+    transform_topcrop = transforms.Compose(
+        [
+            Normalize255(),
+            *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)
-    
+    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,
+        )
+
     raise NotImplementedError(args.dataset)
 
 
@@ -37,32 +57,34 @@ if __name__ == "__main__":
     from fastvideo.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, 
-        'text_max_length': 300, 
-        'max_height': 320,
-        'max_width': 240,
-        'num_frames': 1,
-        'use_image_num': 0, 
-        '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,
 
-    }
+    args = type(
+        "args",
+        (),
+        {
+            "ae": "CausalVAEModel_4x8x8",
+            "dataset": "t2v",
+            "attention_mode": "xformers",
+            "use_rope": True,
+            "text_max_length": 300,
+            "max_height": 320,
+            "max_width": 240,
+            "num_frames": 1,
+            "use_image_num": 0,
+            "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)
@@ -70,7 +92,9 @@ if __name__ == "__main__":
     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]
+        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:
@@ -81,5 +105,7 @@ if __name__ == "__main__":
             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
+    import ipdb
+
+    ipdb.set_trace()
+    print("end")
diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py
index c00d711df..c4143a17f 100644
--- a/fastvideo/dataset/latent_datasets.py
+++ b/fastvideo/dataset/latent_datasets.py
@@ -4,13 +4,14 @@ import json
 import os
 import random
 
+
 class LatentDataset(Dataset):
     def __init__(
-        self, 
-        json_path, 
-        num_latent_t, 
+        self,
+        json_path,
+        num_latent_t,
         cfg_rate,
-        ):
+    ):
         # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
         self.json_path = json_path
         self.cfg_rate = cfg_rate
@@ -18,8 +19,10 @@ class LatentDataset(Dataset):
         self.video_dir = os.path.join(self.datase_dir_path, "video")
         self.latent_dir = os.path.join(self.datase_dir_path, "latent")
         self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed")
-        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
-        with open(self.json_path, 'r') as f:
+        self.prompt_attention_mask_dir = os.path.join(
+            self.datase_dir_path, "prompt_attention_mask"
+        )
+        with open(self.json_path, "r") as f:
             self.data_anno = json.load(f)
         # json.load(f) already keeps the order
         # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
@@ -28,25 +31,44 @@ class LatentDataset(Dataset):
         self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
         # 256 zeros
         self.uncond_prompt_mask = torch.zeros(256).bool()
-        self.lengths = [data_item['length'] if "length" in data_item else 1 for data_item in self.data_anno]
+        self.lengths = [
+            data_item["length"] if "length" in data_item else 1
+            for data_item in self.data_anno
+        ]
+
     def __getitem__(self, idx):
         latent_file = self.data_anno[idx]["latent_path"]
         prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
         prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"]
-        # load 
-        latent = torch.load(os.path.join(self.latent_dir, latent_file), map_location="cpu", weights_only=True)       
-        latent = latent.squeeze(0)[:, -self.num_latent_t:]
+        # load
+        latent = torch.load(
+            os.path.join(self.latent_dir, latent_file),
+            map_location="cpu",
+            weights_only=True,
+        )
+        latent = latent.squeeze(0)[:, -self.num_latent_t :]
         if random.random() < self.cfg_rate:
             prompt_embed = self.uncond_prompt_embed
             prompt_attention_mask = self.uncond_prompt_mask
         else:
-            prompt_embed = torch.load(os.path.join(self.prompt_embed_dir, prompt_embed_file), map_location="cpu", weights_only=True)
-            prompt_attention_mask = torch.load(os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file), map_location="cpu", weights_only=True)
+            prompt_embed = torch.load(
+                os.path.join(self.prompt_embed_dir, prompt_embed_file),
+                map_location="cpu",
+                weights_only=True,
+            )
+            prompt_attention_mask = torch.load(
+                os.path.join(
+                    self.prompt_attention_mask_dir, prompt_attention_mask_file
+                ),
+                map_location="cpu",
+                weights_only=True,
+            )
         return latent, prompt_embed, prompt_attention_mask
-    
+
     def __len__(self):
         return len(self.data_anno)
-    
+
+
 def latent_collate_function(batch):
     # return latent, prompt, latent_attn_mask, text_attn_mask
     # latent_attn_mask: # b t h w
@@ -57,25 +79,48 @@ def latent_collate_function(batch):
     max_t = max([latent.shape[1] for latent in latents])
     max_h = max([latent.shape[2] for latent in latents])
     max_w = max([latent.shape[3] for latent in latents])
-    
+
     # padding
-    latents = [torch.nn.functional.pad(latent, (0, max_t - latent.shape[1], 0, max_h - latent.shape[2], 0, max_w - latent.shape[3])) for latent in latents]
+    latents = [
+        torch.nn.functional.pad(
+            latent,
+            (
+                0,
+                max_t - latent.shape[1],
+                0,
+                max_h - latent.shape[2],
+                0,
+                max_w - latent.shape[3],
+            ),
+        )
+        for latent in latents
+    ]
     # attn mask
     latent_attn_mask = torch.ones(len(latents), max_t, max_h, max_w)
     # set to 0 if padding
     for i, latent in enumerate(latents):
-        latent_attn_mask[i, latent.shape[1]:, :, :] = 0
-        latent_attn_mask[i, :, latent.shape[2]:, :] = 0
-        latent_attn_mask[i, :, :, latent.shape[3]:] = 0
-    
+        latent_attn_mask[i, latent.shape[1] :, :, :] = 0
+        latent_attn_mask[i, :, latent.shape[2] :, :] = 0
+        latent_attn_mask[i, :, :, latent.shape[3] :] = 0
+
     prompt_embeds = torch.stack(prompt_embeds, dim=0)
     prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0)
     latents = torch.stack(latents, dim=0)
     return latents, prompt_embeds, latent_attn_mask, prompt_attention_masks
 
+
 if __name__ == "__main__":
     dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
-    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
+    dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function
+    )
     for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
-        print(latent.shape, prompt_embed.shape, latent_attn_mask.shape, prompt_attention_mask.shape)
-        import pdb; pdb.set_trace()
\ No newline at end of file
+        print(
+            latent.shape,
+            prompt_embed.shape,
+            latent_attn_mask.shape,
+            prompt_attention_mask.shape,
+        )
+        import pdb
+
+        pdb.set_trace()
diff --git a/fastvideo/dataset/t2v_datasets.py b/fastvideo/dataset/t2v_datasets.py
index d1bdbf949..6ab7f030e 100644
--- a/fastvideo/dataset/t2v_datasets.py
+++ b/fastvideo/dataset/t2v_datasets.py
@@ -14,15 +14,10 @@ from PIL import Image
 from accelerate.logging import get_logger
 from fastvideo.utils.dataset_utils import DecordInit
 import torchvision
+
 logger = get_logger(__name__)
 
 
-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
-
 class SingletonMeta(type):
     _instances = {}
 
@@ -55,7 +50,7 @@ class DataSetProg(metaclass=SingletonMeta):
             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]
+            self.worker_elements[i] = self.elements[start:end]
 
     def get_item(self, work_info):
         if work_info is None:
@@ -63,18 +58,20 @@ class DataSetProg(metaclass=SingletonMeta):
         else:
             worker_id = work_info.id
 
-        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_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 filter_resolution(h, w, max_h_div_w_ratio=17/16, min_h_div_w_ratio=8 / 16):
+
+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):
@@ -97,11 +94,11 @@ class T2V_dataset(Dataset):
         self.v_decoder = DecordInit()
         self.video_length_tolerance_range = args.video_length_tolerance_range
         self.support_Chinese = True
-        if not ('mt5' in args.text_encoder_name):
+        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
@@ -123,35 +120,39 @@ class T2V_dataset(Dataset):
             data = self.get_data(idx)
             return data
         except Exception as e:
-            logger.info(f'Error with {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]}")
             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'):
+        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):
-        video_path = dataset_prog.cap_list[idx]['path']
+        video_path = dataset_prog.cap_list[idx]["path"]
         assert os.path.exists(video_path), f"file {video_path} do not exist!"
-        frame_indices = dataset_prog.cap_list[idx]['sample_frame_index']
-        torchvision_video, _, metadata =  torchvision.io.read_video(video_path, output_format="TCHW")
+        frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
+        torchvision_video, _, metadata = torchvision.io.read_video(
+            video_path, output_format="TCHW"
+        )
         video = torchvision_video[frame_indices]
-        video = self.transform(video) 
-        video = rearrange(video, 't c h w -> c t h w')
-        video =  video.to(torch.uint8)
+        video = self.transform(video)
+        video = rearrange(video, "t c h w -> c t h w")
+        video = video.to(torch.uint8)
         assert video.dtype == torch.uint8
 
         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}'
-        
+        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}"
+
         video = video.float() / 127.5 - 1.0
-        
-        text = dataset_prog.cap_list[idx]['cap']
+
+        text = dataset_prog.cap_list[idx]["cap"]
         if not isinstance(text, list):
             text = [text]
         text = [random.choice(text)]
@@ -160,51 +161,70 @@ class T2V_dataset(Dataset):
         text_tokens_and_mask = self.tokenizer(
             text,
             max_length=self.text_max_length,
-            padding='max_length',
+            padding="max_length",
             truncation=True,
             return_attention_mask=True,
             add_special_tokens=True,
-            return_tensors='pt'
+            return_tensors="pt",
+        )
+        input_ids = text_tokens_and_mask["input_ids"]
+        cond_mask = text_tokens_and_mask["attention_mask"]
+        return dict(
+            pixel_values=video,
+            text=text,
+            input_ids=input_ids,
+            cond_mask=cond_mask,
+            path=video_path,
         )
-        input_ids = text_tokens_and_mask['input_ids']
-        cond_mask = text_tokens_and_mask['attention_mask']
-        return dict(pixel_values=video, text=text, input_ids=input_ids, cond_mask=cond_mask, path=video_path)
 
     def get_image(self, idx):
         image_data = dataset_prog.cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
 
-        image = Image.open(image_data['path']).convert('RGB')  # [h, w, c]
+        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]
+        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 = (
+            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 = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-        
+
         image = image.float() / 127.5 - 1.0
-        
-        caps = image_data['cap'] if isinstance(image_data['cap'], list) else [image_data['cap']]
+
+        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)
+        text = caps
         input_ids, cond_mask = [], []
         text = text if random.random() > self.cfg else ""
         text_tokens_and_mask = self.tokenizer(
             text,
             max_length=self.text_max_length,
-            padding='max_length',
+            padding="max_length",
             truncation=True,
             return_attention_mask=True,
             add_special_tokens=True,
-            return_tensors='pt'
+            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,
+            text=text,
+            input_ids=input_ids,
+            cond_mask=cond_mask,
+            path=image_data["path"],
         )
-        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, text=text, input_ids=input_ids, cond_mask=cond_mask, path=image_data['path'])
 
     def define_frame_index(self, cap_list):
-        
         new_cap_list = []
         sample_num_frames = []
         cnt_too_long = 0
@@ -215,78 +235,100 @@ class T2V_dataset(Dataset):
         cnt_movie = 0
         cnt_img = 0
         for i in cap_list:
-            path = i['path']
-            cap = i.get('cap', None)
+            path = i["path"]
+            cap = i.get("cap", None)
             # ======no caption=====
             if cap is None:
                 cnt_no_cap += 1
                 continue
-            if path.endswith('.mp4'):
+            if path.endswith(".mp4"):
                 # ======no fps and duration=====
-                duration = i.get('duration', None)
-                fps = i.get('fps', None)
+                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)
+                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:
+                    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']
+                    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)
+                    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:
                         print("resolution mismatch")
                         cnt_resolution_mismatch += 1
                         continue
 
                 # import ipdb;ipdb.set_trace()
-                i['num_frames'] = math.ceil(fps * duration)
-                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration. 
-                if i['num_frames'] / fps > self.video_length_tolerance_range * (self.num_frames / self.train_fps * self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
+                i["num_frames"] = math.ceil(fps * duration)
+                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
+                if (
+                    i["num_frames"] / fps
+                    > self.video_length_tolerance_range
+                    * (self.num_frames / self.train_fps * self.speed_factor)
+                ):  # too long video is not suitable for this training stage (self.num_frames)
                     cnt_too_long += 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)
+                start_frame_idx = 0
+                frame_indices = np.arange(
+                    start_frame_idx, i["num_frames"], frame_interval
+                ).astype(int)
 
                 # comment out it to enable dynamic frames training
-                if len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio:
+                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[begin_index:end_index]
                     # frame_indices = frame_indices[:self.num_frames]  # head crop
-                i['sample_frame_index'] = frame_indices.tolist()
+                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
+                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'])
+                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")
+                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)}')
+        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, frame_indices):
         decord_vr = self.v_decoder(path)
         video_data = decord_vr.get_batch(frame_indices).asnumpy()
@@ -296,19 +338,20 @@ class T2V_dataset(Dataset):
 
     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]
+        with open(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:
+            with open(anno, "r") as f:
                 sub_list = json.load(f)
-            logger.info(f'Building {anno}...')
+            logger.info(f"Building {anno}...")
             for i in range(len(sub_list)):
-                sub_list[i]['path'] = opj(folder, sub_list[i]['path'])
+                sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
             cap_lists += sub_list
         return cap_lists
 
-
     def get_cap_list(self):
         cap_lists = self.read_jsons(self.data)
         return cap_lists
diff --git a/fastvideo/dataset/transform.py b/fastvideo/dataset/transform.py
index df19d603a..d18f0e9c6 100644
--- a/fastvideo/dataset/transform.py
+++ b/fastvideo/dataset/transform.py
@@ -32,7 +32,9 @@ def center_crop_arr(pil_image, image_size):
     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])
+    return Image.fromarray(
+        arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
+    )
 
 
 def crop(clip, i, j, h, w):
@@ -42,21 +44,37 @@ def crop(clip, i, j, h, w):
     """
     if len(clip.size()) != 4:
         raise ValueError("clip should be a 4D tensor")
-    return clip[..., i: i + h, j: j + w]
+    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)
+        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}")
+        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)
+    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"):
@@ -107,11 +125,10 @@ def center_crop_using_short_edge(clip):
     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
@@ -121,15 +138,16 @@ def center_crop_th_tw(clip, th, tw, top_crop):
     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)
@@ -159,7 +177,9 @@ def normalize_video(clip):
     """
     _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))
+        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
 
@@ -219,7 +239,9 @@ class RandomCropVideo:
         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)}")
+            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
@@ -235,7 +257,7 @@ class RandomCropVideo:
 
 class SpatialStrideCropVideo:
     def __init__(self, stride):
-            self.stride = stride
+        self.stride = stride
 
     def __call__(self, clip):
         """
@@ -258,17 +280,18 @@ class SpatialStrideCropVideo:
     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,
+        skip_low_resolution=False,
+        interpolation_mode="bilinear",
     ):
         self.size = size
         self.skip_low_resolution = skip_low_resolution
@@ -291,27 +314,31 @@ class LongSideResizeVideo:
         else:
             h = int(h * self.size / w)
             w = self.size
-        resize_clip = resize(clip, target_size=(h, w),
-                                         interpolation_mode=self.interpolation_mode)
+        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",
+        self,
+        size,
+        top_crop=False,
+        interpolation_mode="bilinear",
     ):
         if len(size) != 2:
-            raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            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
@@ -325,10 +352,15 @@ class CenterCropResizeVideo:
                 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)
+        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)
+        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:
@@ -336,19 +368,21 @@ class CenterCropResizeVideo:
 
 
 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",
+        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}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}"
+                )
             self.size = size
         else:
             self.size = (size, size)
@@ -363,7 +397,9 @@ class UCFCenterCropVideo:
             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_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
 
@@ -372,18 +408,20 @@ class UCFCenterCropVideo:
 
 
 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",
+        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}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}"
+                )
             self.size = size
         else:
             self.size = (size, size)
@@ -398,13 +436,15 @@ class KineticsRandomCropResizeVideo:
 
 class CenterCropVideo:
     def __init__(
-            self,
-            size,
-            interpolation_mode="bilinear",
+        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}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}"
+                )
             self.size = size
         else:
             self.size = (size, size)
@@ -516,6 +556,7 @@ class TemporalRandomCrop(object):
         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.
 
@@ -530,13 +571,16 @@ class DynamicSampleDuration(object):
     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_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__':
+
+if __name__ == "__main__":
     from torchvision import transforms
     import torchvision.io as io
     import numpy as np
@@ -544,18 +588,20 @@ if __name__ == '__main__':
     import os
 
     vframes, aframes, info = io.read_video(
-        filename='./v_Archery_g01_c03.avi',
-        pts_unit='sec',
-        output_format='TCHW'
+        filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW"
     )
 
-    trans = transforms.Compose([
-        Normalize255(),
-        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)
-    ])
+    trans = transforms.Compose(
+        [
+            Normalize255(),
+            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
@@ -569,7 +615,9 @@ if __name__ == '__main__':
     # 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)
+    frame_indice = np.linspace(
+        start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int
+    )
     print(frame_indice)
 
     select_vframes = vframes[frame_indice]
@@ -580,12 +628,18 @@ if __name__ == '__main__':
     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)
+    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)
+    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))
\ No newline at end of file
+        save_image(
+            select_vframes_trans[i],
+            os.path.join("./test000", "%04d.png" % i),
+            normalize=True,
+            value_range=(-1, 1),
+        )
diff --git a/fastvideo/demo/gradio_web_demo.py b/fastvideo/demo/gradio_web_demo.py
index 889346fe0..bb4bd810d 100644
--- a/fastvideo/demo/gradio_web_demo.py
+++ b/fastvideo/demo/gradio_web_demo.py
@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
-from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
-from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
+from fastvideo.model.pipeline_mochi import MochiPipeline
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import export_to_video
 from fastvideo.distill.solver import PCMFMScheduler
diff --git a/fastvideo/distill.py b/fastvideo/distill.py
index 7f496ebcb..32f30909f 100644
--- a/fastvideo/distill.py
+++ b/fastvideo/distill.py
@@ -1,15 +1,19 @@
 import argparse
 import math
 import os
-from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
-    destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info
+from fastvideo.utils.parallel_states import (
+    initialize_sequence_parallel_state,
+    destroy_sequence_parallel_group,
+    get_sequence_parallel_state,
+    nccl_info,
+)
 from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
 from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_mochi_dit_input
 from fastvideo.utils.validation import log_validation
 import time
 from torch.utils.data import DataLoader
 import torch
-from torch.distributed.fsdp import  ShardingStrategy
+from torch.distributed.fsdp import ShardingStrategy
 from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
     StateDictType,
@@ -38,28 +42,35 @@ from peft import LoraConfig
 from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
 )
-from fastvideo.utils.checkpoint import save_checkpoint, save_lora_checkpoint, resume_lora_optimizer
+from fastvideo.utils.checkpoint import (
+    save_checkpoint,
+    save_lora_checkpoint,
+    resume_lora_optimizer,
+)
+
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.31.0")
 import time
 from collections import deque
 
-            
 def main_print(content):
-    if int(os.environ['LOCAL_RANK']) <= 0: 
+    if int(os.environ["LOCAL_RANK"]) <= 0:
         print(content)
 
+
 def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step):
     main_print(f"--> saving checkpoint at step {step}")
     with FSDP.state_dict_type(
-        transformer, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        transformer,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
     ):
         cpu_state = transformer.state_dict()
-    #todo move to get_state_dict
+    # todo move to get_state_dict
     if rank <= 0:
         save_dir = os.path.join(output_dir, f"checkpoint-{step}")
         os.makedirs(save_dir, exist_ok=True)
-        # save using safetensors 
+        # save using safetensors
         weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
         config_dict = dict(transformer.config)
@@ -67,33 +78,71 @@ def save_checkpoint(transformer: MochiTransformer3DModel, rank, output_dir, step
         # save dict as json
         with open(config_path, "w") as f:
             json.dump(config_dict, f, indent=4)
-    main_print(f"--> checkpoint saved at step {step}")    
-      
+    main_print(f"--> checkpoint saved at step {step}")
+
 
 def reshard_fsdp(model):
     for m in FSDP.fsdp_modules(model):
         if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
             torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
-            
+
+
 def get_norm(model_pred, norms, gradient_accumulation_steps):
-    fro_norm = torch.linalg.matrix_norm(model_pred, ord="fro") / gradient_accumulation_steps
-    largest_singular_value = torch.linalg.matrix_norm(model_pred, ord=2) / gradient_accumulation_steps
+    fro_norm = (
+        torch.linalg.matrix_norm(model_pred, ord="fro") / gradient_accumulation_steps
+    )
+    largest_singular_value = (
+        torch.linalg.matrix_norm(model_pred, ord=2) / gradient_accumulation_steps
+    )
     absolute_mean = torch.mean(torch.abs(model_pred)) / gradient_accumulation_steps
     absolute_max = torch.max(torch.abs(model_pred)) / gradient_accumulation_steps
     dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
     dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
     dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
-    norms["fro"] += torch.mean(fro_norm).item()     
+    norms["fro"] += torch.mean(fro_norm).item()
     norms["largest singular value"] += torch.mean(largest_singular_value).item()
     norms["absolute mean"] += absolute_mean.item()
     norms["absolute max"] += absolute_max.item()
-    
-def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, optimizer, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg, ema_decay, pred_decay_weight, pred_decay_type):
+
+
+def train_one_step_mochi(
+    transformer,
+    teacher_transformer,
+    ema_transformer,
+    optimizer,
+    lr_scheduler,
+    loader,
+    noise_scheduler,
+    solver,
+    noise_random_generator,
+    gradient_accumulation_steps,
+    sp_size,
+    max_grad_norm,
+    uncond_prompt_embed,
+    uncond_prompt_mask,
+    num_euler_timesteps,
+    multiphase,
+    not_apply_cfg_solver,
+    distill_cfg,
+    ema_decay,
+    pred_decay_weight,
+    pred_decay_type,
+):
     total_loss = 0.0
     optimizer.zero_grad()
-    model_pred_norm = {"fro": 0.0, "largest singular value": 0.0, "absolute mean": 0.0, "absolute max": 0.0}
+    model_pred_norm = {
+        "fro": 0.0,
+        "largest singular value": 0.0,
+        "absolute mean": 0.0,
+        "absolute max": 0.0,
+    }
     for _ in range(gradient_accumulation_steps):
-        latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
+        (
+            latents,
+            encoder_hidden_states,
+            latents_attention_mask,
+            encoder_attention_mask,
+        ) = next(loader)
         model_input = normalize_mochi_dit_input(latents)
         noise = torch.randn_like(model_input)
         bsz = model_input.shape[0]
@@ -105,13 +154,9 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti
         # Add noise according to flow matching.
         # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
         sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
-        sigmas_prev = extract_into_tensor(
-            solver.sigmas_prev, index, model_input.shape
-        )
+        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index, model_input.shape)
 
-        timesteps = (
-            sigmas * noise_scheduler.config.num_train_timesteps
-        ).view(-1)
+        timesteps = (sigmas * noise_scheduler.config.num_train_timesteps).view(-1)
         # if squeeze to [], unsqueeze to [1]
 
         timesteps_prev = (
@@ -119,15 +164,14 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti
         ).view(-1)
         noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
 
-        # Predict the noise residual 
+        # Predict the noise residual
         with torch.autocast("cuda", dtype=torch.bfloat16):
-
             model_pred = transformer(
                 noisy_model_input,
                 encoder_hidden_states,
                 timesteps,
-                encoder_attention_mask, # B, L
-                return_dict= False
+                encoder_attention_mask,  # B, L
+                return_dict=False,
             )[0]
 
         # if accelerator.is_main_process:
@@ -141,8 +185,8 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti
                     noisy_model_input,
                     encoder_hidden_states,
                     timesteps,
-                    encoder_attention_mask, # B, L
-                    return_dict= False
+                    encoder_attention_mask,  # B, L
+                    return_dict=False,
                 )[0].float()
             if not_apply_cfg_solver:
                 uncond_teacher_output = cond_teacher_output
@@ -154,14 +198,12 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti
                         uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
                         timesteps,
                         uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
-                        return_dict= False
+                        return_dict=False,
                     )[0].float()
             teacher_output = cond_teacher_output + w * (
                 cond_teacher_output - uncond_teacher_output
             )
-            x_prev = solver.euler_step(
-                noisy_model_input, teacher_output, index
-            )
+            x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
 
         # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
         with torch.no_grad():
@@ -171,74 +213,85 @@ def train_one_step_mochi(transformer, teacher_transformer, ema_transformer, opti
                         x_prev.float(),
                         encoder_hidden_states,
                         timesteps_prev,
-                        encoder_attention_mask, # B, L
-                        return_dict= False
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
                     )[0]
                 else:
                     target_pred = transformer(
                         x_prev.float(),
                         encoder_hidden_states,
                         timesteps_prev,
-                        encoder_attention_mask, # B, L
-                        return_dict= False
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
                     )[0]
 
             target, end_index = solver.euler_style_multiphase_pred(
                 x_prev, target_pred, index, multiphase, True
             )
 
-
-        huber_c = 0.001 
+        huber_c = 0.001
         # loss = loss.mean()
-        loss = torch.mean(
-            torch.sqrt(
-                (model_pred.float() - target.float()) ** 2 + huber_c**2
+        loss = (
+            torch.mean(
+                torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c**2)
+                - huber_c
             )
-            - huber_c
-        ) / gradient_accumulation_steps
+            / gradient_accumulation_steps
+        )
         if pred_decay_weight > 0:
             if pred_decay_type == "l1":
-                pred_decay_loss = torch.mean(torch.sqrt(model_pred.float() ** 2 )) * pred_decay_weight / gradient_accumulation_steps
+                pred_decay_loss = (
+                    torch.mean(torch.sqrt(model_pred.float() ** 2))
+                    * pred_decay_weight
+                    / gradient_accumulation_steps
+                )
                 loss += pred_decay_loss
             elif pred_decay_type == "l2":
                 # essnetially k2?
-                pred_decay_loss = torch.mean(model_pred.float() ** 2 ) * pred_decay_weight / gradient_accumulation_steps
+                pred_decay_loss = (
+                    torch.mean(model_pred.float() ** 2)
+                    * pred_decay_weight
+                    / gradient_accumulation_steps
+                )
                 loss += pred_decay_loss
             else:
                 assert NotImplementedError("pred_decay_type is not implemented")
 
         # calculate model_pred norm and mean
-        get_norm(model_pred.detach().float(), model_pred_norm, gradient_accumulation_steps)
+        get_norm(
+            model_pred.detach().float(), model_pred_norm, gradient_accumulation_steps
+        )
         loss.backward()
-        
+
         avg_loss = loss.detach().clone()
         dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
         dist.all_reduce(pred_decay_loss.detach(), op=dist.ReduceOp.AVG)
-        total_loss += avg_loss.item() 
-        
+        total_loss += avg_loss.item()
 
-    # update ema 
+    # update ema
     if ema_transformer is not None:
         reshard_fsdp(ema_transformer)
-        for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()):
+        for p_averaged, p_model in zip(
+            ema_transformer.parameters(), transformer.parameters()
+        ):
             with torch.no_grad():
-                p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay))
-            
+                p_averaged.copy_(
+                    torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay)
+                )
+
     grad_norm = transformer.clip_grad_norm_(max_grad_norm)
     optimizer.step()
     lr_scheduler.step()
-    
-    
+
     return total_loss, grad_norm.item(), model_pred_norm, pred_decay_loss.item()
 
 
 def main(args):
-
     torch.backends.cuda.matmul.allow_tf32 = True
-    
-    local_rank = int(os.environ['LOCAL_RANK'])
-    rank = int(os.environ['RANK'])
-    world_size = int(os.environ['WORLD_SIZE'])
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
     dist.init_process_group("nccl")
     torch.cuda.set_device(local_rank)
     device = torch.cuda.current_device()
@@ -252,28 +305,28 @@ def main(args):
     noise_random_generator = None
 
     # Handle the repository creation
-    if rank <=0 and args.output_dir is not None:
+    if rank <= 0 and 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.
 
     # Create model:
-    
+
     main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
     # keep the master weight to float32
     if args.dit_model_name_or_path:
         transformer = transformer = MochiTransformer3DModel.from_pretrained(
             args.dit_model_name_or_path,
-            torch_dtype = torch.float32,
-            #torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
+            torch_dtype=torch.float32,
+            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
         )
     else:
         transformer = MochiTransformer3DModel.from_pretrained(
             args.pretrained_model_name_or_path,
             subfolder="transformer",
-            torch_dtype = torch.float32,
-            #torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
+            torch_dtype=torch.float32,
+            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
         )
     teacher_transformer = deepcopy(transformer)
     if args.use_ema:
@@ -291,18 +344,26 @@ def main(args):
         )
         transformer.add_adapter(transformer_lora_config)
 
-    main_print(f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M")
-    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
-    fsdp_kwargs = get_dit_fsdp_kwargs(args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload, args.master_weight_type)
-    
+    main_print(
+        f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M"
+    )
+    main_print(
+        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
+    )
+    fsdp_kwargs = get_dit_fsdp_kwargs(
+        args.fsdp_sharding_startegy,
+        args.use_lora,
+        args.use_cpu_offload,
+        args.master_weight_type,
+    )
+
     if args.use_lora:
         transformer.config.lora_rank = args.lora_rank
         transformer.config.lora_alpha = args.lora_alpha
         transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
         transformer._no_split_modules = ["MochiTransformerBlock"]
-        fsdp_kwargs['auto_wrap_policy'] = fsdp_kwargs['auto_wrap_policy'](transformer)
-    
-    
+        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
+
     transformer = FSDP(
         transformer,
         **fsdp_kwargs,
@@ -330,8 +391,14 @@ def main(args):
         ema_transformer.requires_grad_(False)
     noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift)
     if args.scheduler_type == "pcm_linear_quadratic":
-        linear_steps = int(noise_scheduler.config.num_train_timesteps * args.linear_range)
-        sigmas = linear_quadratic_schedule(noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold, linear_steps)
+        linear_steps = int(
+            noise_scheduler.config.num_train_timesteps * args.linear_range
+        )
+        sigmas = linear_quadratic_schedule(
+            noise_scheduler.config.num_train_timesteps,
+            args.linear_quadratic_threshold,
+            linear_steps,
+        )
         sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
     else:
         sigmas = noise_scheduler.sigmas
@@ -347,7 +414,7 @@ def main(args):
     optimizer = torch.optim.AdamW(
         params_to_optimize,
         lr=args.learning_rate,
-        betas=(0.9,0.999),
+        betas=(0.9, 0.999),
         weight_decay=args.weight_decay,
         eps=1e-8,
     )
@@ -356,32 +423,38 @@ def main(args):
     if args.resume_from_lora_checkpoint:
         transformer, optimizer, init_steps = resume_lora_optimizer(
             transformer, args.resume_from_lora_checkpoint, optimizer
-        )      
+        )
     main_print(f"optimizer: {optimizer}")
 
-    #todo add lr scheduler
+    # todo add lr scheduler
     lr_scheduler = get_scheduler(
-                args.lr_scheduler,
-                optimizer=optimizer,
-                num_warmup_steps=args.lr_warmup_steps * world_size,
-                num_training_steps=args.max_train_steps * world_size,
-                num_cycles=args.lr_num_cycles,
-                power=args.lr_power,
-                last_epoch=init_steps - 1,
-            )
-    
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * world_size,
+        num_training_steps=args.max_train_steps * world_size,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+        last_epoch=init_steps - 1,
+    )
+
     train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
     uncond_prompt_embed = train_dataset.uncond_prompt_embed
     uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = LengthGroupedSampler(
-                args.train_batch_size,
-                rank=rank,
-                world_size=world_size,
-                lengths=train_dataset.lengths, 
-                group_frame=args.group_frame, 
-                group_resolution=args.group_resolution, 
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=False)
-    
+    sampler = (
+        LengthGroupedSampler(
+            args.train_batch_size,
+            rank=rank,
+            world_size=world_size,
+            lengths=train_dataset.lengths,
+            group_frame=args.group_frame,
+            group_resolution=args.group_resolution,
+        )
+        if (args.group_frame or args.group_resolution)
+        else DistributedSampler(
+            train_dataset, rank=rank, num_replicas=world_size, shuffle=False
+        )
+    )
+
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
@@ -389,54 +462,73 @@ def main(args):
         pin_memory=True,
         batch_size=args.train_batch_size,
         num_workers=args.dataloader_num_workers,
-        drop_last=True, 
+        drop_last=True,
     )
 
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader)
+        / args.gradient_accumulation_steps
+        * args.sp_size
+        / args.train_sp_batch_size
+    )
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
-
     if rank <= 0:
         project = args.tracker_project_name or "fastvideo"
         wandb.init(project=project, config=args)
 
     # Train!
-    total_batch_size = args.train_batch_size * world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
+    total_batch_size = (
+        args.train_batch_size
+        * world_size
+        * args.gradient_accumulation_steps
+        / args.sp_size
+        * args.train_sp_batch_size
+    )
     main_print("***** Running training *****")
     main_print(f"  Num examples = {len(train_dataset)}")
     main_print(f"  Dataloader size = {len(train_dataloader)}")
     main_print(f"  Num Epochs = {args.num_train_epochs}")
     main_print(f"  Resume training from step {init_steps}")
     main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
+    main_print(
+        f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
+    )
     main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
     main_print(f"  Total optimization steps = {args.max_train_steps}")
-    main_print(f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
+    main_print(
+        f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B"
+    )
     # print dtype
     main_print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
 
-
     # Potentially load in the weights and states from a previous save
     if args.resume_from_checkpoint:
         assert NotImplementedError("resume_from_checkpoint is not supported now.")
-        # TODO 
-
+        # TODO
 
     progress_bar = tqdm(
         range(0, args.max_train_steps),
         initial=init_steps,
         desc="Steps",
         # Only show the progress bar once on each machine.
-        disable= local_rank > 0,
+        disable=local_rank > 0,
+    )
+
+    loader = sp_parallel_dataloader_wrapper(
+        train_dataloader,
+        device,
+        args.train_batch_size,
+        args.sp_size,
+        args.train_sp_batch_size,
     )
 
-    loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size)
-    
     step_times = deque(maxlen=100)
 
-    #todo future 
+    # todo future
     for i in range(init_steps):
         next(loader)
+
     # log_validation(args, transformer, device,
     #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
     def get_num_phases(multi_phased_distill_schedule, step):
@@ -448,41 +540,71 @@ def main(args):
             if step <= int(phase_step):
                 return int(phase)
         return phase
-    for step in range(init_steps + 1, args.max_train_steps+1):
+
+    for step in range(init_steps + 1, args.max_train_steps + 1):
         start_time = time.time()
         assert args.multi_phased_distill_schedule is not None
         num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
 
-        loss, grad_norm, pred_norm, aux_loss = train_one_step_mochi(transformer,teacher_transformer, ema_transformer, optimizer, lr_scheduler, loader, noise_scheduler,solver, noise_random_generator, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, args.num_euler_timesteps, num_phases, args.not_apply_cfg_solver,args.distill_cfg, args.ema_decay , args.pred_decay_weight, args.pred_decay_type)
+        loss, grad_norm, pred_norm, aux_loss = train_one_step_mochi(
+            transformer,
+            teacher_transformer,
+            ema_transformer,
+            optimizer,
+            lr_scheduler,
+            loader,
+            noise_scheduler,
+            solver,
+            noise_random_generator,
+            args.gradient_accumulation_steps,
+            args.sp_size,
+            args.max_grad_norm,
+            uncond_prompt_embed,
+            uncond_prompt_mask,
+            args.num_euler_timesteps,
+            num_phases,
+            args.not_apply_cfg_solver,
+            args.distill_cfg,
+            args.ema_decay,
+            args.pred_decay_weight,
+            args.pred_decay_type,
+        )
 
         step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 
-        progress_bar.set_postfix({
-        "loss": f"{loss:.4f}", 
-        "step_time": f"{step_time:.2f}s",
-        "grad_norm": grad_norm,
-        "phases": num_phases,
-    })
+        progress_bar.set_postfix(
+            {
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+                "phases": num_phases,
+            }
+        )
         progress_bar.update(1)
         if rank <= 0:
-            wandb.log({
-            "train_loss": loss,
-            "learning_rate": lr_scheduler.get_last_lr()[0],
-            "step_time": step_time,
-            "avg_step_time": avg_step_time,
-            "grad_norm": grad_norm ,
-            "pred_fro_norm": pred_norm["fro"],
-            "pred_largest_singular_value": pred_norm["largest singular value"],
-            "pred_absolute_mean": pred_norm["absolute mean"],
-            "pred_absolute_max": pred_norm["absolute max"],
-            "aux_loss": aux_loss,
-        }, step=step)
-        if step  % args.checkpointing_steps == 0:
+            wandb.log(
+                {
+                    "train_loss": loss,
+                    "learning_rate": lr_scheduler.get_last_lr()[0],
+                    "step_time": step_time,
+                    "avg_step_time": avg_step_time,
+                    "grad_norm": grad_norm,
+                    "pred_fro_norm": pred_norm["fro"],
+                    "pred_largest_singular_value": pred_norm["largest singular value"],
+                    "pred_absolute_mean": pred_norm["absolute mean"],
+                    "pred_absolute_max": pred_norm["absolute max"],
+                    "aux_loss": aux_loss,
+                },
+                step=step,
+            )
+        if step % args.checkpointing_steps == 0:
             if args.use_lora:
                 # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
+                save_lora_checkpoint(
+                    transformer, optimizer, rank, args.output_dir, step
+                )
             else:
                 # Your existing checkpoint saving code
                 if args.use_ema:
@@ -490,18 +612,42 @@ def main(args):
                 else:
                     save_checkpoint(transformer, rank, args.output_dir, step)
             dist.barrier()
-        if args.log_validation and step  % args.validation_steps == 0:
-            log_validation(args, transformer, device,
-                            torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps,  linear_quadratic_threshold=args.linear_quadratic_threshold,linear_range=args.linear_range, ema=False)
+        if args.log_validation and step % args.validation_steps == 0:
+            log_validation(
+                args,
+                transformer,
+                device,
+                torch.bfloat16,
+                step,
+                scheduler_type=args.scheduler_type,
+                shift=args.shift,
+                num_euler_timesteps=args.num_euler_timesteps,
+                linear_quadratic_threshold=args.linear_quadratic_threshold,
+                linear_range=args.linear_range,
+                ema=False,
+            )
             if args.use_ema:
-                log_validation(args, ema_transformer, device,
-                                torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,linear_range=args.linear_range, ema=True)
+                log_validation(
+                    args,
+                    ema_transformer,
+                    device,
+                    torch.bfloat16,
+                    step,
+                    scheduler_type=args.scheduler_type,
+                    shift=args.shift,
+                    num_euler_timesteps=args.num_euler_timesteps,
+                    linear_quadratic_threshold=args.linear_quadratic_threshold,
+                    linear_range=args.linear_range,
+                    ema=True,
+                )
 
     if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
+        save_lora_checkpoint(
+            transformer, optimizer, rank, args.output_dir, args.max_train_steps
+        )
     else:
-        save_checkpoint(transformer,  rank, args.output_dir, args.max_train_steps)
-        
+        save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
+
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
 
@@ -512,116 +658,244 @@ if __name__ == "__main__":
     # dataset & dataloader
     parser.add_argument("--data_json_path", type=str, required=True)
     parser.add_argument("--num_frames", type=int, default=163)
-    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("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true") # TODO
-    parser.add_argument("--group_resolution", action="store_true") # TODO
+    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(
+        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
+    )
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
 
     # text encoder & vae & diffusion model
     parser.add_argument("--pretrained_model_name_or_path", type=str)
     parser.add_argument("--dit_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
 
     # diffusion setting
     parser.add_argument("--ema_decay", type=float, default=0.95)
     parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument('--cfg', type=float, default=0.1)
-    parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.")
-    
+    parser.add_argument("--cfg", type=float, default=0.1)
+
     # validation & logs
     parser.add_argument("--validation_prompt_dir", type=str)
     parser.add_argument("--validation_sampling_steps", type=str, default="64")
-    parser.add_argument('--validation_guidance_scale', type=str, default="4.5")
+    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
 
-    parser.add_argument('--validation_steps', type=float, default=64)
+    parser.add_argument("--validation_steps", type=float, default=64)
     parser.add_argument("--log_validation", action="store_true")
     parser.add_argument("--tracker_project_name", type=str, default=None)
-    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("--shift", type=float, default=1.0 )
-    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("--resume_from_lora_checkpoint", type=str, default=None,
-                        help=(
-                            "Whether training should be resumed from a previous lora 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(
+        "--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("--shift", type=float, default=1.0)
+    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(
+        "--resume_from_lora_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous lora 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***."
+        ),
+    )
 
     # optimizer & scheduler & Training
     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("--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=10, help="Number of steps for the warmup in the lr scheduler.")
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
+    parser.add_argument(
+        "--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(
+        "--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=10,
+        help="Number of steps for the warmup in the lr scheduler.",
+    )
+    parser.add_argument(
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
     parser.add_argument("--selective_checkpointing", type=float, default=1.0)
-    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("--use_cpu_offload", action="store_true", help="Whether to use CPU offload for param & gradient & optimizer states.")
+    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(
+        "--use_cpu_offload",
+        action="store_true",
+        help="Whether to use CPU offload for param & gradient & optimizer states.",
+    )
 
     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")
+    parser.add_argument(
+        "--train_sp_batch_size",
+        type=int,
+        default=1,
+        help="Batch size for sequence parallel training",
+    )
 
-    parser.add_argument("--use_lora", action="store_true", default=False, help="Whether to use LoRA for finetuning.") 
-    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        default=False,
+        help="Whether to use LoRA for finetuning.",
+    )
+    parser.add_argument(
+        "--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA."
+    )
+    parser.add_argument(
+        "--lora_rank", type=int, default=128, help="LoRA rank parameter. "
+    )
     parser.add_argument("--fsdp_sharding_startegy", default="full")
 
     # lr_scheduler
-    parser.add_argument("--lr_scheduler", type=str, default="constant",
+    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("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument("--lr_num_cycles", type=int, default=1, help="Number of cycles in the learning rate scheduler.")
-    parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.",)
-    parser.add_argument("--not_apply_cfg_solver", action="store_true", help="Whether to apply the cfg_solver.")
-    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of cycles in the learning rate scheduler.",
+    )
+    parser.add_argument(
+        "--lr_power",
+        type=float,
+        default=1.0,
+        help="Power factor of the polynomial scheduler.",
+    )
+    parser.add_argument(
+        "--not_apply_cfg_solver",
+        action="store_true",
+        help="Whether to apply the cfg_solver.",
+    )
+    parser.add_argument(
+        "--distill_cfg", type=float, default=3.0, help="Distillation coefficient."
+    )
     # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
-    parser.add_argument("--linear_quadratic_threshold", type=float, default=0.025, help="Threshold for linear quadratic scheduler.")
-    parser.add_argument("--linear_range", type=float, default=0.5, help="Range for linear quadratic scheduler.")
-    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
+    parser.add_argument(
+        "--scheduler_type", type=str, default="pcm", help="The scheduler type to use."
+    )
+    parser.add_argument(
+        "--linear_quadratic_threshold",
+        type=float,
+        default=0.025,
+        help="Threshold for linear quadratic scheduler.",
+    )
+    parser.add_argument(
+        "--linear_range",
+        type=float,
+        default=0.5,
+        help="Range for linear quadratic scheduler.",
+    )
+    parser.add_argument(
+        "--weight_decay", type=float, default=0.001, help="Weight decay to apply."
+    )
     parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA.")
-    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None) 
+    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
     parser.add_argument("--finetune_weight", type=float, default=0.0)
     parser.add_argument("--pred_decay_weight", type=float, default=0.0)
     parser.add_argument("--pred_decay_type", default="l1")
-    parser.add_argument("--master_weight_type", type=str, default="fp32", help="Weight type to use - fp32 or bf16.")
+    parser.add_argument(
+        "--master_weight_type",
+        type=str,
+        default="fp32",
+        help="Weight type to use - fp32 or bf16.",
+    )
     args = parser.parse_args()
-    main(args)
\ No newline at end of file
+    main(args)
diff --git a/fastvideo/distill/discriminator.py b/fastvideo/distill/discriminator.py
index 8eb827350..d586bfb70 100644
--- a/fastvideo/distill/discriminator.py
+++ b/fastvideo/distill/discriminator.py
@@ -23,7 +23,6 @@ from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-
 class DiscriminatorHead(nn.Module):
     def __init__(self, input_channel, output_channel=1):
         super().__init__()
@@ -48,9 +47,9 @@ class DiscriminatorHead(nn.Module):
     def forward(self, x):
         b, twh, c = x.shape
         t = twh // (30 * 53)
-        x = x.view(-1, 30 *53, c)
+        x = x.view(-1, 30 * 53, c)
         x = x.permute(0, 2, 1)
-        x = x.view(b*t, c, 30, 53)
+        x = x.view(b * t, c, 30, 53)
         x = self.conv1(x)
         x = self.conv2(x) + x
         x = self.conv_out(x)
@@ -58,10 +57,9 @@ class DiscriminatorHead(nn.Module):
 
 
 class Discriminator(nn.Module):
-
     def __init__(
         self,
-        stride = 8,
+        stride=8,
         num_h_per_head=1,
         adapter_channel_dims=[3072],
     ):
@@ -82,24 +80,23 @@ class Discriminator(nn.Module):
             ]
         )
 
-
-
     def forward(self, features):
         outputs = []
+
         def create_custom_forward(module):
             def custom_forward(*inputs):
                 return module(*inputs)
+
             return custom_forward
+
         assert len(features) // self.stride == len(self.heads)
         for i in range(0, len(features), self.stride):
-            for h in self.heads[i//self.stride]:
+            for h in self.heads[i // self.stride]:
                 # out = torch.utils.checkpoint.checkpoint(
                 #     create_custom_forward(h),
                 #     features[i],
                 #     use_reentrant=False
                 # )
-                out=h(features[i])
+                out = h(features[i])
                 outputs.append(out)
         return outputs
-
-
diff --git a/fastvideo/distill/solver.py b/fastvideo/distill/solver.py
index e53bfebdb..8420f32e7 100644
--- a/fastvideo/distill/solver.py
+++ b/fastvideo/distill/solver.py
@@ -17,13 +17,14 @@ logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 class PCMFMSchedulerOutput(BaseOutput):
     prev_sample: torch.FloatTensor
 
+
 def extract_into_tensor(a, t, x_shape):
     b, *_ = t.shape
     out = a.gather(-1, t)
     return out.reshape(b, *((1,) * (len(x_shape) - 1)))
 
-class PCMFMScheduler(SchedulerMixin, ConfigMixin):
 
+class PCMFMScheduler(SchedulerMixin, ConfigMixin):
     _compatibles = []
     order = 1
 
@@ -34,13 +35,14 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
         shift: float = 1.0,
         pcm_timesteps: int = 50,
         linear_quadratic=False,
-        linear_quadratic_threshold=0.025, 
+        linear_quadratic_threshold=0.025,
         linear_range=0.5,
     ):
- 
         if linear_quadratic:
             linear_steps = int(num_train_timesteps * linear_range)
-            sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps)
+            sigmas = linear_quadratic_schedule(
+                num_train_timesteps, linear_quadratic_threshold, linear_steps
+            )
             sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
         else:
             timesteps = np.linspace(
@@ -238,6 +240,7 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
     def __len__(self):
         return self.config.num_train_timesteps
 
+
 class EulerSolver:
     def __init__(self, sigmas, timesteps=1000, euler_timesteps=50):
         self.step_ratio = timesteps // euler_timesteps
@@ -279,7 +282,6 @@ class EulerSolver:
         multiphase,
         is_target=False,
     ):
-
         inference_indices = np.linspace(
             0, len(self.euler_timesteps), num=multiphase, endpoint=False
         )
@@ -305,4 +307,3 @@ class EulerSolver:
         x_prev = sample + (sigma_prev - sigma) * model_pred
 
         return x_prev, timestep_index_end
-
diff --git a/fastvideo/distill_adv.py b/fastvideo/distill_adv.py
index 31fa853fd..15c2da684 100644
--- a/fastvideo/distill_adv.py
+++ b/fastvideo/distill_adv.py
@@ -5,8 +5,12 @@ import math
 import os
 import shutil
 from pathlib import Path
-from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
-    destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info
+from fastvideo.utils.parallel_states import (
+    initialize_sequence_parallel_state,
+    destroy_sequence_parallel_group,
+    get_sequence_parallel_state,
+    nccl_info,
+)
 from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
 from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_mochi_dit_input
 from fastvideo.utils.validation import log_validation
@@ -26,7 +30,11 @@ from fastvideo.utils.dataset_utils import LengthGroupedSampler
 import wandb
 from accelerate.utils import set_seed
 from tqdm.auto import tqdm
-from fastvideo.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing, get_discriminator_fsdp_kwargs
+from fastvideo.fsdp_util import (
+    get_dit_fsdp_kwargs,
+    apply_fsdp_checkpointing,
+    get_discriminator_fsdp_kwargs,
+)
 import diffusers
 from diffusers import (
     FlowMatchEulerDiscreteScheduler,
@@ -43,17 +51,22 @@ from peft import LoraConfig
 from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
 )
-from fastvideo.utils.checkpoint import save_checkpoint, save_lora_checkpoint, resume_lora_optimizer, resume_training, save_checkpoint_generator_discriminator, resume_training_generator_discriminator
+from fastvideo.utils.checkpoint import (
+    save_checkpoint,
+    save_lora_checkpoint,
+    resume_lora_optimizer,
+    resume_training,
+    save_checkpoint_generator_discriminator,
+    resume_training_generator_discriminator,
+)
 from fastvideo.utils.logging import main_print
+
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
 check_min_version("0.31.0")
 import time
 from collections import deque
 
 
-
-
-    
 def gan_d_loss(
     discriminator,
     teacher_transformer,
@@ -73,7 +86,7 @@ def gan_d_loss(
             timestep,
             encoder_attention_mask,
             output_attn=True,
-            return_dict= False
+            return_dict=False,
         )[1]
         real_features = teacher_transformer(
             sample_real,
@@ -81,15 +94,11 @@ def gan_d_loss(
             timestep,
             encoder_attention_mask,
             output_attn=True,
-            return_dict= False
+            return_dict=False,
         )[1]
-        
-    fake_outputs = discriminator(
-        fake_features
-    )
-    real_outputs = discriminator(
-        real_features
-    )
+
+    fake_outputs = discriminator(fake_features)
+    real_outputs = discriminator(real_features)
     for fake_output, real_output in zip(fake_outputs, real_outputs):
         loss += (
             torch.mean(weight * torch.relu(fake_output.float() + 1))
@@ -97,6 +106,7 @@ def gan_d_loss(
         ) / (discriminator.head_num * discriminator.num_h_per_head)
     return loss
 
+
 def gan_g_loss(
     discriminator,
     teacher_transformer,
@@ -108,12 +118,12 @@ def gan_g_loss(
 ):
     loss = 0.0
     features = teacher_transformer(
-            sample_fake,
-            encoder_hidden_states,
-            timestep,
-            encoder_attention_mask,
-            output_attn=True,
-            return_dict= False
+        sample_fake,
+        encoder_hidden_states,
+        timestep,
+        encoder_attention_mask,
+        output_attn=True,
+        return_dict=False,
     )[1]
     fake_outputs = discriminator(
         features,
@@ -124,13 +134,39 @@ def gan_g_loss(
         )
     return loss
 
-def train_one_step_mochi(transformer, teacher_transformer , optimizer, discriminator, discriminator_optimizer,global_step, lr_scheduler,loader, noise_scheduler, solver,noise_random_generator, sp_size, precondition_outputs, max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, num_euler_timesteps, multiphase, not_apply_cfg_solver, distill_cfg, adv_weight):
-
 
+def train_one_step_mochi(
+    transformer,
+    teacher_transformer,
+    optimizer,
+    discriminator,
+    discriminator_optimizer,
+    global_step,
+    lr_scheduler,
+    loader,
+    noise_scheduler,
+    solver,
+    noise_random_generator,
+    sp_size,
+    precondition_outputs,
+    max_grad_norm,
+    uncond_prompt_embed,
+    uncond_prompt_mask,
+    num_euler_timesteps,
+    multiphase,
+    not_apply_cfg_solver,
+    distill_cfg,
+    adv_weight,
+):
     optimizer.zero_grad()
     discriminator_optimizer.zero_grad()
 
-    latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
+    (
+        latents,
+        encoder_hidden_states,
+        latents_attention_mask,
+        encoder_attention_mask,
+    ) = next(loader)
     model_input = normalize_mochi_dit_input(latents)
     noise = torch.randn_like(model_input)
     bsz = model_input.shape[0]
@@ -142,36 +178,29 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
     # Add noise according to flow matching.
     # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
     sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
-    sigmas_prev = extract_into_tensor(
-        solver.sigmas_prev, index, model_input.shape
-    )
+    sigmas_prev = extract_into_tensor(solver.sigmas_prev, index, model_input.shape)
 
-    timesteps = (
-        sigmas * noise_scheduler.config.num_train_timesteps
-    ).view(-1)
+    timesteps = (sigmas * noise_scheduler.config.num_train_timesteps).view(-1)
     # if squeeze to [], unsqueeze to [1]
 
-    timesteps_prev = (
-        sigmas_prev * noise_scheduler.config.num_train_timesteps
-    ).view(-1)
+    timesteps_prev = (sigmas_prev * noise_scheduler.config.num_train_timesteps).view(-1)
     noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
 
-    # Predict the noise residual 
+    # Predict the noise residual
     with torch.autocast("cuda", dtype=torch.bfloat16):
-
         model_pred = transformer(
             noisy_model_input,
             encoder_hidden_states,
             timesteps,
-            encoder_attention_mask, # B, L
-            return_dict= False
+            encoder_attention_mask,  # B, L
+            return_dict=False,
         )[0]
 
     # if accelerator.is_main_process:
     model_pred, end_index = solver.euler_style_multiphase_pred(
         noisy_model_input, model_pred, index, multiphase
     )
-    
+
     weighting = 1.0
     # # simplified flow matching aka 0-rectified flow matching loss
     # # target = model_input - noise
@@ -180,27 +209,17 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
     for i in range(end_index.size(0)):
         adv_index[i] = torch.randint(
             end_index[i].item(),
-            end_index[i].item()
-            + num_euler_timesteps // multiphase,
+            end_index[i].item() + num_euler_timesteps // multiphase,
             (1,),
             dtype=end_index.dtype,
             device=end_index.device,
         )
-        
-    sigmas_end = extract_into_tensor(
-        solver.sigmas_prev, end_index, model_input.shape
-    )
-    sigmas_adv = extract_into_tensor(
-        solver.sigmas_prev, adv_index, model_input.shape
-    )
-    timesteps_end = (
-        sigmas_end * noise_scheduler.config.num_train_timesteps
-    ).view(-1)
-    timesteps_adv = (
-        sigmas_adv * noise_scheduler.config.num_train_timesteps
-    ).view(-1)
-    
-    
+
+    sigmas_end = extract_into_tensor(solver.sigmas_prev, end_index, model_input.shape)
+    sigmas_adv = extract_into_tensor(solver.sigmas_prev, adv_index, model_input.shape)
+    timesteps_end = (sigmas_end * noise_scheduler.config.num_train_timesteps).view(-1)
+    timesteps_adv = (sigmas_adv * noise_scheduler.config.num_train_timesteps).view(-1)
+
     with torch.no_grad():
         w = distill_cfg
         with torch.autocast("cuda", dtype=torch.bfloat16):
@@ -208,8 +227,8 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
                 noisy_model_input,
                 encoder_hidden_states,
                 timesteps,
-                encoder_attention_mask, # B, L
-                return_dict= False
+                encoder_attention_mask,  # B, L
+                return_dict=False,
             )[0].float()
         if not_apply_cfg_solver:
             uncond_teacher_output = cond_teacher_output
@@ -221,14 +240,12 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
                     uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
                     timesteps,
                     uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
-                    return_dict= False
+                    return_dict=False,
                 )[0].float()
         teacher_output = cond_teacher_output + w * (
             cond_teacher_output - uncond_teacher_output
         )
-        x_prev = solver.euler_step(
-            noisy_model_input, teacher_output, index
-        )
+        x_prev = solver.euler_step(noisy_model_input, teacher_output, index)
 
     # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
     with torch.no_grad():
@@ -237,8 +254,8 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
                 x_prev.float(),
                 encoder_hidden_states,
                 timesteps_prev,
-                encoder_attention_mask, # B, L
-                return_dict= False
+                encoder_attention_mask,  # B, L
+                return_dict=False,
             )[0]
 
         target, end_index = solver.euler_style_multiphase_pred(
@@ -246,24 +263,16 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
         )
 
     real_adv = (
-        (1 - sigmas_adv) * target
-        + (sigmas_adv - sigmas_end) * torch.randn_like(target)
+        (1 - sigmas_adv) * target + (sigmas_adv - sigmas_end) * torch.randn_like(target)
     ) / (1 - sigmas_end)
     fake_adv = (
         (1 - sigmas_adv) * model_pred
         + (sigmas_adv - sigmas_end) * torch.randn_like(model_pred)
     ) / (1 - sigmas_end)
-    
-    
 
-        
     huber_c = 0.001
     g_loss = torch.mean(
-        torch.sqrt(
-            (model_pred.float() - target.float()) ** 2
-            + huber_c**2
-        )
-        - huber_c
+        torch.sqrt((model_pred.float() - target.float()) ** 2 + huber_c**2) - huber_c
     )
     discriminator.requires_grad_(False)
     with torch.autocast("cuda", dtype=torch.bfloat16):
@@ -278,16 +287,15 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
         )
     g_loss += g_gan_loss
     g_loss.backward()
-    
+
     g_loss = g_loss.detach().clone()
     dist.all_reduce(g_loss, op=dist.ReduceOp.AVG)
-        
 
     g_grad_norm = transformer.clip_grad_norm_(max_grad_norm).item()
     optimizer.step()
     lr_scheduler.step()
     optimizer.zero_grad()
-    
+
     discriminator_optimizer.zero_grad()
     discriminator.requires_grad_(True)
     with torch.autocast("cuda", dtype=torch.bfloat16):
@@ -300,23 +308,22 @@ def train_one_step_mochi(transformer, teacher_transformer , optimizer, discrimin
             encoder_hidden_states,
             encoder_attention_mask,
             1.0,
-        ) 
-    
+        )
+
     d_loss.backward()
     d_grad_norm = discriminator.clip_grad_norm_(max_grad_norm).item()
     discriminator_optimizer.step()
     discriminator_optimizer.zero_grad()
-    
+
     return g_loss, g_grad_norm, d_loss, d_grad_norm
 
 
 def main(args):
-
     torch.backends.cuda.matmul.allow_tf32 = True
-    
-    local_rank = int(os.environ['LOCAL_RANK'])
-    rank = int(os.environ['RANK'])
-    world_size = int(os.environ['WORLD_SIZE'])
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
     dist.init_process_group("nccl")
     torch.cuda.set_device(local_rank)
     device = torch.cuda.current_device()
@@ -330,33 +337,32 @@ def main(args):
     noise_random_generator = None
 
     # Handle the repository creation
-    if rank <=0 and args.output_dir is not None:
+    if rank <= 0 and 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.
 
     # Create model:
-    
+
     main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
     # keep the master weight to float32
     if args.dit_model_name_or_path:
         transformer = transformer = MochiTransformer3DModel.from_pretrained(
             args.dit_model_name_or_path,
-            torch_dtype = torch.float32,
-            #torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
+            torch_dtype=torch.float32,
+            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
         )
     else:
         transformer = MochiTransformer3DModel.from_pretrained(
             args.pretrained_model_name_or_path,
             subfolder="transformer",
-            torch_dtype = torch.float32,
-            #torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
+            torch_dtype=torch.float32,
+            # torch_dtype=torch.bfloat16 if args.use_lora else torch.float32,
         )
     teacher_transformer = deepcopy(transformer)
     discriminator = Discriminator(args.discriminator_head_stride)
-    
-    
+
     if args.use_lora:
         transformer.requires_grad_(False)
         transformer_lora_config = LoraConfig(
@@ -367,20 +373,27 @@ def main(args):
         )
         transformer.add_adapter(transformer_lora_config)
 
-    main_print(f"  Total transformer parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M")
+    main_print(
+        f"  Total transformer parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M"
+    )
     # discriminator
-    main_print(f"  Total discriminator parameters = {sum(p.numel() for p in discriminator.parameters() if p.requires_grad) / 1e6} M")
-    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
-    fsdp_kwargs = get_dit_fsdp_kwargs(args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload)
+    main_print(
+        f"  Total discriminator parameters = {sum(p.numel() for p in discriminator.parameters() if p.requires_grad) / 1e6} M"
+    )
+    main_print(
+        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
+    )
+    fsdp_kwargs = get_dit_fsdp_kwargs(
+        args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload
+    )
     discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs(args.master_weight_type)
     if args.use_lora:
         transformer.config.lora_rank = args.lora_rank
         transformer.config.lora_alpha = args.lora_alpha
         transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
         transformer._no_split_modules = ["MochiTransformerBlock"]
-        fsdp_kwargs['auto_wrap_policy'] = fsdp_kwargs['auto_wrap_policy'](transformer)
-    
-    
+        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
+
     transformer = FSDP(
         transformer,
         **fsdp_kwargs,
@@ -403,7 +416,9 @@ def main(args):
     teacher_transformer.requires_grad_(False)
     noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=args.shift)
     if args.scheduler_type == "pcm_linear_quadratic":
-        sigmas = linear_quadratic_schedule(noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold)
+        sigmas = linear_quadratic_schedule(
+            noise_scheduler.config.num_train_timesteps, args.linear_quadratic_threshold
+        )
         sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
     else:
         sigmas = noise_scheduler.sigmas
@@ -419,11 +434,11 @@ def main(args):
     optimizer = torch.optim.AdamW(
         params_to_optimize,
         lr=args.learning_rate,
-        betas=(0.9,0.999),
+        betas=(0.9, 0.999),
         weight_decay=1e-3,
         eps=1e-8,
     )
-    
+
     discriminator_optimizer = torch.optim.AdamW(
         discriminator.parameters(),
         lr=args.discriminator_learning_rate,
@@ -431,42 +446,58 @@ def main(args):
         weight_decay=1e-3,
         eps=1e-8,
     )
-        
-        
+
     init_steps = 0
     if args.resume_from_lora_checkpoint:
         transformer, optimizer, init_steps = resume_lora_optimizer(
             transformer, args.resume_from_lora_checkpoint, optimizer
-        )   
-    elif args.resume_from_checkpoint:
-        transformer, optimizer,discriminator, discriminator_optimizer, init_steps = resume_training_generator_discriminator(
-            transformer, optimizer,discriminator, discriminator_optimizer, args.resume_from_checkpoint, rank
         )
-  
+    elif args.resume_from_checkpoint:
+        (
+            transformer,
+            optimizer,
+            discriminator,
+            discriminator_optimizer,
+            init_steps,
+        ) = resume_training_generator_discriminator(
+            transformer,
+            optimizer,
+            discriminator,
+            discriminator_optimizer,
+            args.resume_from_checkpoint,
+            rank,
+        )
+
     main_print(f"optimizer: {optimizer}")
 
     lr_scheduler = get_scheduler(
-                args.lr_scheduler,
-                optimizer=optimizer,
-                num_warmup_steps=args.lr_warmup_steps * world_size,
-                num_training_steps=args.max_train_steps * world_size,
-                num_cycles=args.lr_num_cycles,
-                power=args.lr_power,
-                last_epoch=init_steps - 1,
-            )
-    
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * world_size,
+        num_training_steps=args.max_train_steps * world_size,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+        last_epoch=init_steps - 1,
+    )
+
     train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
     uncond_prompt_embed = train_dataset.uncond_prompt_embed
     uncond_prompt_mask = train_dataset.uncond_prompt_mask
-    sampler = LengthGroupedSampler(
-                args.train_batch_size,
-                rank=rank,
-                world_size=world_size,
-                lengths=train_dataset.lengths, 
-                group_frame=args.group_frame, 
-                group_resolution=args.group_resolution, 
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=False)
-    
+    sampler = (
+        LengthGroupedSampler(
+            args.train_batch_size,
+            rank=rank,
+            world_size=world_size,
+            lengths=train_dataset.lengths,
+            group_frame=args.group_frame,
+            group_resolution=args.group_resolution,
+        )
+        if (args.group_frame or args.group_resolution)
+        else DistributedSampler(
+            train_dataset, rank=rank, num_replicas=world_size, shuffle=False
+        )
+    )
+
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
@@ -474,96 +505,176 @@ def main(args):
         pin_memory=True,
         batch_size=args.train_batch_size,
         num_workers=args.dataloader_num_workers,
-        drop_last=True, 
+        drop_last=True,
+    )
+    assert args.gradient_accumulation_steps == 1
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader)
+        / args.gradient_accumulation_steps
+        * args.sp_size
+        / args.train_sp_batch_size
     )
-    assert args.gradient_accumulation_steps == 1 
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
-
     if rank <= 0:
         project = args.tracker_project_name or "fastvideo"
         wandb.init(project=project, config=args)
 
     # Train!
-    total_batch_size = args.train_batch_size * world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
+    total_batch_size = (
+        args.train_batch_size
+        * world_size
+        * args.gradient_accumulation_steps
+        / args.sp_size
+        * args.train_sp_batch_size
+    )
     main_print("***** Running training *****")
     main_print(f"  Num examples = {len(train_dataset)}")
     main_print(f"  Dataloader size = {len(train_dataloader)}")
     main_print(f"  Num Epochs = {args.num_train_epochs}")
     main_print(f"  Resume training from step {init_steps}")
     main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
+    main_print(
+        f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
+    )
     main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
     main_print(f"  Total optimization steps = {args.max_train_steps}")
-    main_print(f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
+    main_print(
+        f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B"
+    )
     # print dtype
     main_print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
 
-
-
-
     progress_bar = tqdm(
         range(0, args.max_train_steps),
         initial=init_steps,
         desc="Steps",
         # Only show the progress bar once on each machine.
-        disable= local_rank > 0,
+        disable=local_rank > 0,
     )
-    loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size)
-    
+    loader = sp_parallel_dataloader_wrapper(
+        train_dataloader,
+        device,
+        args.train_batch_size,
+        args.sp_size,
+        args.train_sp_batch_size,
+    )
+
     step_times = deque(maxlen=100)
     # log_validation(args, transformer, device,
     #         torch.bfloat16, init_steps, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False)
 
     for i in range(init_steps):
         _ = next(loader)
-    for step in range(init_steps + 1, args.max_train_steps+1):
+    for step in range(init_steps + 1, args.max_train_steps + 1):
         start_time = time.time()
-        generator_loss, generator_grad_norm, discriminator_loss, discriminator_grad_norm= train_one_step_mochi(transformer,teacher_transformer, optimizer, discriminator, discriminator_optimizer, step,lr_scheduler, loader, noise_scheduler,solver, noise_random_generator , args.sp_size, args.precondition_outputs, args.max_grad_norm, uncond_prompt_embed, uncond_prompt_mask, args.num_euler_timesteps, args.validation_sampling_steps, args.not_apply_cfg_solver,args.distill_cfg, args.adv_weight)
+        (
+            generator_loss,
+            generator_grad_norm,
+            discriminator_loss,
+            discriminator_grad_norm,
+        ) = train_one_step_mochi(
+            transformer,
+            teacher_transformer,
+            optimizer,
+            discriminator,
+            discriminator_optimizer,
+            step,
+            lr_scheduler,
+            loader,
+            noise_scheduler,
+            solver,
+            noise_random_generator,
+            args.sp_size,
+            args.precondition_outputs,
+            args.max_grad_norm,
+            uncond_prompt_embed,
+            uncond_prompt_mask,
+            args.num_euler_timesteps,
+            args.validation_sampling_steps,
+            args.not_apply_cfg_solver,
+            args.distill_cfg,
+            args.adv_weight,
+        )
 
         step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 
-        progress_bar.set_postfix({
-        "g_loss": f"{generator_loss:.4f}", 
-        "d_loss": f"{discriminator_loss:.4f}",
-        "g_grad_norm": generator_grad_norm,
-        "d_grad_norm": discriminator_grad_norm,
-        "step_time": f"{step_time:.2f}s",
-    })
+        progress_bar.set_postfix(
+            {
+                "g_loss": f"{generator_loss:.4f}",
+                "d_loss": f"{discriminator_loss:.4f}",
+                "g_grad_norm": generator_grad_norm,
+                "d_grad_norm": discriminator_grad_norm,
+                "step_time": f"{step_time:.2f}s",
+            }
+        )
         progress_bar.update(1)
         if rank <= 0:
-            wandb.log({
-            "generator_loss": generator_loss,
-            "discriminator_loss": discriminator_loss,
-            "generator_grad_norm": generator_grad_norm,
-            "discriminator_grad_norm": discriminator_grad_norm,
-            "learning_rate": lr_scheduler.get_last_lr()[0],
-            "step_time": step_time,
-            "avg_step_time": avg_step_time,
-        }, step=step)
-        if step  % args.checkpointing_steps == 0:
+            wandb.log(
+                {
+                    "generator_loss": generator_loss,
+                    "discriminator_loss": discriminator_loss,
+                    "generator_grad_norm": generator_grad_norm,
+                    "discriminator_grad_norm": discriminator_grad_norm,
+                    "learning_rate": lr_scheduler.get_last_lr()[0],
+                    "step_time": step_time,
+                    "avg_step_time": avg_step_time,
+                },
+                step=step,
+            )
+        if step % args.checkpointing_steps == 0:
             main_print(f"--> saving checkpoint at step {step}")
             if args.use_lora:
                 # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
+                save_lora_checkpoint(
+                    transformer, optimizer, rank, args.output_dir, step
+                )
             else:
                 # Your existing checkpoint saving code
-                save_checkpoint_generator_discriminator(transformer, optimizer, discriminator, discriminator_optimizer, rank, args.output_dir, step)
+                save_checkpoint_generator_discriminator(
+                    transformer,
+                    optimizer,
+                    discriminator,
+                    discriminator_optimizer,
+                    rank,
+                    args.output_dir,
+                    step,
+                )
             main_print(f"--> checkpoint saved at step {step}")
             dist.barrier()
-        if args.log_validation and step  % args.validation_steps == 0:
-            log_validation(args, transformer, device,
-                            torch.bfloat16, step, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps,linear_quadratic_threshold=args.linear_quadratic_threshold, ema=False)
+        if args.log_validation and step % args.validation_steps == 0:
+            log_validation(
+                args,
+                transformer,
+                device,
+                torch.bfloat16,
+                step,
+                scheduler_type=args.scheduler_type,
+                shift=args.shift,
+                num_euler_timesteps=args.num_euler_timesteps,
+                linear_quadratic_threshold=args.linear_quadratic_threshold,
+                ema=False,
+            )
 
     if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
+        save_lora_checkpoint(
+            transformer, optimizer, rank, args.output_dir, args.max_train_steps
+        )
     else:
-        save_checkpoint(transformer,  optimizer, rank, args.output_dir, args.max_train_steps)
-        save_checkpoint(discriminator, discriminator_optimizer, rank, args.output_dir, step, discriminator=True)
-        
+        save_checkpoint(
+            transformer, optimizer, rank, args.output_dir, args.max_train_steps
+        )
+        save_checkpoint(
+            discriminator,
+            discriminator_optimizer,
+            rank,
+            args.output_dir,
+            step,
+            discriminator=True,
+        )
+
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
 
@@ -574,111 +685,252 @@ if __name__ == "__main__":
     # dataset & dataloader
     parser.add_argument("--data_json_path", type=str, required=True)
     parser.add_argument("--num_frames", type=int, default=163)
-    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("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true") # TODO
-    parser.add_argument("--group_resolution", action="store_true") # TODO
+    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(
+        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
+    )
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
 
     # text encoder & vae & diffusion model
     parser.add_argument("--pretrained_model_name_or_path", type=str)
     parser.add_argument("--dit_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
 
     # diffusion setting
     parser.add_argument("--ema_decay", type=float, default=0.999)
     parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument('--cfg', type=float, default=0.1)
-    parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.")
-    
+    parser.add_argument("--cfg", type=float, default=0.1)
+    parser.add_argument(
+        "--precondition_outputs",
+        action="store_true",
+        help="Whether to precondition the outputs of the model.",
+    )
+
     # validation & logs
     parser.add_argument("--validation_prompt_dir", type=str)
     parser.add_argument("--validation_sampling_steps", type=int, default=64)
-    parser.add_argument('--validation_guidance_scale', type=float, default=4.5)
-    parser.add_argument('--validation_steps', type=float, default=64)
+    parser.add_argument("--validation_guidance_scale", type=float, default=4.5)
+    parser.add_argument("--validation_steps", type=float, default=64)
     parser.add_argument("--log_validation", action="store_true")
     parser.add_argument("--tracker_project_name", type=str, default=None)
-    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("--shift", type=float, default=1.0 )
-    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("--resume_from_lora_checkpoint", type=str, default=None,
-                        help=(
-                            "Whether training should be resumed from a previous lora 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(
+        "--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("--shift", type=float, default=1.0)
+    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(
+        "--resume_from_lora_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous lora 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***."
+        ),
+    )
 
     # optimizer & scheduler & Training
     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("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")
-    parser.add_argument("--discriminator_learning_rate", type=float, default=1e-5, 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=10, help="Number of steps for the warmup in the lr scheduler.")
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--discriminator_learning_rate",
+        type=float,
+        default=1e-5,
+        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=10,
+        help="Number of steps for the warmup in the lr scheduler.",
+    )
+    parser.add_argument(
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
     parser.add_argument("--selective_checkpointing", type=float, default=1.0)
-    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("--use_cpu_offload", action="store_true", help="Whether to use CPU offload for param & gradient & optimizer states.")
+    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(
+        "--use_cpu_offload",
+        action="store_true",
+        help="Whether to use CPU offload for param & gradient & optimizer states.",
+    )
 
     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")
+    parser.add_argument(
+        "--train_sp_batch_size",
+        type=int,
+        default=1,
+        help="Batch size for sequence parallel training",
+    )
 
-    parser.add_argument("--use_lora", action="store_true", default=False, help="Whether to use LoRA for finetuning.") 
-    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        default=False,
+        help="Whether to use LoRA for finetuning.",
+    )
+    parser.add_argument(
+        "--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA."
+    )
+    parser.add_argument(
+        "--lora_rank", type=int, default=128, help="LoRA rank parameter. "
+    )
     parser.add_argument("--fsdp_sharding_startegy", default="full")
-    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
 
     # lr_scheduler
-    parser.add_argument("--lr_scheduler", type=str, default="constant",
+    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("--num_euler_timesteps", type=int, default=100)
-    parser.add_argument("--lr_num_cycles", type=int, default=1, help="Number of cycles in the learning rate scheduler.")
-    parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.",)
-    parser.add_argument("--not_apply_cfg_solver", action="store_true", help="Whether to apply the cfg_solver.")
-    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of cycles in the learning rate scheduler.",
+    )
+    parser.add_argument(
+        "--lr_power",
+        type=float,
+        default=1.0,
+        help="Power factor of the polynomial scheduler.",
+    )
+    parser.add_argument(
+        "--not_apply_cfg_solver",
+        action="store_true",
+        help="Whether to apply the cfg_solver.",
+    )
+    parser.add_argument(
+        "--distill_cfg", type=float, default=3.0, help="Distillation coefficient."
+    )
     # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
-    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
-    parser.add_argument("--adv_weight", type=float, default=0.1, help="The weight of the adversarial loss.")
-    parser.add_argument("--discriminator_head_stride", type=int, default=2, help="The stride of the discriminator head.")
-    parser.add_argument("--linear_quadratic_threshold", type=float, default=0.025, help="The threshold of the linear quadratic scheduler.")
-    parser.add_argument("--master_weight_type", type=str, default="fp32", help="Weight type to use - fp32 or bf16.")
+    parser.add_argument(
+        "--scheduler_type", type=str, default="pcm", help="The scheduler type to use."
+    )
+    parser.add_argument(
+        "--adv_weight",
+        type=float,
+        default=0.1,
+        help="The weight of the adversarial loss.",
+    )
+    parser.add_argument(
+        "--discriminator_head_stride",
+        type=int,
+        default=2,
+        help="The stride of the discriminator head.",
+    )
+    parser.add_argument(
+        "--linear_quadratic_threshold",
+        type=float,
+        default=0.025,
+        help="The threshold of the linear quadratic scheduler.",
+    )
+    parser.add_argument(
+        "--master_weight_type",
+        type=str,
+        default="fp32",
+        help="Weight type to use - fp32 or bf16.",
+    )
     args = parser.parse_args()
-    main(args)
\ No newline at end of file
+    main(args)
diff --git a/fastvideo/fsdp_util.py b/fastvideo/fsdp_util.py
index 6a9e9eadb..96a459767 100644
--- a/fastvideo/fsdp_util.py
+++ b/fastvideo/fsdp_util.py
@@ -58,27 +58,30 @@ def apply_fsdp_checkpointing(model, p=1):
                 cut_off += 1
                 return True
         return False
-    
-    
+
     apply_activation_checkpointing(
-        model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=selective_checkpointing
+        model,
+        checkpoint_wrapper_fn=non_reentrant_wrapper,
+        check_fn=selective_checkpointing,
     )
 
-def get_mixed_precision(master_weight_type='fp32'):
-    weight_type = torch.float32 if master_weight_type == 'fp32' else torch.bfloat16
+
+def get_mixed_precision(master_weight_type="fp32"):
+    weight_type = torch.float32 if master_weight_type == "fp32" else torch.bfloat16
     mixed_precision = MixedPrecision(
         param_dtype=weight_type,
         # Gradient communication precision.
         reduce_dtype=weight_type,
         # Buffer precision.
         buffer_dtype=weight_type,
-        cast_forward_inputs=False
+        cast_forward_inputs=False,
     )
     return mixed_precision
 
 
-
-def get_dit_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False, master_weight_type='fp32'):
+def get_dit_fsdp_kwargs(
+    sharding_strategy, use_lora=False, cpu_offload=False, master_weight_type="fp32"
+):
     if use_lora:
         auto_wrap_policy = fsdp_auto_wrap_policy
     else:
@@ -88,10 +91,10 @@ def get_dit_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False, m
                 MochiTransformerBlock,
             },
         )
-    
+
     # we use float32 for fsdp but autocast during training
     mixed_precision = get_mixed_precision(master_weight_type)
-    
+
     if sharding_strategy == "full":
         sharding_strategy = ShardingStrategy.FULL_SHARD
     elif sharding_strategy == "hybrid_full":
@@ -100,10 +103,12 @@ def get_dit_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False, m
         sharding_strategy = ShardingStrategy.NO_SHARD
         auto_wrap_policy = None
     elif sharding_strategy == "hybrid_zero2":
-        sharding_strategy  = ShardingStrategy._HYBRID_SHARD_ZERO2
-    
+        sharding_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2
+
     device_id = torch.cuda.current_device()
-    cpu_offload=torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None
+    cpu_offload = (
+        torch.distributed.fsdp.CPUOffload(offload_params=True) if cpu_offload else None
+    )
     fsdp_kwargs = {
         "auto_wrap_policy": auto_wrap_policy,
         "mixed_precision": mixed_precision,
@@ -112,29 +117,26 @@ def get_dit_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False, m
         "limit_all_gathers": True,
         "cpu_offload": cpu_offload,
     }
-    
+
     # Add LoRA-specific settings when LoRA is enabled
     if use_lora:
-        fsdp_kwargs.update({
-            "use_orig_params": False,  # Required for LoRA memory savings
-            "sync_module_states": True,
-        })
-    
+        fsdp_kwargs.update(
+            {
+                "use_orig_params": False,  # Required for LoRA memory savings
+                "sync_module_states": True,
+            }
+        )
+
     return fsdp_kwargs
-    
-    
-        
-        
 
-def get_discriminator_fsdp_kwargs(master_weight_type='fp32'):
 
+def get_discriminator_fsdp_kwargs(master_weight_type="fp32"):
     auto_wrap_policy = None
 
-
     # Use existing mixed precision settings
 
     mixed_precision = get_mixed_precision(master_weight_type)
-    sharding_strategy  = ShardingStrategy.NO_SHARD
+    sharding_strategy = ShardingStrategy.NO_SHARD
     device_id = torch.cuda.current_device()
     fsdp_kwargs = {
         "auto_wrap_policy": auto_wrap_policy,
@@ -143,8 +145,5 @@ def get_discriminator_fsdp_kwargs(master_weight_type='fp32'):
         "device_id": device_id,
         "limit_all_gathers": True,
     }
-    
+
     return fsdp_kwargs
-    
-    
-        
\ No newline at end of file
diff --git a/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py b/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py
new file mode 100644
index 000000000..afc2b671f
--- /dev/null
+++ b/fastvideo/models/mochi_hf/convert_diffusers_to_mochi.py
@@ -0,0 +1,431 @@
+import torch
+import argparse
+from safetensors.torch import save_file
+import os
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--diffusers_path", required=True, type=str)
+parser.add_argument("--transformer_path", type=str, default=None, help="Path to save transformer model")
+parser.add_argument("--vae_encoder_path", type=str, default=None, help="Path to save VAE encoder model")
+parser.add_argument("--vae_decoder_path", type=str, default=None, help="Path to save VAE decoder model")
+
+args = parser.parse_args()
+
+def reverse_scale_shift(weight, dim):
+    scale, shift = weight.chunk(2, dim=0)
+    new_weight = torch.cat([shift, scale], dim=0)
+    return new_weight
+
+def reverse_proj_gate(weight):
+    gate, proj = weight.chunk(2, dim=0)
+    new_weight = torch.cat([proj, gate], dim=0)
+    return new_weight
+
+def convert_diffusers_transformer_to_mochi(state_dict):
+    original_state_dict = state_dict.copy()
+    new_state_dict = {}
+
+    # Convert patch_embed
+    new_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("patch_embed.proj.weight")
+    new_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("patch_embed.proj.bias")
+
+    # Convert time_embed
+    new_state_dict["t_embedder.mlp.0.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.weight")
+    new_state_dict["t_embedder.mlp.0.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_1.bias")
+    new_state_dict["t_embedder.mlp.2.weight"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.weight")
+    new_state_dict["t_embedder.mlp.2.bias"] = original_state_dict.pop("time_embed.timestep_embedder.linear_2.bias")
+    new_state_dict["t5_y_embedder.to_kv.weight"] = original_state_dict.pop("time_embed.pooler.to_kv.weight")
+    new_state_dict["t5_y_embedder.to_kv.bias"] = original_state_dict.pop("time_embed.pooler.to_kv.bias")
+    new_state_dict["t5_y_embedder.to_q.weight"] = original_state_dict.pop("time_embed.pooler.to_q.weight")
+    new_state_dict["t5_y_embedder.to_q.bias"] = original_state_dict.pop("time_embed.pooler.to_q.bias")
+    new_state_dict["t5_y_embedder.to_out.weight"] = original_state_dict.pop("time_embed.pooler.to_out.weight")
+    new_state_dict["t5_y_embedder.to_out.bias"] = original_state_dict.pop("time_embed.pooler.to_out.bias")
+    new_state_dict["t5_yproj.weight"] = original_state_dict.pop("time_embed.caption_proj.weight")
+    new_state_dict["t5_yproj.bias"] = original_state_dict.pop("time_embed.caption_proj.bias")
+
+    # Convert transformer blocks
+    num_layers = 48
+    for i in range(num_layers):
+        block_prefix = f"transformer_blocks.{i}."
+        new_prefix = f"blocks.{i}."
+
+        # norm1
+        new_state_dict[new_prefix + "mod_x.weight"] = original_state_dict.pop(block_prefix + "norm1.linear.weight")
+        new_state_dict[new_prefix + "mod_x.bias"] = original_state_dict.pop(block_prefix + "norm1.linear.bias")
+
+        if i < num_layers - 1:
+            new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(
+                block_prefix + "norm1_context.linear.weight"
+            )
+            new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(
+                block_prefix + "norm1_context.linear.bias"
+            )
+        else:
+            new_state_dict[new_prefix + "mod_y.weight"] = original_state_dict.pop(
+                block_prefix + "norm1_context.linear_1.weight"
+            )
+            new_state_dict[new_prefix + "mod_y.bias"] = original_state_dict.pop(
+                block_prefix + "norm1_context.linear_1.bias"
+            )
+
+        # Visual attention
+        q = original_state_dict.pop(block_prefix + "attn1.to_q.weight")
+        k = original_state_dict.pop(block_prefix + "attn1.to_k.weight") 
+        v = original_state_dict.pop(block_prefix + "attn1.to_v.weight")
+        qkv_weight = torch.cat([q, k, v], dim=0)
+        new_state_dict[new_prefix + "attn.qkv_x.weight"] = qkv_weight
+
+        new_state_dict[new_prefix + "attn.q_norm_x.weight"] = original_state_dict.pop(
+            block_prefix + "attn1.norm_q.weight"
+        )
+        new_state_dict[new_prefix + "attn.k_norm_x.weight"] = original_state_dict.pop(
+            block_prefix + "attn1.norm_k.weight"
+        )
+        new_state_dict[new_prefix + "attn.proj_x.weight"] = original_state_dict.pop(
+            block_prefix + "attn1.to_out.0.weight"
+        )
+        new_state_dict[new_prefix + "attn.proj_x.bias"] = original_state_dict.pop(
+            block_prefix + "attn1.to_out.0.bias"
+        )
+
+        # Context attention
+        q = original_state_dict.pop(block_prefix + "attn1.add_q_proj.weight")
+        k = original_state_dict.pop(block_prefix + "attn1.add_k_proj.weight")
+        v = original_state_dict.pop(block_prefix + "attn1.add_v_proj.weight")
+        qkv_weight = torch.cat([q, k, v], dim=0)
+        new_state_dict[new_prefix + "attn.qkv_y.weight"] = qkv_weight
+
+        new_state_dict[new_prefix + "attn.q_norm_y.weight"] = original_state_dict.pop(
+            block_prefix + "attn1.norm_added_q.weight"
+        )
+        new_state_dict[new_prefix + "attn.k_norm_y.weight"] = original_state_dict.pop(
+            block_prefix + "attn1.norm_added_k.weight"
+        )
+        if i < num_layers - 1:
+            new_state_dict[new_prefix + "attn.proj_y.weight"] = original_state_dict.pop(
+                block_prefix + "attn1.to_add_out.weight"
+            )
+            new_state_dict[new_prefix + "attn.proj_y.bias"] = original_state_dict.pop(
+                block_prefix + "attn1.to_add_out.bias"
+            )
+
+        # MLP
+        new_state_dict[new_prefix + "mlp_x.w1.weight"] = reverse_proj_gate(
+            original_state_dict.pop(block_prefix + "ff.net.0.proj.weight")
+        )
+        new_state_dict[new_prefix + "mlp_x.w2.weight"] = original_state_dict.pop(block_prefix + "ff.net.2.weight")
+        if i < num_layers - 1:
+            new_state_dict[new_prefix + "mlp_y.w1.weight"] = reverse_proj_gate(
+                original_state_dict.pop(block_prefix + "ff_context.net.0.proj.weight")
+            )
+            new_state_dict[new_prefix + "mlp_y.w2.weight"] = original_state_dict.pop(
+                block_prefix + "ff_context.net.2.weight"
+            )
+
+    # Output layers
+    new_state_dict["final_layer.mod.weight"] = reverse_scale_shift(
+        original_state_dict.pop("norm_out.linear.weight"), dim=0
+    )
+    new_state_dict["final_layer.mod.bias"] = reverse_scale_shift(
+        original_state_dict.pop("norm_out.linear.bias"), dim=0
+    )
+    new_state_dict["final_layer.linear.weight"] = original_state_dict.pop("proj_out.weight")
+    new_state_dict["final_layer.linear.bias"] = original_state_dict.pop("proj_out.bias")
+
+    new_state_dict["pos_frequencies"] = original_state_dict.pop("pos_frequencies")
+
+    print("Remaining Keys:", original_state_dict.keys())
+
+    return new_state_dict
+
+def convert_diffusers_vae_to_mochi(state_dict):
+    original_state_dict = state_dict.copy()
+    encoder_state_dict = {}
+    decoder_state_dict = {}
+
+    # Convert encoder
+    prefix = "encoder."
+
+    encoder_state_dict["layers.0.weight"] = original_state_dict.pop(f"{prefix}proj_in.weight")
+    encoder_state_dict["layers.0.bias"] = original_state_dict.pop(f"{prefix}proj_in.bias")
+
+    # Convert block_in
+    for i in range(3):
+        encoder_state_dict[f"layers.{i+1}.stack.0.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.0.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.2.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv1.conv.weight"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.2.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv1.conv.bias"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.3.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.3.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.5.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv2.conv.weight"
+        )
+        encoder_state_dict[f"layers.{i+1}.stack.5.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv2.conv.bias"
+        )
+
+    # Convert down_blocks
+    down_block_layers = [3, 4, 6]
+    for block in range(3):
+        encoder_state_dict[f"layers.{block+4}.layers.0.weight"] = original_state_dict.pop(
+            f"{prefix}down_blocks.{block}.conv_in.conv.weight"
+        )
+        encoder_state_dict[f"layers.{block+4}.layers.0.bias"] = original_state_dict.pop(
+            f"{prefix}down_blocks.{block}.conv_in.conv.bias"
+        )
+
+        for i in range(down_block_layers[block]):
+            # Convert resnets
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.weight"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.0.bias"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.weight"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.2.bias"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.weight"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.3.bias"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.weight"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.stack.5.bias"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias"
+            )
+
+            # Convert attentions
+            q = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight")
+            k = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight")
+            v = original_state_dict.pop(f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight")
+            qkv_weight = torch.cat([q, k, v], dim=0)
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.qkv.weight"] = qkv_weight
+
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.weight"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.bias"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.weight"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight"
+            )
+            encoder_state_dict[f"layers.{block+4}.layers.{i+1}.attn_block.norm.bias"] = original_state_dict.pop(
+                f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias"
+            )
+
+    # Convert block_out
+    for i in range(3):
+        encoder_state_dict[f"layers.{i+7}.stack.0.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.0.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.2.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv1.conv.weight"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.2.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv1.conv.bias"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.3.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.3.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.5.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv2.conv.weight"
+        )
+        encoder_state_dict[f"layers.{i+7}.stack.5.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv2.conv.bias"
+        )
+
+        q = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_q.weight")
+        k = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_k.weight")
+        v = original_state_dict.pop(f"{prefix}block_out.attentions.{i}.to_v.weight")
+        qkv_weight = torch.cat([q, k, v], dim=0)
+        encoder_state_dict[f"layers.{i+7}.attn_block.attn.qkv.weight"] = qkv_weight
+
+        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.attentions.{i}.to_out.0.weight"
+        )
+        encoder_state_dict[f"layers.{i+7}.attn_block.attn.out.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.attentions.{i}.to_out.0.bias"
+        )
+        encoder_state_dict[f"layers.{i+7}.attn_block.norm.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.norms.{i}.norm_layer.weight"
+        )
+        encoder_state_dict[f"layers.{i+7}.attn_block.norm.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.norms.{i}.norm_layer.bias"
+        )
+
+    # Convert output layers
+    encoder_state_dict["output_norm.weight"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.weight")
+    encoder_state_dict["output_norm.bias"] = original_state_dict.pop(f"{prefix}norm_out.norm_layer.bias")
+    encoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight")
+
+    # Convert decoder
+    prefix = "decoder."
+
+    decoder_state_dict["blocks.0.0.weight"] = original_state_dict.pop(f"{prefix}conv_in.weight")
+    decoder_state_dict["blocks.0.0.bias"] = original_state_dict.pop(f"{prefix}conv_in.bias")
+
+    # Convert block_in
+    for i in range(3):
+        decoder_state_dict[f"blocks.0.{i+1}.stack.0.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.0.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.2.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv1.conv.weight"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.2.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv1.conv.bias"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.3.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.3.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.5.weight"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv2.conv.weight"
+        )
+        decoder_state_dict[f"blocks.0.{i+1}.stack.5.bias"] = original_state_dict.pop(
+            f"{prefix}block_in.resnets.{i}.conv2.conv.bias"
+        )
+
+    # Convert up_blocks
+    up_block_layers = [6, 4, 3]
+    for block in range(3):
+        for i in range(up_block_layers[block]):
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.weight"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.0.bias"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.weight"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.2.bias"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.weight"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.3.bias"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.weight"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight"
+            )
+            decoder_state_dict[f"blocks.{block+1}.blocks.{i}.stack.5.bias"] = original_state_dict.pop(
+                f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias"
+            )
+        decoder_state_dict[f"blocks.{block+1}.proj.weight"] = original_state_dict.pop(
+            f"{prefix}up_blocks.{block}.proj.weight"
+        )
+        decoder_state_dict[f"blocks.{block+1}.proj.bias"] = original_state_dict.pop(
+            f"{prefix}up_blocks.{block}.proj.bias"
+        )
+
+    # Convert block_out
+    for i in range(3):
+        decoder_state_dict[f"blocks.4.{i}.stack.0.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.0.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.2.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv1.conv.weight"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.2.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv1.conv.bias"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.3.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.3.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.5.weight"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv2.conv.weight"
+        )
+        decoder_state_dict[f"blocks.4.{i}.stack.5.bias"] = original_state_dict.pop(
+            f"{prefix}block_out.resnets.{i}.conv2.conv.bias"
+        )
+
+    # Convert output layers  
+    decoder_state_dict["output_proj.weight"] = original_state_dict.pop(f"{prefix}proj_out.weight")
+    decoder_state_dict["output_proj.bias"] = original_state_dict.pop(f"{prefix}proj_out.bias")
+
+    return encoder_state_dict, decoder_state_dict
+
+def ensure_safetensors_extension(path):
+    if not path.endswith('.safetensors'):
+        path = path + '.safetensors'
+    return path
+
+def ensure_directory_exists(path):
+    directory = os.path.dirname(path)
+    if directory:
+        os.makedirs(directory, exist_ok=True)
+
+def main(args):
+    from diffusers import MochiPipeline
+    
+    pipe = MochiPipeline.from_pretrained(args.diffusers_path)
+
+    if args.transformer_path:
+        transformer_path = ensure_safetensors_extension(args.transformer_path)
+        ensure_directory_exists(transformer_path)
+
+        print(f"Converting transformer model...")
+        transformer_state_dict = convert_diffusers_transformer_to_mochi(pipe.transformer.state_dict())
+        save_file(transformer_state_dict, transformer_path)
+        print(f"Saved transformer to {transformer_path}")
+
+    if args.vae_encoder_path and args.vae_decoder_path:
+        encoder_path = ensure_safetensors_extension(args.vae_encoder_path)
+        decoder_path = ensure_safetensors_extension(args.vae_decoder_path)
+
+        ensure_directory_exists(encoder_path)
+        ensure_directory_exists(decoder_path)
+
+        print(f"Converting VAE models...")
+        encoder_state_dict, decoder_state_dict = convert_diffusers_vae_to_mochi(pipe.vae.state_dict())
+
+        save_file(encoder_state_dict, encoder_path)
+        print(f"Saved VAE encoder to {encoder_path}")
+
+        save_file(decoder_state_dict, decoder_path)
+        print(f"Saved VAE decoder to {decoder_path}")
+    elif args.vae_encoder_path or args.vae_decoder_path:
+        print("Warning: Both VAE encoder and decoder paths must be specified to convert VAE models.")
+
+if __name__ == "__main__":
+    main(args)
\ No newline at end of file
diff --git a/fastvideo/models/mochi_hf/mochi_latents_utils.py b/fastvideo/models/mochi_hf/mochi_latents_utils.py
index 0919e4b1b..48a57f5d4 100644
--- a/fastvideo/models/mochi_hf/mochi_latents_utils.py
+++ b/fastvideo/models/mochi_hf/mochi_latents_utils.py
@@ -1,39 +1,42 @@
 import torch
-mochi_latents_mean = torch.tensor([
-    -0.06730895953510081,
-    -0.038011381506090416,
-    -0.07477820912866141,
-    -0.05565264470995561,
-    0.012767231469026969,
-    -0.04703542746246419,
-    0.043896967884726704,
-    -0.09346305707025976,
-    -0.09918314763016893,
-    -0.008729793427399178,
-    -0.011931556316503654,
-    -0.0321993391887285
-  ]).view(1, 12, 1, 1, 1)
-mochi_latents_std = torch.tensor([
-    0.9263795028493863,
-    0.9248894543193766,
-    0.9393059390890617,
-    0.959253732819592,
-    0.8244560132752793,
-    0.917259975397747,
-    0.9294154431013696,
-    1.3720942357788521,
-    0.881393668867029,
-    0.9168315692124348,
-    0.9185249279345552,
-    0.9274757570805041
-    ]).view(1, 12, 1, 1, 1)
+
+mochi_latents_mean = torch.tensor(
+    [
+        -0.06730895953510081,
+        -0.038011381506090416,
+        -0.07477820912866141,
+        -0.05565264470995561,
+        0.012767231469026969,
+        -0.04703542746246419,
+        0.043896967884726704,
+        -0.09346305707025976,
+        -0.09918314763016893,
+        -0.008729793427399178,
+        -0.011931556316503654,
+        -0.0321993391887285,
+    ]
+).view(1, 12, 1, 1, 1)
+mochi_latents_std = torch.tensor(
+    [
+        0.9263795028493863,
+        0.9248894543193766,
+        0.9393059390890617,
+        0.959253732819592,
+        0.8244560132752793,
+        0.917259975397747,
+        0.9294154431013696,
+        1.3720942357788521,
+        0.881393668867029,
+        0.9168315692124348,
+        0.9185249279345552,
+        0.9274757570805041,
+    ]
+).view(1, 12, 1, 1, 1)
 mochi_scaling_factor = 1.0
 
-   
-def normalize_mochi_dit_input(latents):
-  latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
-  latents_std = mochi_latents_std.to(latents.device, latents.dtype)
-  latents = (latents - latents_mean) / latents_std
-  return latents
 
-        
\ No newline at end of file
+def normalize_mochi_dit_input(latents):
+    latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
+    latents_std = mochi_latents_std.to(latents.device, latents.dtype)
+    latents = (latents - latents_mean) / latents_std
+    return latents
diff --git a/fastvideo/models/mochi_hf/modeling_mochi.py b/fastvideo/models/mochi_hf/modeling_mochi.py
index 2064278fe..2d674be8d 100644
--- a/fastvideo/models/mochi_hf/modeling_mochi.py
+++ b/fastvideo/models/mochi_hf/modeling_mochi.py
@@ -19,15 +19,29 @@ import torch.nn as nn
 import diffusers
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.utils import is_torch_version, logging
-from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_version,
+    logging,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
 from diffusers.utils.torch_utils import maybe_allow_in_graph
 from diffusers.models.attention import FeedForward as HF_FeedForward
 from diffusers.models.attention_processor import Attention
-from diffusers.models.embeddings import MochiCombinedTimestepCaptionEmbedding, PatchEmbed
+from diffusers.models.embeddings import (
+    MochiCombinedTimestepCaptionEmbedding,
+    PatchEmbed,
+)
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.loaders import PeftAdapterMixin
-from fastvideo.models.mochi_hf.norm import  MochiLayerNormContinuous, MochiRMSNormZero, MochiModulatedRMSNorm, MochiRMSNorm
+from fastvideo.models.mochi_hf.norm import (
+    MochiLayerNormContinuous,
+    MochiRMSNormZero,
+    MochiModulatedRMSNorm,
+    MochiRMSNorm,
+)
 from diffusers.models.normalization import AdaLayerNormContinuous
 
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
@@ -43,6 +57,8 @@ from flash_attn.bert_padding import pad_input, unpad_input
 from liger_kernel.ops.swiglu import LigerSiLUMulFunction
 from sageattention import sageattn
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
 class FeedForward(HF_FeedForward):
     def __init__(
         self,
@@ -55,37 +71,50 @@ class FeedForward(HF_FeedForward):
         inner_dim=None,
         bias: bool = True,
     ):
-        super().__init__(dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias)
+        super().__init__(
+            dim, dim_out, mult, dropout, activation_fn, final_dropout, inner_dim, bias
+        )
         assert activation_fn == "swiglu"
-        
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.net[0].proj(hidden_states)
         hidden_states, gate = hidden_states.chunk(2, dim=-1)
-        
-        return self.net[2](
-            LigerSiLUMulFunction.apply(gate, hidden_states)
-        )
 
-def flash_attn_no_pad(qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None):  
+        return self.net[2](LigerSiLUMulFunction.apply(gate, hidden_states))
+
+
+def flash_attn_no_pad(
+    qkv, key_padding_mask, causal=False, dropout_p=0.0, softmax_scale=None
+):
     # adapted from https://github.com/Dao-AILab/flash-attention/blob/13403e81157ba37ca525890f2f0f2137edf75311/flash_attn/flash_attention.py#L27
     batch_size = qkv.shape[0]
     seqlen = qkv.shape[1]
     nheads = qkv.shape[-2]
-    x = rearrange(qkv, 'b s three h d -> b s (three h d)')
-    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(x, key_padding_mask)
-    
+    x = rearrange(qkv, "b s three h d -> b s (three h d)")
+    x_unpad, indices, cu_seqlens, max_s, used_seqlens_in_batch = unpad_input(
+        x, key_padding_mask
+    )
 
-    x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+    x_unpad = rearrange(x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads)
     output_unpad = flash_attn_varlen_qkvpacked_func(
-        x_unpad, cu_seqlens, max_s, dropout_p,
-        softmax_scale=softmax_scale, causal=causal
+        x_unpad,
+        cu_seqlens,
+        max_s,
+        dropout_p,
+        softmax_scale=softmax_scale,
+        causal=causal,
+    )
+    output = rearrange(
+        pad_input(
+            rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, batch_size, seqlen
+        ),
+        "b s (h d) -> b s h d",
+        h=nheads,
     )
-    output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
-                                indices, batch_size, seqlen),
-                    'b s (h d) -> b s h d', h=nheads)
     return output
 
-class MochiAttention(nn.Module):          
+
+class MochiAttention(nn.Module):
     def __init__(
         self,
         query_dim: int,
@@ -119,17 +148,25 @@ class MochiAttention(nn.Module):
         self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
         self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
 
-        self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
-        self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+        self.add_k_proj = nn.Linear(
+            added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
+        )
+        self.add_v_proj = nn.Linear(
+            added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
+        )
         if self.context_pre_only is not None:
-            self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+            self.add_q_proj = nn.Linear(
+                added_kv_proj_dim, self.inner_dim, bias=added_proj_bias
+            )
 
         self.to_out = nn.ModuleList([])
         self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
         self.to_out.append(nn.Dropout(dropout))
 
         if not self.context_pre_only:
-            self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+            self.to_add_out = nn.Linear(
+                self.inner_dim, self.out_context_dim, bias=out_bias
+            )
 
         self.processor = processor
 
@@ -190,7 +227,9 @@ class MochiAttnProcessor2_0:
 
     def __init__(self):
         if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
+            raise ImportError(
+                "MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0."
+            )
 
     def __call__(
         self,
@@ -211,12 +250,11 @@ class MochiAttnProcessor2_0:
         key = key.unflatten(2, (attn.heads, -1))
         value = value.unflatten(2, (attn.heads, -1))
 
-
         if attn.norm_q is not None:
             query = attn.norm_q(query)
         if attn.norm_k is not None:
             key = attn.norm_k(key)
-        # [b, 256, h * d] 
+        # [b, 256, h * d]
         encoder_query = attn.add_q_proj(encoder_hidden_states)
         encoder_key = attn.add_k_proj(encoder_hidden_states)
         encoder_value = attn.add_v_proj(encoder_hidden_states)
@@ -225,37 +263,37 @@ class MochiAttnProcessor2_0:
         encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
         encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
         encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
-        
-        
+
         if attn.norm_added_q is not None:
             encoder_query = attn.norm_added_q(encoder_query)
         if attn.norm_added_k is not None:
             encoder_key = attn.norm_added_k(encoder_key)
-            
+
         if image_rotary_emb is not None:
             freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
         # shard the head dimension
         if get_sequence_parallel_state():
             # B, S, H, D to (S, B,) H, D
-            # batch_size, seq_len, attn_heads, head_dim 
+            # batch_size, seq_len, attn_heads, head_dim
             query = all_to_all_4D(query, scatter_dim=2, gather_dim=1)
-            key = all_to_all_4D(key,  scatter_dim=2, gather_dim=1)
+            key = all_to_all_4D(key, scatter_dim=2, gather_dim=1)
             value = all_to_all_4D(value, scatter_dim=2, gather_dim=1)
 
-            
             def shrink_head(encoder_state, dim):
                 local_heads = encoder_state.shape[dim] // nccl_info.sp_size
-                return encoder_state.narrow(dim, nccl_info.rank_within_group * local_heads, local_heads)
+                return encoder_state.narrow(
+                    dim, nccl_info.rank_within_group * local_heads, local_heads
+                )
+
             encoder_query = shrink_head(encoder_query, dim=2)
             encoder_key = shrink_head(encoder_key, dim=2)
             encoder_value = shrink_head(encoder_value, dim=2)
             if image_rotary_emb is not None:
                 freqs_cos = shrink_head(freqs_cos, dim=1)
                 freqs_sin = shrink_head(freqs_sin, dim=1)
-    
 
-    
         if image_rotary_emb is not None:
+
             def apply_rotary_emb(x, freqs_cos, freqs_sin):
                 x_even = x[..., 0::2].float()
                 x_odd = x[..., 1::2].float()
@@ -263,9 +301,10 @@ class MochiAttnProcessor2_0:
                 sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
 
                 return torch.stack([cos, sin], dim=-1).flatten(-2)
+
             query = apply_rotary_emb(query, freqs_cos, freqs_sin)
             key = apply_rotary_emb(key, freqs_cos, freqs_sin)
-            
+
         # query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
         # encoder_query, encoder_key, encoder_value = (
         #     encoder_query.transpose(1, 2),
@@ -293,11 +332,11 @@ class MochiAttnProcessor2_0:
             hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None)
         
         # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False)
-        
+
         # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length
         # def no_padding_mask(score, b, h, q_idx, kv_idx):
         #     return torch.where(kv_idx < valid_lengths[b],score,  -float("inf"))
-            
+
         # hidden_states = flex_attention(query, key, value, score_mod=no_padding_mask)
         if get_sequence_parallel_state():
             hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
@@ -305,7 +344,9 @@ class MochiAttnProcessor2_0:
             )
             # B, S, H, D
             hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=2).contiguous()
+            encoder_hidden_states = all_gather(
+                encoder_hidden_states, dim=2
+            ).contiguous()
             hidden_states = hidden_states.flatten(2, 3)
             hidden_states = hidden_states.to(query.dtype)
             encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
@@ -317,8 +358,6 @@ class MochiAttnProcessor2_0:
             hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
                 (sequence_length, encoder_sequence_length), dim=1
             )
-            
-
 
         # linear proj
         hidden_states = attn.to_out[0](hidden_states)
@@ -330,6 +369,7 @@ class MochiAttnProcessor2_0:
 
         return hidden_states, encoder_hidden_states
 
+
 @maybe_allow_in_graph
 class MochiTransformerBlock(nn.Module):
     r"""
@@ -372,7 +412,9 @@ class MochiTransformerBlock(nn.Module):
         self.norm1 = MochiRMSNormZero(dim, 4 * dim, eps=eps, elementwise_affine=False)
 
         if not context_pre_only:
-            self.norm1_context = MochiRMSNormZero(dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False)
+            self.norm1_context = MochiRMSNormZero(
+                dim, 4 * pooled_projection_dim, eps=eps, elementwise_affine=False
+            )
         else:
             self.norm1_context = MochiLayerNormContinuous(
                 embedding_dim=pooled_projection_dim,
@@ -396,12 +438,18 @@ class MochiTransformerBlock(nn.Module):
 
         # TODO(aryan): norm_context layers are not needed when `context_pre_only` is True
         self.norm2 = MochiModulatedRMSNorm(eps=eps)
-        self.norm2_context = MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None
+        self.norm2_context = (
+            MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None
+        )
 
         self.norm3 = MochiModulatedRMSNorm(eps)
-        self.norm3_context = MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None
+        self.norm3_context = (
+            MochiModulatedRMSNorm(eps=eps) if not self.context_pre_only else None
+        )
 
-        self.ff = FeedForward(dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False)
+        self.ff = FeedForward(
+            dim, inner_dim=self.ff_inner_dim, activation_fn=activation_fn, bias=False
+        )
         self.ff_context = None
         if not context_pre_only:
             self.ff_context = FeedForward(
@@ -421,14 +469,19 @@ class MochiTransformerBlock(nn.Module):
         encoder_attention_mask: torch.Tensor,
         temb: torch.Tensor,
         image_rotary_emb: Optional[torch.Tensor] = None,
-        output_attn = False, 
+        output_attn=False,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-        norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+        norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(
+            hidden_states, temb
+        )
 
         if not self.context_pre_only:
-            norm_encoder_hidden_states, enc_gate_msa, enc_scale_mlp, enc_gate_mlp = self.norm1_context(
-                encoder_hidden_states, temb
-            )
+            (
+                norm_encoder_hidden_states,
+                enc_gate_msa,
+                enc_scale_mlp,
+                enc_gate_mlp,
+            ) = self.norm1_context(encoder_hidden_states, temb)
         else:
             norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)
 
@@ -436,20 +489,27 @@ class MochiTransformerBlock(nn.Module):
             hidden_states=norm_hidden_states,
             encoder_hidden_states=norm_encoder_hidden_states,
             image_rotary_emb=image_rotary_emb,
-            encoder_attention_mask=encoder_attention_mask
+            encoder_attention_mask=encoder_attention_mask,
         )
 
-        hidden_states = hidden_states + self.norm2(attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1))
-        norm_hidden_states = self.norm3(hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32)))
+        hidden_states = hidden_states + self.norm2(
+            attn_hidden_states, torch.tanh(gate_msa).unsqueeze(1)
+        )
+        norm_hidden_states = self.norm3(
+            hidden_states, (1 + scale_mlp.unsqueeze(1).to(torch.float32))
+        )
         ff_output = self.ff(norm_hidden_states)
-        hidden_states = hidden_states + self.norm4(ff_output, torch.tanh(gate_mlp).unsqueeze(1))
+        hidden_states = hidden_states + self.norm4(
+            ff_output, torch.tanh(gate_mlp).unsqueeze(1)
+        )
 
         if not self.context_pre_only:
             encoder_hidden_states = encoder_hidden_states + self.norm2_context(
                 context_attn_hidden_states, torch.tanh(enc_gate_msa).unsqueeze(1)
             )
             norm_encoder_hidden_states = self.norm3_context(
-                encoder_hidden_states, (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32))
+                encoder_hidden_states,
+                (1 + enc_scale_mlp.unsqueeze(1).to(torch.float32)),
             )
             context_ff_output = self.ff_context(norm_encoder_hidden_states)
             encoder_hidden_states = encoder_hidden_states + self.norm4_context(
@@ -491,18 +551,22 @@ class MochiRoPE(nn.Module):
     ) -> torch.Tensor:
         scale = (self.target_area / (height * width)) ** 0.5
         t = torch.arange(num_frames * nccl_info.sp_size, device=device, dtype=dtype)
-        h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
+        h = self._centers(
+            -height * scale / 2, height * scale / 2, height, device, dtype
+        )
         w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
 
         grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
 
         positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
         return positions
-    
+
     def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
         with torch.autocast(freqs.device.type, enabled=False):
             # Always run ROPE freqs computation in FP32
-            freqs = torch.einsum("nd,dhf->nhf", pos.to(torch.float32), freqs.to(torch.float32))
+            freqs = torch.einsum(
+                "nd,dhf->nhf", pos.to(torch.float32), freqs.to(torch.float32)
+            )
         freqs_cos = torch.cos(freqs)
         freqs_sin = torch.sin(freqs)
         return freqs_cos, freqs_sin
@@ -589,7 +653,9 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             num_attention_heads=8,
         )
 
-        self.pos_frequencies = nn.Parameter(torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0))
+        self.pos_frequencies = nn.Parameter(
+            torch.full((3, num_attention_heads, attention_head_dim // 2), 0.0)
+        )
         self.rope = MochiRoPE()
 
         self.transformer_blocks = nn.ModuleList(
@@ -608,7 +674,11 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         )
 
         self.norm_out = AdaLayerNormContinuous(
-            inner_dim, inner_dim, elementwise_affine=False, eps=1e-6, norm_type="layer_norm"
+            inner_dim,
+            inner_dim,
+            elementwise_affine=False,
+            eps=1e-6,
+            norm_type="layer_norm",
         )
         self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels)
 
@@ -624,12 +694,14 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
         encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         encoder_attention_mask: torch.Tensor,
-        output_attn = False,
+        output_attn=False,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = False,
     ) -> torch.Tensor:
-        assert return_dict is False, "return_dict is not supported in MochiTransformer3DModel"
-        
+        assert (
+            return_dict is False
+        ), "return_dict is not supported in MochiTransformer3DModel"
+
         if attention_kwargs is not None:
             attention_kwargs = attention_kwargs.copy()
             lora_scale = attention_kwargs.pop("scale", 1.0)
@@ -640,20 +712,26 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
             # weight the lora layers by setting `lora_scale` for each PEFT layer
             scale_lora_layers(self, lora_scale)
         else:
-            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+            if (
+                attention_kwargs is not None
+                and attention_kwargs.get("scale", None) is not None
+            ):
                 logger.warning(
                     "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
                 )
-    
+
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p = self.config.patch_size
 
         post_patch_height = height // p
         post_patch_width = width // p
-        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux 
+        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux
         timestep = 1000 - timestep
         temb, encoder_hidden_states = self.time_embed(
-            timestep, encoder_hidden_states, encoder_attention_mask, hidden_dtype=hidden_states.dtype
+            timestep,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            hidden_dtype=hidden_states.dtype,
         )
 
         hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
@@ -678,8 +756,14 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
 
                     return custom_forward
 
-                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                hidden_states, encoder_hidden_states, attn_outputs = torch.utils.checkpoint.checkpoint(
+                ckpt_kwargs: Dict[str, Any] = (
+                    {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                )
+                (
+                    hidden_states,
+                    encoder_hidden_states,
+                    attn_outputs,
+                ) = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(block),
                     hidden_states,
                     encoder_hidden_states,
@@ -696,24 +780,26 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
                     encoder_attention_mask=encoder_attention_mask,
                     temb=temb,
                     image_rotary_emb=image_rotary_emb,
-                    output_attn = output_attn,
+                    output_attn=output_attn,
                 )
             attn_outputs_list.append(attn_outputs)
 
         hidden_states = self.norm_out(hidden_states, temb)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1)
+        hidden_states = hidden_states.reshape(
+            batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1
+        )
         hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
         output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
-        
+
         if USE_PEFT_BACKEND:
             # remove `lora_scale` from each PEFT layer
             unscale_lora_layers(self, lora_scale)
-            
-        if not output_attn :
-            attn_outputs_list = None 
+
+        if not output_attn:
+            attn_outputs_list = None
         else:
             attn_outputs_list = torch.stack(attn_outputs_list, dim=0)
-        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux 
-        return (-output, attn_outputs_list)
\ No newline at end of file
+        # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux
+        return (-output, attn_outputs_list)
diff --git a/fastvideo/models/mochi_hf/norm.py b/fastvideo/models/mochi_hf/norm.py
index e86fa6d7a..65b034652 100644
--- a/fastvideo/models/mochi_hf/norm.py
+++ b/fastvideo/models/mochi_hf/norm.py
@@ -38,7 +38,7 @@ class MochiModulatedRMSNorm(nn.Module):
         hidden_states = hidden_states.to(hidden_states_dtype)
 
         return hidden_states
-    
+
 
 class MochiRMSNorm(nn.Module):
     def __init__(self, dim, eps: float, elementwise_affine=True):
@@ -63,7 +63,7 @@ class MochiRMSNorm(nn.Module):
         hidden_states = hidden_states.to(hidden_states_dtype)
 
         return hidden_states
-    
+
 
 class MochiLayerNormContinuous(nn.Module):
     def __init__(
@@ -92,7 +92,7 @@ class MochiLayerNormContinuous(nn.Module):
         x = self.norm(x, (1 + scale.unsqueeze(1).to(torch.float32)))
 
         return x.to(input_dtype)
-    
+
 
 class MochiRMSNormZero(nn.Module):
     r"""
@@ -102,7 +102,11 @@ class MochiRMSNormZero(nn.Module):
     """
 
     def __init__(
-        self, embedding_dim: int, hidden_dim: int, eps: float = 1e-5, elementwise_affine: bool = False
+        self,
+        embedding_dim: int,
+        hidden_dim: int,
+        eps: float = 1e-5,
+        elementwise_affine: bool = False,
     ) -> None:
         super().__init__()
 
@@ -118,7 +122,9 @@ class MochiRMSNormZero(nn.Module):
         emb = self.linear(self.silu(emb))
         scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
 
-        hidden_states = self.norm(hidden_states, (1 + scale_msa[:, None].to(torch.float32)))
+        hidden_states = self.norm(
+            hidden_states, (1 + scale_msa[:, None].to(torch.float32))
+        )
         hidden_states = hidden_states.to(hidden_states_dtype)
 
-        return hidden_states, gate_msa, scale_mlp, gate_mlp
\ No newline at end of file
+        return hidden_states, gate_msa, scale_mlp, gate_mlp
diff --git a/fastvideo/models/mochi_hf/pipeline_mochi.py b/fastvideo/models/mochi_hf/pipeline_mochi.py
index c14bf260a..79822635a 100644
--- a/fastvideo/models/mochi_hf/pipeline_mochi.py
+++ b/fastvideo/models/mochi_hf/pipeline_mochi.py
@@ -35,7 +35,7 @@ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
 from einops import rearrange
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-from fastvideo.utils.communications import  all_gather
+from fastvideo.utils.communications import all_gather
 from diffusers.loaders import Mochi1LoraLoaderMixin
 
 if is_torch_xla_available():
@@ -81,14 +81,19 @@ def calculate_shift(
 def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
     if linear_steps is None:
         linear_steps = num_steps // 2
-    linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+    linear_sigma_schedule = [
+        i * threshold_noise / linear_steps for i in range(linear_steps)
+    ]
     threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
     quadratic_steps = num_steps - linear_steps
     quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
-    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
+    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (
+        quadratic_steps**2
+    )
     const = quadratic_coef * (linear_steps**2)
     quadratic_sigma_schedule = [
-        quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
+        quadratic_coef * (i**2) + linear_coef * i + const
+        for i in range(linear_steps, num_steps)
     ]
     sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
     sigma_schedule = [1.0 - x for x in sigma_schedule]
@@ -128,9 +133,13 @@ def retrieve_timesteps(
         second element is the number of inference steps.
     """
     if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+        raise ValueError(
+            "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+        )
     if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        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"
@@ -140,7 +149,9 @@ def retrieve_timesteps(
         timesteps = scheduler.timesteps
         num_inference_steps = len(timesteps)
     elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
         if not accept_sigmas:
             raise ValueError(
                 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
@@ -204,9 +215,13 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         self.vae_temporal_scale_factor = 6
         self.patch_size = 2
 
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor)
+        self.video_processor = VideoProcessor(
+            vae_scale_factor=self.vae_spatial_scale_factor
+        )
         self.tokenizer_max_length = (
-            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+            self.tokenizer.model_max_length
+            if hasattr(self, "tokenizer") and self.tokenizer is not None
+            else 77
         )
         self.default_height = 480
         self.default_width = 848
@@ -238,22 +253,32 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         prompt_attention_mask = text_inputs.attention_mask
         prompt_attention_mask = prompt_attention_mask.bool().to(device)
 
-        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").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_sequence_length - 1 : -1])
+        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_sequence_length - 1 : -1]
+            )
             logger.warning(
                 "The following part of your input was truncated because `max_sequence_length` is set to "
                 f" {max_sequence_length} tokens: {removed_text}"
             )
 
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)[0]
+        prompt_embeds = self.text_encoder(
+            text_input_ids.to(device), attention_mask=prompt_attention_mask
+        )[0]
         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
 
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         _, seq_len, _ = prompt_embeds.shape
         prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+        prompt_embeds = prompt_embeds.view(
+            batch_size * num_videos_per_prompt, seq_len, -1
+        )
 
         prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
         prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1)
@@ -320,7 +345,11 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
 
         if do_classifier_free_guidance and negative_prompt_embeds is None:
             negative_prompt = negative_prompt or ""
-            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt = (
+                batch_size * [negative_prompt]
+                if isinstance(negative_prompt, str)
+                else negative_prompt
+            )
 
             if prompt is not None and type(prompt) is not type(negative_prompt):
                 raise TypeError(
@@ -334,7 +363,10 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
                     " the batch size of `prompt`."
                 )
 
-            negative_prompt_embeds, negative_prompt_attention_mask = self._get_t5_prompt_embeds(
+            (
+                negative_prompt_embeds,
+                negative_prompt_attention_mask,
+            ) = self._get_t5_prompt_embeds(
                 prompt=negative_prompt,
                 num_videos_per_prompt=num_videos_per_prompt,
                 max_sequence_length=max_sequence_length,
@@ -342,7 +374,12 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
                 dtype=dtype,
             )
 
-        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
+        return (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        )
 
     def check_inputs(
         self,
@@ -356,10 +393,13 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         negative_prompt_attention_mask=None,
     ):
         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}.")
+            raise ValueError(
+                f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
+            )
 
         if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+            k in self._callback_tensor_inputs
+            for k in callback_on_step_end_tensor_inputs
         ):
             raise ValueError(
                 f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
@@ -374,14 +414,25 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
             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)}")
+        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_embeds is not None and prompt_attention_mask is None:
-            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
+            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 (
+            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:
@@ -470,7 +521,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
     @property
     def attention_kwargs(self):
         return self._attention_kwargs
-    
+
     @property
     def interrupt(self):
         return self._interrupt
@@ -500,7 +551,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         max_sequence_length: int = 256,
-        return_all_states = False,
+        return_all_states=False,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -628,7 +679,9 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         )
         if self.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)
+            prompt_attention_mask = torch.cat(
+                [negative_prompt_attention_mask, prompt_attention_mask], dim=0
+            )
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.config.in_channels
@@ -645,9 +698,10 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         )
         world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
         if get_sequence_parallel_state():
-            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
+            latents = rearrange(
+                latents, "b t (n s) h w -> b t n s h w", n=world_size
+            ).contiguous()
             latents = latents[:, :, rank, :, :, :]
-            
 
         original_noise = copy.deepcopy(latents)
         # 5. Prepare timestep
@@ -656,7 +710,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
         sigmas = linear_quadratic_schedule(num_inference_steps, threshold_noise)
         sigmas = np.array(sigmas)
         # check if of type FlowMatchEulerDiscreteScheduler
-        if  isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+        if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
             timesteps, num_inference_steps = retrieve_timesteps(
                 self.scheduler,
                 num_inference_steps,
@@ -670,7 +724,9 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
                 num_inference_steps,
                 device,
             )
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
         self._num_timesteps = len(timesteps)
 
         # 6. Denoising loop
@@ -679,7 +735,11 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
                 if self.interrupt:
                     continue
 
-                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                latent_model_input = (
+                    torch.cat([latents] * 2)
+                    if self.do_classifier_free_guidance
+                    else latents
+                )
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
 
@@ -691,16 +751,20 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
                     attention_kwargs=attention_kwargs,
                     return_dict=False,
                 )[0]
-                
+
                 # Mochi CFG + Sampling runs in FP32
                 noise_pred = noise_pred.to(torch.float32)
                 if self.do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
 
                 # compute the previous noisy sample x_t -> x_t-1
                 latents_dtype = latents.dtype
-                latents = self.scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0]
+                latents = self.scheduler.step(
+                    noise_pred, t, latents.to(torch.float32), return_dict=False
+                )[0]
                 latents = latents.to(latents_dtype)
 
                 if latents.dtype != latents_dtype:
@@ -718,7 +782,9 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
                     prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
 
                 # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
                     progress_bar.update()
 
                 if XLA_AVAILABLE:
@@ -726,34 +792,49 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
 
         if get_sequence_parallel_state():
             latents = all_gather(latents, dim=2)
-            #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)
+            # 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 output_type == "latent":
             video = latents
         else:
             # unscale/denormalize the latents
             # denormalize with the mean and std if available and not None
-            has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
-            has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
+            has_latents_mean = (
+                hasattr(self.vae.config, "latents_mean")
+                and self.vae.config.latents_mean is not None
+            )
+            has_latents_std = (
+                hasattr(self.vae.config, "latents_std")
+                and self.vae.config.latents_std is not None
+            )
             if has_latents_mean and has_latents_std:
                 latents_mean = (
-                    torch.tensor(self.vae.config.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+                    torch.tensor(self.vae.config.latents_mean)
+                    .view(1, 12, 1, 1, 1)
+                    .to(latents.device, latents.dtype)
                 )
                 latents_std = (
-                    torch.tensor(self.vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+                    torch.tensor(self.vae.config.latents_std)
+                    .view(1, 12, 1, 1, 1)
+                    .to(latents.device, latents.dtype)
+                )
+                latents = (
+                    latents * latents_std / self.vae.config.scaling_factor
+                    + latents_mean
                 )
-                latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
             else:
                 latents = latents / self.vae.config.scaling_factor
 
             video = self.vae.decode(latents, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-            
+            video = self.video_processor.postprocess_video(
+                video, output_type=output_type
+            )
+
         # Offload all models
         self.maybe_free_model_hooks()
         if return_all_states:
@@ -761,7 +842,7 @@ class MochiPipeline(DiffusionPipeline, Mochi1LoraLoaderMixin):
             # prompt_embeds with shape torch.Size([2, 256]), where prompt_embeds[1] is the prompt_embeds for the actual prompt
             # prompt_embeds[0] is for negative prompt
             return original_noise, video, latents, prompt_embeds, prompt_attention_mask
-        
+
         if not return_dict:
             return (video,)
 
diff --git a/fastvideo/sample/generate_synthetic.py b/fastvideo/sample/generate_synthetic.py
index 1ad4e3f50..f5186a808 100644
--- a/fastvideo/sample/generate_synthetic.py
+++ b/fastvideo/sample/generate_synthetic.py
@@ -7,7 +7,10 @@ import os
 from diffusers.utils import export_to_video
 import argparse
 
-def generate_video_and_latent(pipe, prompt, height, width, num_frames, num_inference_steps, guidance_scale):
+
+def generate_video_and_latent(
+    pipe, prompt, height, width, num_frames, num_inference_steps, guidance_scale
+):
     # Set the random seed for reproducibility
     generator = torch.Generator("cuda").manual_seed(12345)
     # Generate videos from the input prompt
@@ -19,17 +22,16 @@ def generate_video_and_latent(pipe, prompt, height, width, num_frames, num_infer
         generator=generator,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
-        output_type="latent_and_video"
+        output_type="latent_and_video",
     )
     # prompt_embed has negative prompt at index 0
     return noise[0], video[0], latent[0], prompt_embed[1], prompt_attention_mask[1]
-    
+
     # return dummy tensor to debug first
     # return torch.zeros(1, 3, 480, 848), torch.zeros(1, 256, 16, 16)
 
 
 if __name__ == "__main__":
-
     parser = argparse.ArgumentParser()
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument("--height", type=int, default=480)
@@ -37,47 +39,71 @@ if __name__ == "__main__":
     parser.add_argument("--num_inference_steps", type=int, default=64)
     parser.add_argument("--guidance_scale", type=float, default=4.5)
     parser.add_argument("--model_path", type=str, default="data/mochi")
-    parser.add_argument("--prompt_path", type=str, default="data/dummyVid/videos2caption.json")
+    parser.add_argument(
+        "--prompt_path", type=str, default="data/dummyVid/videos2caption.json"
+    )
     parser.add_argument("--dataset_output_dir", type=str, default="data/dummySynthetic")
     args = parser.parse_args()
-    
-    local_rank = int(os.getenv('RANK', 0))
-    world_size = int(os.getenv('WORLD_SIZE', 1))
-    print('world_size', world_size, 'local rank', local_rank)
-    torch.cuda.set_device(local_rank)
-    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
 
-    
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    print("world_size", world_size, "local rank", local_rank)
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(
+        backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
+    )
+
     if not isinstance(args.prompt_path, list):
         args.prompt_path = [args.prompt_path]
-    if len(args.prompt_path) == 1 and args.prompt_path[0].endswith('txt'):
-        text_prompt = open(args.prompt_path[0], 'r').readlines()
+    if len(args.prompt_path) == 1 and args.prompt_path[0].endswith("txt"):
+        text_prompt = open(args.prompt_path[0], "r").readlines()
         text_prompt = [i.strip() for i in text_prompt]
-        
-        
+
     pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
     pipe.enable_vae_tiling()
     pipe.enable_model_cpu_offload(gpu_id=local_rank)
-    # make dir if not exist 
-    
+    # make dir if not exist
+
     os.makedirs(args.dataset_output_dir, exist_ok=True)
     os.makedirs(os.path.join(args.dataset_output_dir, "noise"), exist_ok=True)
     os.makedirs(os.path.join(args.dataset_output_dir, "video"), exist_ok=True)
     os.makedirs(os.path.join(args.dataset_output_dir, "latent"), exist_ok=True)
     os.makedirs(os.path.join(args.dataset_output_dir, "prompt_embed"), exist_ok=True)
-    os.makedirs(os.path.join(args.dataset_output_dir, "prompt_attention_mask"), exist_ok=True)
+    os.makedirs(
+        os.path.join(args.dataset_output_dir, "prompt_attention_mask"), exist_ok=True
+    )
     data = []
     for i, prompt in enumerate(text_prompt):
         if i % world_size != local_rank:
             continue
-        noise, video, latent, prompt_embed, prompt_attention_mask = generate_video_and_latent(pipe, prompt, args.height, args.width, args.num_frames, args.num_inference_steps, args.guidance_scale)
+        (
+            noise,
+            video,
+            latent,
+            prompt_embed,
+            prompt_attention_mask,
+        ) = generate_video_and_latent(
+            pipe,
+            prompt,
+            args.height,
+            args.width,
+            args.num_frames,
+            args.num_inference_steps,
+            args.guidance_scale,
+        )
         # save latent
         video_name = str(i)
         noise_path = os.path.join(args.dataset_output_dir, "noise", video_name + ".pt")
-        latent_path = os.path.join(args.dataset_output_dir, "latent", video_name + ".pt")
-        prompt_embed_path = os.path.join(args.dataset_output_dir, "prompt_embed", video_name + ".pt")
+        latent_path = os.path.join(
+            args.dataset_output_dir, "latent", video_name + ".pt"
+        )
+        prompt_embed_path = os.path.join(
+            args.dataset_output_dir, "prompt_embed", video_name + ".pt"
+        )
         video_path = os.path.join(args.dataset_output_dir, "video", video_name + ".mp4")
-        prompt_attention_mask_path = os.path.join(args.dataset_output_dir, "prompt_attention_mask", video_name + ".pt")
+        prompt_attention_mask_path = os.path.join(
+            args.dataset_output_dir, "prompt_attention_mask", video_name + ".pt"
+        )
         # save latent
         torch.save(noise, noise_path)
         torch.save(latent, latent_path)
@@ -85,7 +111,7 @@ if __name__ == "__main__":
         torch.save(prompt_attention_mask, prompt_attention_mask_path)
         export_to_video(video, video_path, fps=30)
         item = {}
-        
+
         item["cap"] = prompt
         item["video"] = video_name + ".mp4"
         item["noise"] = video_name + ".pt"
@@ -97,11 +123,11 @@ if __name__ == "__main__":
     local_data = data
     gathered_data = [None] * world_size
     dist.all_gather_object(gathered_data, local_data)
-    
+
     # save json
     if local_rank == 0:
         all_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.dataset_output_dir, "videos2caption.json"), 'w') as f:
+        with open(
+            os.path.join(args.dataset_output_dir, "videos2caption.json"), "w"
+        ) as f:
             json.dump(all_data, f, indent=4)
-
-    
\ No newline at end of file
diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py
index 5a6582aae..ac93fd51d 100644
--- a/fastvideo/sample/sample_t2v_mochi.py
+++ b/fastvideo/sample/sample_t2v_mochi.py
@@ -3,7 +3,10 @@ from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
 import torch.distributed as dist
 
 from diffusers.utils import export_to_video
-from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
+from fastvideo.utils.parallel_states import (
+    initialize_sequence_parallel_state,
+    nccl_info,
+)
 import argparse
 import os
 from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
@@ -19,14 +22,19 @@ from typing import Dict
 from diffusers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import convert_unet_state_dict_to_peft
 from fastvideo.distill.solver import PCMFMScheduler
+
+
 def initialize_distributed():
-    local_rank = int(os.getenv('RANK', 0))
-    world_size = int(os.getenv('WORLD_SIZE', 1))
-    print('world_size', world_size)
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    print("world_size", world_size)
     torch.cuda.set_device(local_rank)
-    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+    dist.init_process_group(
+        backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
+    )
     initialize_sequence_parallel_state(world_size)
 
+
 def main(args):
     initialize_distributed()
     print(nccl_info.sp_size)
@@ -37,60 +45,81 @@ def main(args):
         scheduler = FlowMatchEulerDiscreteScheduler()
     else:
         linear_quadratic = True if "linear_quadratic" in args.scheduler_type else False
-        scheduler = PCMFMScheduler(1000, args.shift, args.num_euler_timesteps, linear_quadratic,args.linear_threshold, args.linear_range)
+        scheduler = PCMFMScheduler(
+            1000,
+            args.shift,
+            args.num_euler_timesteps,
+            linear_quadratic,
+            args.linear_threshold,
+            args.linear_range,
+        )
     if args.transformer_path is not None:
         transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
     else:
-        transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder = 'transformer/')
-    
-    pipe = MochiPipeline.from_pretrained(args.model_path, transformer = transformer,scheduler=scheduler)
-    
+        transformer = MochiTransformer3DModel.from_pretrained(
+            args.model_path, subfolder="transformer/"
+        )
+
+    pipe = MochiPipeline.from_pretrained(
+        args.model_path, transformer=transformer, scheduler=scheduler
+    )
+
     pipe.enable_vae_tiling()
-    
+
     if args.lora_checkpoint_dir is not None:
         print(f"Loading LoRA weights from {args.lora_checkpoint_dir}")
         config_path = os.path.join(args.lora_checkpoint_dir, "lora_config.json")
-        with open(config_path, 'r') as f:
+        with open(config_path, "r") as f:
             lora_config_dict = json.load(f)
-        rank=lora_config_dict['lora_params']['lora_rank']
-        lora_alpha=lora_config_dict['lora_params']['lora_alpha']
+        rank = lora_config_dict["lora_params"]["lora_rank"]
+        lora_alpha = lora_config_dict["lora_params"]["lora_alpha"]
         lora_scaling = lora_alpha / rank
         pipe.load_lora_weights(args.lora_checkpoint_dir, adapter_name="default")
         pipe.set_adapters(["default"], [lora_scaling])
         print(f"Successfully Loaded LoRA weights from {args.lora_checkpoint_dir}")
     # pipe.to(device)
-    
+
     pipe.enable_model_cpu_offload(device)
-    
+
     # Generate videos from the input prompt
 
     if args.prompt_embed_path is not None:
-        prompt_embeds = torch.load(args.prompt_embed_path, map_location="cpu", weights_only=True).to(device).unsqueeze(0)
-        encoder_attention_mask = torch.load(args.encoder_attention_mask_path, map_location="cpu", weights_only=True).to(device).unsqueeze(0)
+        prompt_embeds = (
+            torch.load(args.prompt_embed_path, map_location="cpu", weights_only=True)
+            .to(device)
+            .unsqueeze(0)
+        )
+        encoder_attention_mask = (
+            torch.load(
+                args.encoder_attention_mask_path, map_location="cpu", weights_only=True
+            )
+            .to(device)
+            .unsqueeze(0)
+        )
         prompts = None
     elif args.prompt_path is not None:
         prompts = [line.strip() for line in open(args.prompt_path, "r")]
         prompt_embeds = None
         encoder_attention_mask = None
-    else: 
+    else:
         prompts = args.prompts
         prompt_embeds = None
         encoder_attention_mask = None
-        
+
     if prompts is not None:
         videos = []
         with torch.autocast("cuda", dtype=torch.bfloat16):
-                for prompt in prompts:
-                    video = pipe(
-                        prompt=[prompt],
-                        height=args.height,
-                        width=args.width,
-                        num_frames=args.num_frames,
-                        num_inference_steps=args.num_inference_steps,
-                        guidance_scale=args.guidance_scale,
-                        generator=generator,
-                    ).frames
-                    videos.append(video[0])
+            for prompt in prompts:
+                video = pipe(
+                    prompt=[prompt],
+                    height=args.height,
+                    width=args.width,
+                    num_frames=args.num_frames,
+                    num_inference_steps=args.num_inference_steps,
+                    guidance_scale=args.guidance_scale,
+                    generator=generator,
+                ).frames
+                videos.append(video[0])
     else:
         with torch.autocast("cuda", dtype=torch.bfloat16):
             videos = pipe(
@@ -106,18 +135,21 @@ def main(args):
 
     if nccl_info.global_rank <= 0:
         if prompts is not None:
-            # mkdir 
+            # mkdir
             os.makedirs(args.output_path, exist_ok=True)
             for video, prompt in zip(videos, prompts):
                 suffix = prompt.split(".")[0]
-                export_to_video(video, os.path.join(args.output_path, f"{suffix}.mp4"), fps=30)
+                export_to_video(
+                    video, os.path.join(args.output_path, f"{suffix}.mp4"), fps=30
+                )
         else:
             export_to_video(videos[0], args.output_path + ".mp4", fps=30)
 
+
 if __name__ == "__main__":
-    # arg parse 
+    # arg parse
     parser = argparse.ArgumentParser()
-    parser.add_argument("--prompts", nargs='+', default=[])
+    parser.add_argument("--prompts", nargs="+", default=[])
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument("--height", type=int, default=480)
     parser.add_argument("--width", type=int, default=848)
@@ -131,7 +163,12 @@ if __name__ == "__main__":
     parser.add_argument("--prompt_path", type=str, default=None)
     parser.add_argument("--scheduler_type", type=str, default="euler")
     parser.add_argument("--encoder_attention_mask_path", type=str, default=None)
-    parser.add_argument('--lora_checkpoint_dir', type=str, default=None, help='Path to the directory containing LoRA checkpoints')
+    parser.add_argument(
+        "--lora_checkpoint_dir",
+        type=str,
+        default=None,
+        help="Path to the directory containing LoRA checkpoints",
+    )
     parser.add_argument("--shift", type=float, default=8.0)
     parser.add_argument("--num_euler_timesteps", type=int, default=100)
     parser.add_argument("--linear_threshold", type=float, default=0.025)
diff --git a/fastvideo/sample/sample_t2v_mochi_no_sp.py b/fastvideo/sample/sample_t2v_mochi_no_sp.py
index e4b4b0748..3fc4e5b31 100644
--- a/fastvideo/sample/sample_t2v_mochi_no_sp.py
+++ b/fastvideo/sample/sample_t2v_mochi_no_sp.py
@@ -4,6 +4,8 @@ from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel
 from diffusers.utils import export_to_video, load_image, load_video
 import argparse
 from diffusers import FlowMatchEulerDiscreteScheduler
+
+
 def main(args):
     # Set the random seed for reproducibility
     generator = torch.Generator("cuda").manual_seed(args.seed)
@@ -12,8 +14,12 @@ def main(args):
     if args.transformer_path is not None:
         transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path)
     else:
-        transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder = 'transformer/')
-    pipe = MochiPipeline.from_pretrained(args.model_path, transformer = transformer, scheduler = scheduler)
+        transformer = MochiTransformer3DModel.from_pretrained(
+            args.model_path, subfolder="transformer/"
+        )
+    pipe = MochiPipeline.from_pretrained(
+        args.model_path, transformer=transformer, scheduler=scheduler
+    )
     pipe.enable_vae_tiling()
     # pipe.to("cuda:1")
     pipe.enable_model_cpu_offload()
@@ -29,14 +35,15 @@ def main(args):
             num_inference_steps=args.num_inference_steps,
             guidance_scale=args.guidance_scale,
         ).frames
-    
-    for prompt,video in zip(args.prompts, videos):
+
+    for prompt, video in zip(args.prompts, videos):
         export_to_video(video, args.output_path + f"_{prompt}.mp4", fps=30)
 
+
 if __name__ == "__main__":
-    # arg parse 
+    # arg parse
     parser = argparse.ArgumentParser()
-    parser.add_argument("--prompts", nargs='+', default=[])
+    parser.add_argument("--prompts", nargs="+", default=[])
     parser.add_argument("--num_frames", type=int, default=163)
     parser.add_argument("--height", type=int, default=480)
     parser.add_argument("--width", type=int, default=848)
diff --git a/fastvideo/train.py b/fastvideo/train.py
index ffae9e65a..a1d48d680 100644
--- a/fastvideo/train.py
+++ b/fastvideo/train.py
@@ -5,8 +5,12 @@ import math
 import os
 import shutil
 from pathlib import Path
-from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
-    destroy_sequence_parallel_group, get_sequence_parallel_state, nccl_info
+from fastvideo.utils.parallel_states import (
+    initialize_sequence_parallel_state,
+    destroy_sequence_parallel_group,
+    get_sequence_parallel_state,
+    nccl_info,
+)
 from fastvideo.utils.communications import sp_parallel_dataloader_wrapper, broadcast
 from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_mochi_dit_input
 from fastvideo.utils.validation import log_validation
@@ -39,7 +43,11 @@ from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dic
 from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
 )
-from fastvideo.utils.checkpoint import save_checkpoint, save_lora_checkpoint, resume_lora_optimizer
+from fastvideo.utils.checkpoint import (
+    save_checkpoint,
+    save_lora_checkpoint,
+    resume_lora_optimizer,
+)
 from fastvideo.utils.logging import main_print
 from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
@@ -48,10 +56,13 @@ import time
 from collections import deque
 
 
-  
-                
 def compute_density_for_timestep_sampling(
-    weighting_scheme: str, batch_size: int, generator, logit_mean: float = None, logit_std: float = None, mode_scale: float = None
+    weighting_scheme: str,
+    batch_size: int,
+    generator,
+    logit_mean: float = None,
+    logit_std: float = None,
+    mode_scale: float = None,
 ):
     """
     Compute the density for sampling the timesteps when doing SD3 training.
@@ -62,7 +73,13 @@ def compute_density_for_timestep_sampling(
     """
     if weighting_scheme == "logit_normal":
         # See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
-        u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu", generator=generator)
+        u = torch.normal(
+            mean=logit_mean,
+            std=logit_std,
+            size=(batch_size,),
+            device="cpu",
+            generator=generator,
+        )
         u = torch.nn.functional.sigmoid(u)
     elif weighting_scheme == "mode":
         u = torch.rand(size=(batch_size,), device="cpu", generator=generator)
@@ -71,6 +88,7 @@ def compute_density_for_timestep_sampling(
         u = torch.rand(size=(batch_size,), device="cpu", generator=generator)
     return u
 
+
 def get_sigmas(noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32):
     sigmas = noise_scheduler.sigmas.to(device=device, dtype=dtype)
     schedule_timesteps = noise_scheduler.timesteps.to(device)
@@ -83,15 +101,35 @@ def get_sigmas(noise_scheduler, device, timesteps, n_dim=4, dtype=torch.float32)
     return sigma
 
 
-def train_one_step_mochi(transformer, optimizer, lr_scheduler,loader, noise_scheduler, noise_random_generator, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale, rank):
+def train_one_step_mochi(
+    transformer,
+    optimizer,
+    lr_scheduler,
+    loader,
+    noise_scheduler,
+    noise_random_generator,
+    gradient_accumulation_steps,
+    sp_size,
+    precondition_outputs,
+    max_grad_norm,
+    weighting_scheme,
+    logit_mean,
+    logit_std,
+    mode_scale,
+):
     total_loss = 0.0
     optimizer.zero_grad()
     for _ in range(gradient_accumulation_steps):
-        latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
+        (
+            latents,
+            encoder_hidden_states,
+            latents_attention_mask,
+            encoder_attention_mask,
+        ) = next(loader)
         latents = normalize_mochi_dit_input(latents)
         batch_size = latents.shape[0]
         noise = torch.randn_like(latents)
-        u =   compute_density_for_timestep_sampling(
+        u = compute_density_for_timestep_sampling(
             weighting_scheme=weighting_scheme,
             batch_size=batch_size,
             generator=noise_random_generator,
@@ -104,7 +142,13 @@ def train_one_step_mochi(transformer, optimizer, lr_scheduler,loader, noise_sche
         if sp_size > 1:
             # Make sure that the timesteps are the same across all sp processes.
             broadcast(timesteps)
-        sigmas = get_sigmas(noise_scheduler, latents.device, timesteps, n_dim=latents.ndim, dtype=latents.dtype)
+        sigmas = get_sigmas(
+            noise_scheduler,
+            latents.device,
+            timesteps,
+            n_dim=latents.ndim,
+            dtype=latents.dtype,
+        )
         noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
         # if rank<=0:
         #     print("2222222222222222222222222222222222222222222222")
@@ -115,38 +159,41 @@ def train_one_step_mochi(transformer, optimizer, lr_scheduler,loader, noise_sche
                 noisy_model_input,
                 encoder_hidden_states,
                 timesteps,
-                encoder_attention_mask, # B, L
-                return_dict= False
+                encoder_attention_mask,  # B, L
+                return_dict=False,
             )[0]
         # if rank<=0:
         #     print("333333333333333333333333333333333333333333333333")
         if precondition_outputs:
-            model_pred = noisy_model_input -  model_pred * sigmas 
+            model_pred = noisy_model_input - model_pred * sigmas
         if precondition_outputs:
             target = latents
         else:
-            target =  noise - latents
-        
-        loss = torch.mean((model_pred.float() - target.float()) ** 2) / gradient_accumulation_steps
- 
+            target = noise - latents
+
+        loss = (
+            torch.mean((model_pred.float() - target.float()) ** 2)
+            / gradient_accumulation_steps
+        )
+
         loss.backward()
-        
+
         avg_loss = loss.detach().clone()
         dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
-        total_loss += avg_loss.item() 
-        
+        total_loss += avg_loss.item()
 
     grad_norm = transformer.clip_grad_norm_(max_grad_norm)
     optimizer.step()
     lr_scheduler.step()
     return total_loss, grad_norm.item()
 
+
 def main(args):
     torch.backends.cuda.matmul.allow_tf32 = True
-    
-    local_rank = int(os.environ['LOCAL_RANK'])
-    rank = int(os.environ['RANK'])
-    world_size = int(os.environ['WORLD_SIZE'])
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
     dist.init_process_group("nccl")
     torch.cuda.set_device(local_rank)
     device = torch.cuda.current_device()
@@ -160,22 +207,24 @@ def main(args):
     noise_random_generator = None
 
     # Handle the repository creation
-    if rank <=0 and args.output_dir is not None:
+    if rank <= 0 and 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.
 
     # Create model:
-    
+
     main_print(f"--> loading model from {args.pretrained_model_name_or_path}")
     # keep the master weight to float32
     transformer = MochiTransformer3DModel.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="transformer",
-        torch_dtype = torch.float32 if args.master_weight_type == 'fp32' else torch.bfloat16
+        torch_dtype=torch.float32
+        if args.master_weight_type == "fp32"
+        else torch.bfloat16,
     )
-    
+
     if args.use_lora:
         transformer.requires_grad_(False)
         transformer_lora_config = LoraConfig(
@@ -187,12 +236,18 @@ def main(args):
         transformer.add_adapter(transformer_lora_config)
 
     if args.resume_from_lora_checkpoint:
-        lora_state_dict = MochiPipeline.lora_state_dict(args.resume_from_lora_checkpoint)
+        lora_state_dict = MochiPipeline.lora_state_dict(
+            args.resume_from_lora_checkpoint
+        )
         transformer_state_dict = {
-            f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.")
+            f'{k.replace("transformer.", "")}': v
+            for k, v in lora_state_dict.items()
+            if k.startswith("transformer.")
         }
         transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)
-        incompatible_keys = set_peft_model_state_dict(transformer, transformer_state_dict, adapter_name="default")
+        incompatible_keys = set_peft_model_state_dict(
+            transformer, transformer_state_dict, adapter_name="default"
+        )
         if incompatible_keys is not None:
             # check only for unexpected keys
             unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
@@ -202,18 +257,26 @@ def main(args):
                     f" {unexpected_keys}. "
                 )
 
-    main_print(f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M")
-    main_print(f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}")
-    fsdp_kwargs = get_dit_fsdp_kwargs(args.fsdp_sharding_startegy, args.use_lora, args.use_cpu_offload, args.master_weight_type)
-    
+    main_print(
+        f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6} M"
+    )
+    main_print(
+        f"--> Initializing FSDP with sharding strategy: {args.fsdp_sharding_startegy}"
+    )
+    fsdp_kwargs = get_dit_fsdp_kwargs(
+        args.fsdp_sharding_startegy,
+        args.use_lora,
+        args.use_cpu_offload,
+        args.master_weight_type,
+    )
+
     if args.use_lora:
         transformer.config.lora_rank = args.lora_rank
         transformer.config.lora_alpha = args.lora_alpha
         transformer.config.lora_target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
         transformer._no_split_modules = ["MochiTransformerBlock"]
-        fsdp_kwargs['auto_wrap_policy'] = fsdp_kwargs['auto_wrap_policy'](transformer)
-    
-    
+        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)
+
     transformer = FSDP(
         transformer,
         **fsdp_kwargs,
@@ -234,7 +297,7 @@ def main(args):
     optimizer = torch.optim.AdamW(
         params_to_optimize,
         lr=args.learning_rate,
-        betas=(0.9,0.999),
+        betas=(0.9, 0.999),
         weight_decay=args.weight_decay,
         eps=1e-8,
     )
@@ -243,30 +306,35 @@ def main(args):
     if args.resume_from_lora_checkpoint:
         transformer, optimizer, init_steps = resume_lora_optimizer(
             transformer, args.resume_from_lora_checkpoint, optimizer
-        )   
+        )
     main_print(f"optimizer: {optimizer}")
 
-
     lr_scheduler = get_scheduler(
-                args.lr_scheduler,
-                optimizer=optimizer,
-                num_warmup_steps=args.lr_warmup_steps,
-                num_training_steps=args.max_train_steps,
-                num_cycles=args.lr_num_cycles,
-                power=args.lr_power,
-                last_epoch=init_steps - 1,
-            )
-    
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps,
+        num_training_steps=args.max_train_steps,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+        last_epoch=init_steps - 1,
+    )
+
     train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
-    sampler = LengthGroupedSampler(
-                args.train_batch_size,
-                rank=rank,
-                world_size=world_size,
-                lengths=train_dataset.lengths, 
-                group_frame=args.group_frame, 
-                group_resolution=args.group_resolution, 
-    ) if (args.group_frame or args.group_resolution) else DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=False)
-    
+    sampler = (
+        LengthGroupedSampler(
+            args.train_batch_size,
+            rank=rank,
+            world_size=world_size,
+            lengths=train_dataset.lengths,
+            group_frame=args.group_frame,
+            group_resolution=args.group_resolution,
+        )
+        if (args.group_frame or args.group_resolution)
+        else DistributedSampler(
+            train_dataset, rank=rank, num_replicas=world_size, shuffle=False
+        )
+    )
+
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
@@ -274,93 +342,136 @@ def main(args):
         pin_memory=True,
         batch_size=args.train_batch_size,
         num_workers=args.dataloader_num_workers,
-        drop_last=True, 
+        drop_last=True,
     )
 
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader)
+        / args.gradient_accumulation_steps
+        * args.sp_size
+        / args.train_sp_batch_size
+    )
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
-
     if rank <= 0:
         project = args.tracker_project_name or "fastvideo"
         wandb.init(project=project, config=args)
 
     # Train!
-    total_batch_size = args.train_batch_size * world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size
+    total_batch_size = (
+        args.train_batch_size
+        * world_size
+        * args.gradient_accumulation_steps
+        / args.sp_size
+        * args.train_sp_batch_size
+    )
     main_print("***** Running training *****")
     main_print(f"  Num examples = {len(train_dataset)}")
     main_print(f"  Dataloader size = {len(train_dataloader)}")
     main_print(f"  Num Epochs = {args.num_train_epochs}")
     main_print(f"  Resume training from step {init_steps}")
     main_print(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    main_print(f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}")
+    main_print(
+        f"  Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
+    )
     main_print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
     main_print(f"  Total optimization steps = {args.max_train_steps}")
-    main_print(f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
+    main_print(
+        f"  Total training parameters per FSDP shard = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B"
+    )
     # print dtype
     main_print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")
 
-
     # Potentially load in the weights and states from a previous save
     if args.resume_from_checkpoint:
         assert NotImplementedError("resume_from_checkpoint is not supported now.")
-        # TODO 
-
+        # TODO
 
     progress_bar = tqdm(
         range(0, args.max_train_steps),
         initial=init_steps,
         desc="Steps",
         # Only show the progress bar once on each machine.
-        disable= local_rank > 0,
+        disable=local_rank > 0,
+    )
+
+    loader = sp_parallel_dataloader_wrapper(
+        train_dataloader,
+        device,
+        args.train_batch_size,
+        args.sp_size,
+        args.train_sp_batch_size,
     )
 
-    loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size)
-    
     step_times = deque(maxlen=100)
 
-    #todo future 
+    # todo future
     for i in range(init_steps):
         next(loader)
-    for step in range(init_steps + 1, args.max_train_steps+1):
+    for step in range(init_steps + 1, args.max_train_steps + 1):
         start_time = time.time()
-        loss, grad_norm= train_one_step_mochi(transformer, optimizer, lr_scheduler, loader, noise_scheduler, noise_random_generator, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale, rank)
+        loss, grad_norm = train_one_step_mochi(
+            transformer,
+            optimizer,
+            lr_scheduler,
+            loader,
+            noise_scheduler,
+            noise_random_generator,
+            args.gradient_accumulation_steps,
+            args.sp_size,
+            args.precondition_outputs,
+            args.max_grad_norm,
+            args.weighting_scheme,
+            args.logit_mean,
+            args.logit_std,
+            args.mode_scale,
+        )
 
         step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 
-        progress_bar.set_postfix({
-        "loss": f"{loss:.4f}", 
-        "step_time": f"{step_time:.2f}s",
-        "grad_norm": grad_norm
-    })
+        progress_bar.set_postfix(
+            {
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            }
+        )
         progress_bar.update(1)
         if rank <= 0:
-            wandb.log({
-            "train_loss": loss,
-            "learning_rate": lr_scheduler.get_last_lr()[0],
-            "step_time": step_time,
-            "avg_step_time": avg_step_time,
-            "grad_norm": grad_norm 
-        }, step=step)
-        if step  % args.checkpointing_steps == 0:
+            wandb.log(
+                {
+                    "train_loss": loss,
+                    "learning_rate": lr_scheduler.get_last_lr()[0],
+                    "step_time": step_time,
+                    "avg_step_time": avg_step_time,
+                    "grad_norm": grad_norm,
+                },
+                step=step,
+            )
+        if step % args.checkpointing_steps == 0:
             if args.use_lora:
                 # Save LoRA weights
-                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
+                save_lora_checkpoint(
+                    transformer, optimizer, rank, args.output_dir, step
+                )
             else:
                 # Your existing checkpoint saving code
                 save_checkpoint(transformer, optimizer, rank, args.output_dir, step)
             dist.barrier()
-        if args.log_validation and step  % args.validation_steps == 0:
-            log_validation(args, transformer, device,
-                            torch.bfloat16, step)
+        if args.log_validation and step % args.validation_steps == 0:
+            log_validation(args, transformer, device, torch.bfloat16, step)
 
     if args.use_lora:
-        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
+        save_lora_checkpoint(
+            transformer, optimizer, rank, args.output_dir, args.max_train_steps
+        )
     else:
-        save_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
-        
+        save_checkpoint(
+            transformer, optimizer, rank, args.output_dir, args.max_train_steps
+        )
+
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
 
@@ -371,90 +482,195 @@ if __name__ == "__main__":
     # dataset & dataloader
     parser.add_argument("--data_json_path", type=str, required=True)
     parser.add_argument("--num_frames", type=int, default=163)
-    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("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
-    parser.add_argument("--group_frame", action="store_true") # TODO
-    parser.add_argument("--group_resolution", action="store_true") # TODO
+    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(
+        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
+    )
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
 
     # text encoder & vae & diffusion model
     parser.add_argument("--pretrained_model_name_or_path", type=str)
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
 
     # diffusion setting
     parser.add_argument("--ema_decay", type=float, default=0.999)
     parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument('--cfg', type=float, default=0.1)
-    parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.")
-    
+    parser.add_argument("--cfg", type=float, default=0.1)
+    parser.add_argument(
+        "--precondition_outputs",
+        action="store_true",
+        help="Whether to precondition the outputs of the model.",
+    )
+
     # validation & logs
     parser.add_argument("--validation_prompt_dir", type=str)
     parser.add_argument("--uncond_prompt_dir", type=str)
-    parser.add_argument("--validation_sampling_steps", type=str, default="64", help="use ',' to split multi sampling steps")
-    parser.add_argument('--validation_guidance_scale', type=str, default="4.5", help="use ',' to split multi scale")
-    parser.add_argument('--validation_steps', type=int, default=50)
+    parser.add_argument(
+        "--validation_sampling_steps",
+        type=str,
+        default="64",
+        help="use ',' to split multi sampling steps",
+    )
+    parser.add_argument(
+        "--validation_guidance_scale",
+        type=str,
+        default="4.5",
+        help="use ',' to split multi scale",
+    )
+    parser.add_argument("--validation_steps", type=int, default=50)
     parser.add_argument("--log_validation", action="store_true")
     parser.add_argument("--tracker_project_name", type=str, default=None)
-    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("--resume_from_lora_checkpoint", type=str, default=None,
-                        help=(
-                            "Whether training should be resumed from a previous lora 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(
+        "--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(
+        "--resume_from_lora_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous lora 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***."
+        ),
+    )
 
     # optimizer & scheduler & Training
     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("--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=10, help="Number of steps for the warmup in the lr scheduler.")
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
+    parser.add_argument(
+        "--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(
+        "--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=10,
+        help="Number of steps for the warmup in the lr scheduler.",
+    )
+    parser.add_argument(
+        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
     parser.add_argument("--selective_checkpointing", type=float, default=1.0)
-    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("--use_cpu_offload", action="store_true", help="Whether to use CPU offload for param & gradient & optimizer states.")
+    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(
+        "--use_cpu_offload",
+        action="store_true",
+        help="Whether to use CPU offload for param & gradient & optimizer states.",
+    )
 
     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")
+    parser.add_argument(
+        "--train_sp_batch_size",
+        type=int,
+        default=1,
+        help="Batch size for sequence parallel training",
+    )
 
-    parser.add_argument("--use_lora", action="store_true", default=False, help="Whether to use LoRA for finetuning.") 
-    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
-    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        default=False,
+        help="Whether to use LoRA for finetuning.",
+    )
+    parser.add_argument(
+        "--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA."
+    )
+    parser.add_argument(
+        "--lora_rank", type=int, default=128, help="LoRA rank parameter. "
+    )
     parser.add_argument("--fsdp_sharding_startegy", default="full")
 
     parser.add_argument(
@@ -464,10 +680,16 @@ if __name__ == "__main__":
         choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
     )
     parser.add_argument(
-        "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
+        "--logit_mean",
+        type=float,
+        default=0.0,
+        help="mean to use when using the `'logit_normal'` weighting scheme.",
     )
     parser.add_argument(
-        "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+        "--logit_std",
+        type=float,
+        default=1.0,
+        help="std to use when using the `'logit_normal'` weighting scheme.",
     )
     parser.add_argument(
         "--mode_scale",
@@ -476,16 +698,36 @@ if __name__ == "__main__":
         help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
     )
     # lr_scheduler
-    parser.add_argument("--lr_scheduler", type=str, default="constant",
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
         help=(
             'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
             ' "constant", "constant_with_warmup"]'
         ),
     )
-    parser.add_argument("--lr_num_cycles", type=int, default=1, help="Number of cycles in the learning rate scheduler.")
-    parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.",)
-    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay to apply.")
-    parser.add_argument("--master_weight_type", type=str, default="fp32", help="Weight type to use - fp32 or bf16.")
-    
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of cycles in the learning rate scheduler.",
+    )
+    parser.add_argument(
+        "--lr_power",
+        type=float,
+        default=1.0,
+        help="Power factor of the polynomial scheduler.",
+    )
+    parser.add_argument(
+        "--weight_decay", type=float, default=0.01, help="Weight decay to apply."
+    )
+    parser.add_argument(
+        "--master_weight_type",
+        type=str,
+        default="fp32",
+        help="Weight type to use - fp32 or bf16.",
+    )
+
     args = parser.parse_args()
-    main(args)
\ No newline at end of file
+    main(args)
diff --git a/fastvideo/utils/checkpoint.py b/fastvideo/utils/checkpoint.py
index 900b20782..d9bc3089e 100644
--- a/fastvideo/utils/checkpoint.py
+++ b/fastvideo/utils/checkpoint.py
@@ -1,31 +1,42 @@
-# import 
+# import
 import os
 import json
 import torch
 from fastvideo.utils.logging import main_print
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, StateDictType, FullStateDictConfig
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    StateDictType,
+    FullStateDictConfig,
+)
 from safetensors.torch import save_file, load_file
 import torch.distributed.checkpoint as dist_cp
-from torch.distributed.checkpoint.default_planner import DefaultSavePlanner, DefaultLoadPlanner
+from torch.distributed.checkpoint.default_planner import (
+    DefaultSavePlanner,
+    DefaultLoadPlanner,
+)
 from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
 from torch.distributed.fsdp import FullOptimStateDictConfig
 from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict
 from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline
-def save_checkpoint(model, optimizer, rank, output_dir, step, discriminator=False):
 
+
+def save_checkpoint(model, optimizer, rank, output_dir, step, discriminator=False):
     with FSDP.state_dict_type(
-        model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True), FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        model,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+        FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
     ):
         cpu_state = model.state_dict()
         optim_state = FSDP.optim_state_dict(
-            model, 
+            model,
             optimizer,
         )
-    
-    #todo move to get_state_dict
+
+    # todo move to get_state_dict
     save_dir = os.path.join(output_dir, f"checkpoint-{step}")
     os.makedirs(save_dir, exist_ok=True)
-    # save using safetensors 
+    # save using safetensors
     if rank <= 0 and not discriminator:
         weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
@@ -41,21 +52,30 @@ def save_checkpoint(model, optimizer, rank, output_dir, step, discriminator=Fals
         save_file(cpu_state, weight_path)
         optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt")
         torch.save(optim_state, optimizer_path)
-        
-        
-        
-def save_checkpoint_generator_discriminator(model, optimizer, discriminator, discriminator_optimizer, rank, output_dir, step,):
+
+
+def save_checkpoint_generator_discriminator(
+    model,
+    optimizer,
+    discriminator,
+    discriminator_optimizer,
+    rank,
+    output_dir,
+    step,
+):
     with FSDP.state_dict_type(
-        model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        model,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
     ):
         cpu_state = model.state_dict()
-        
-    #todo move to get_state_dict
+
+    # todo move to get_state_dict
     save_dir = os.path.join(output_dir, f"checkpoint-{step}")
     os.makedirs(save_dir, exist_ok=True)
     hf_weight_dir = os.path.join(save_dir, "hf_weights")
     os.makedirs(hf_weight_dir, exist_ok=True)
-    # save using safetensors 
+    # save using safetensors
     if rank <= 0:
         config_dict = dict(model.config)
         config_path = os.path.join(hf_weight_dir, "config.json")
@@ -64,8 +84,7 @@ def save_checkpoint_generator_discriminator(model, optimizer, discriminator, dis
             json.dump(config_dict, f, indent=4)
         weight_path = os.path.join(hf_weight_dir, "diffusion_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
-    
-    
+
     main_print(f"--> saved HF weight checkpoint at path {hf_weight_dir}")
     model_weight_dir = os.path.join(save_dir, "model_weights_state")
     os.makedirs(model_weight_dir, exist_ok=True)
@@ -76,44 +95,53 @@ def save_checkpoint_generator_discriminator(model, optimizer, discriminator, dis
         model_state = model.state_dict()
         weight_state_dict = {"model": model_state}
         dist_cp.save_state_dict(
-                state_dict=weight_state_dict,
-                storage_writer=dist_cp.FileSystemWriter(model_weight_dir),
-                planner=DefaultSavePlanner(),
+            state_dict=weight_state_dict,
+            storage_writer=dist_cp.FileSystemWriter(model_weight_dir),
+            planner=DefaultSavePlanner(),
         )
         optimizer_state_dict = {"optimizer": optim_state}
         dist_cp.save_state_dict(
-                state_dict=optimizer_state_dict,
-                storage_writer=dist_cp.FileSystemWriter(model_optimizer_dir),
-                planner=DefaultSavePlanner(),
+            state_dict=optimizer_state_dict,
+            storage_writer=dist_cp.FileSystemWriter(model_optimizer_dir),
+            planner=DefaultSavePlanner(),
         )
-        
-        
+
     discriminator_fsdp_state_dir = os.path.join(save_dir, "discriminator_fsdp_state")
     os.makedirs(discriminator_fsdp_state_dir, exist_ok=True)
-    with FSDP.state_dict_type(discriminator, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True), FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)):
+    with FSDP.state_dict_type(
+        discriminator,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+        FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
+    ):
         optim_state = FSDP.optim_state_dict(discriminator, discriminator_optimizer)
         model_state = discriminator.state_dict()
         state_dict = {"optimizer": optim_state, "model": model_state}
-        if rank <=0:
-            discriminator_fsdp_state_fil = os.path.join(discriminator_fsdp_state_dir, "discriminator_state.pt")
+        if rank <= 0:
+            discriminator_fsdp_state_fil = os.path.join(
+                discriminator_fsdp_state_dir, "discriminator_state.pt"
+            )
             torch.save(state_dict, discriminator_fsdp_state_fil)
-    
+
     main_print("--> saved FSDP state checkpoint")
-    
+
+
 def load_sharded_model(model, optimizer, model_dir, optimizer_dir):
     with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
         weight_state_dict = {"model": model.state_dict()}
-        
+
         optim_state = load_sharded_optimizer_state_dict(
             model_state_dict=weight_state_dict["model"],
             optimizer_key="optimizer",
             storage_reader=dist_cp.FileSystemReader(optimizer_dir),
         )
         optim_state = optim_state["optimizer"]
-        flattened_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optim_state)
+        flattened_osd = FSDP.optim_state_dict_to_load(
+            model=model, optim=optimizer, optim_state_dict=optim_state
+        )
         optimizer.load_state_dict(flattened_osd)
         dist_cp.load_state_dict(
-            state_dict = weight_state_dict,
+            state_dict=weight_state_dict,
             storage_reader=dist_cp.FileSystemReader(model_dir),
             planner=DefaultLoadPlanner(),
         )
@@ -122,38 +150,62 @@ def load_sharded_model(model, optimizer, model_dir, optimizer_dir):
     main_print(f"--> loaded model and optimizer from path {model_dir}")
     return model, optimizer
 
+
 def load_full_state_model(model, optimizer, checkpoint_file, rank):
-    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True), FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)):
+    with FSDP.state_dict_type(
+        model,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+        FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
+    ):
         discriminator_state = torch.load(checkpoint_file)
         model_state = discriminator_state["model"]
-        if  rank <= 0:
+        if rank <= 0:
             optim_state = discriminator_state["optimizer"]
         else:
             optim_state = None
         model.load_state_dict(model_state)
-        discriminator_optim_state = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=optim_state)
+        discriminator_optim_state = FSDP.optim_state_dict_to_load(
+            model=model, optim=optimizer, optim_state_dict=optim_state
+        )
         optimizer.load_state_dict(discriminator_optim_state)
-    main_print(f"--> loaded discriminator and discriminator optimizer from path {checkpoint_file}")
+    main_print(
+        f"--> loaded discriminator and discriminator optimizer from path {checkpoint_file}"
+    )
     return model, optimizer
-    
-def resume_training_generator_discriminator(model, optimizer, discriminator, discriminator_optimizer, checkpoint_dir, rank):
+
+
+def resume_training_generator_discriminator(
+    model, optimizer, discriminator, discriminator_optimizer, checkpoint_dir, rank
+):
     step = int(checkpoint_dir.split("-")[-1])
     model_weight_dir = os.path.join(checkpoint_dir, "model_weights_state")
     model_optimizer_dir = os.path.join(checkpoint_dir, "model_optimizer_state")
-    model, optimizer = load_sharded_model(model, optimizer, model_weight_dir, model_optimizer_dir)
-    discriminator_ckpt_file = os.path.join(checkpoint_dir, "discriminator_fsdp_state", "discriminator_state.pt")
-    discriminator, discriminator_optimizer = load_full_state_model(discriminator, discriminator_optimizer, discriminator_ckpt_file, rank)
+    model, optimizer = load_sharded_model(
+        model, optimizer, model_weight_dir, model_optimizer_dir
+    )
+    discriminator_ckpt_file = os.path.join(
+        checkpoint_dir, "discriminator_fsdp_state", "discriminator_state.pt"
+    )
+    discriminator, discriminator_optimizer = load_full_state_model(
+        discriminator, discriminator_optimizer, discriminator_ckpt_file, rank
+    )
     return model, optimizer, discriminator, discriminator_optimizer, step
-    
-    
+
+
 def resume_training(model, optimizer, checkpoint_dir, discriminator=False):
     weight_path = os.path.join(checkpoint_dir, "diffusion_pytorch_model.safetensors")
     if discriminator:
-        weight_path = os.path.join(checkpoint_dir, "discriminator_pytorch_model.safetensors")
+        weight_path = os.path.join(
+            checkpoint_dir, "discriminator_pytorch_model.safetensors"
+        )
     model_weights = load_file(weight_path)
 
     with FSDP.state_dict_type(
-        model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True), FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        model,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+        FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True),
     ):
         current_state = model.state_dict()
         current_state.update(model_weights)
@@ -164,78 +216,67 @@ def resume_training(model, optimizer, checkpoint_dir, discriminator=False):
         optim_path = os.path.join(checkpoint_dir, "optimizer.pt")
     optimizer_state_dict = torch.load(optim_path, weights_only=False)
     optim_state = FSDP.optim_state_dict_to_load(
-        model=model,
-        optim=optimizer,
-        optim_state_dict=optimizer_state_dict
+        model=model, optim=optimizer, optim_state_dict=optimizer_state_dict
     )
     optimizer.load_state_dict(optim_state)
     step = int(checkpoint_dir.split("-")[-1])
     return model, optimizer, step
 
 
-def save_lora_checkpoint(
-    transformer, 
-    optimizer,
-    rank, 
-    output_dir, 
-    step
-):
+def save_lora_checkpoint(transformer, optimizer, rank, output_dir, step):
     with FSDP.state_dict_type(
-        transformer, 
-        StateDictType.FULL_STATE_DICT, 
-        FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        transformer,
+        StateDictType.FULL_STATE_DICT,
+        FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
     ):
         full_state_dict = transformer.state_dict()
         lora_optim_state = FSDP.optim_state_dict(
-            transformer, 
+            transformer,
             optimizer,
         )
 
     if rank <= 0:
         save_dir = os.path.join(output_dir, f"lora-checkpoint-{step}")
         os.makedirs(save_dir, exist_ok=True)
-        
+
         # save optimizer
         optim_path = os.path.join(save_dir, "lora_optimizer.pt")
         torch.save(lora_optim_state, optim_path)
         # save lora weight
         main_print(f"--> saving LoRA checkpoint at step {step}")
-        transformer_lora_layers = get_peft_model_state_dict(model=transformer, state_dict=full_state_dict)
+        transformer_lora_layers = get_peft_model_state_dict(
+            model=transformer, state_dict=full_state_dict
+        )
         MochiPipeline.save_lora_weights(
             save_directory=save_dir,
             transformer_lora_layers=transformer_lora_layers,
-            is_main_process=True
+            is_main_process=True,
         )
         # save config
         lora_config = {
-            'step': step,
-            'lora_params': {
-                'lora_rank': transformer.config.lora_rank, 
-                'lora_alpha': transformer.config.lora_alpha,
-                'target_modules': transformer.config.lora_target_modules
-            }
+            "step": step,
+            "lora_params": {
+                "lora_rank": transformer.config.lora_rank,
+                "lora_alpha": transformer.config.lora_alpha,
+                "target_modules": transformer.config.lora_target_modules,
+            },
         }
         config_path = os.path.join(save_dir, "lora_config.json")
         with open(config_path, "w") as f:
             json.dump(lora_config, f, indent=4)
     main_print(f"--> LoRA checkpoint saved at step {step}")
 
-def resume_lora_optimizer(
-    transformer,
-    checkpoint_dir,
-    optimizer
-):
+
+def resume_lora_optimizer(transformer, checkpoint_dir, optimizer):
     config_path = os.path.join(checkpoint_dir, "lora_config.json")
     with open(config_path, "r") as f:
         config_dict = json.load(f)
     optim_path = os.path.join(checkpoint_dir, "lora_optimizer.pt")
     optimizer_state_dict = torch.load(optim_path, weights_only=False)
     optim_state = FSDP.optim_state_dict_to_load(
-            model=transformer,
-            optim=optimizer,
-            optim_state_dict=optimizer_state_dict
-        )
+        model=transformer, optim=optimizer, optim_state_dict=optimizer_state_dict
+    )
     optimizer.load_state_dict(optim_state)
-    step = config_dict['step']
+    step = config_dict["step"]
     main_print(f"-->  Successfully resuming LoRA optimizer from step {step}")
-    return transformer, optimizer, step
\ No newline at end of file
+    return transformer, optimizer, step
diff --git a/fastvideo/utils/communications.py b/fastvideo/utils/communications.py
index c09d30cca..b61bc174e 100644
--- a/fastvideo/utils/communications.py
+++ b/fastvideo/utils/communications.py
@@ -10,11 +10,12 @@ from typing import Any, Tuple
 from torch import Tensor
 from torch.nn import Module
 
+
 def broadcast(input_: torch.Tensor):
     src = nccl_info.group_id * nccl_info.sp_size
     dist.broadcast(input_, src=src, group=nccl_info.group)
-    
-    
+
+
 def _all_to_all_4D(
     input: torch.tensor, scatter_idx: int = 2, gather_idx: int = 1, group=None
 ) -> torch.tensor:
@@ -112,7 +113,6 @@ class SeqAllToAll4D(torch.autograd.Function):
         scatter_idx: int,
         gather_idx: int,
     ) -> Tensor:
-
         ctx.group = group
         ctx.scatter_idx = scatter_idx
         ctx.gather_idx = gather_idx
@@ -129,18 +129,16 @@ class SeqAllToAll4D(torch.autograd.Function):
             None,
             None,
         )
-        
-        
+
+
 def all_to_all_4D(
     input_: torch.Tensor,
     scatter_dim: int = 2,
     gather_dim: int = 1,
 ):
-    return SeqAllToAll4D.apply( nccl_info.group,input_, scatter_dim, gather_dim)
+    return SeqAllToAll4D.apply(nccl_info.group, input_, scatter_dim, gather_dim)
 
 
-
-    
 def _all_to_all(
     input_: torch.Tensor,
     world_size: int,
@@ -148,7 +146,9 @@ def _all_to_all(
     scatter_dim: int,
     gather_dim: int,
 ):
-    input_list = [t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)]
+    input_list = [
+        t.contiguous() for t in torch.tensor_split(input_, world_size, scatter_dim)
+    ]
     output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]
     dist.all_to_all(output_list, input_list, group=group)
     return torch.cat(output_list, dim=gather_dim).contiguous()
@@ -170,7 +170,9 @@ class _AllToAll(torch.autograd.Function):
         ctx.scatter_dim = scatter_dim
         ctx.gather_dim = gather_dim
         ctx.world_size = dist.get_world_size(process_group)
-        output = _all_to_all(input_, ctx.world_size, process_group, scatter_dim, gather_dim)
+        output = _all_to_all(
+            input_, ctx.world_size, process_group, scatter_dim, gather_dim
+        )
         return output
 
     @staticmethod
@@ -198,7 +200,6 @@ def all_to_all(
     return _AllToAll.apply(input_, nccl_info.group, scatter_dim, gather_dim)
 
 
-
 class _AllGather(torch.autograd.Function):
     """All-gather communication with autograd support.
 
@@ -237,6 +238,7 @@ class _AllGather(torch.autograd.Function):
 
         return grad_input, None
 
+
 def all_gather(input_: torch.Tensor, dim: int = 1):
     """Performs an all-gather operation on the input tensor along the specified dimension.
 
@@ -250,49 +252,83 @@ def all_gather(input_: torch.Tensor, dim: int = 1):
     return _AllGather.apply(input_, dim)
 
 
-def prepare_sequence_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
+def prepare_sequence_parallel_data(
+    hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
+):
     if nccl_info.sp_size == 1:
-        return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
-    def prepare(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
-        
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            attention_mask,
+            encoder_attention_mask,
+        )
+
+    def prepare(
+        hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
+    ):
         hidden_states = all_to_all(hidden_states, scatter_dim=2, gather_dim=0)
-        encoder_hidden_states = all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0)
+        encoder_hidden_states = all_to_all(
+            encoder_hidden_states, scatter_dim=1, gather_dim=0
+        )
         attention_mask = all_to_all(attention_mask, scatter_dim=1, gather_dim=0)
-        encoder_attention_mask = all_to_all(encoder_attention_mask, scatter_dim=1, gather_dim=0)
-        return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
+        encoder_attention_mask = all_to_all(
+            encoder_attention_mask, scatter_dim=1, gather_dim=0
+        )
+        return (
+            hidden_states,
+            encoder_hidden_states,
+            attention_mask,
+            encoder_attention_mask,
+        )
 
     sp_size = nccl_info.sp_size
     frame = hidden_states.shape[2]
     assert frame % sp_size == 0, "frame should be a multiple of sp_size"
 
-    hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask = prepare(hidden_states,
-                                                                                            encoder_hidden_states.repeat(1, sp_size,  1),
-                                                                                            attention_mask.repeat(1, sp_size, 1, 1),
-                                                                                            encoder_attention_mask.repeat(1, sp_size))
-
+    (
+        hidden_states,
+        encoder_hidden_states,
+        attention_mask,
+        encoder_attention_mask,
+    ) = prepare(
+        hidden_states,
+        encoder_hidden_states.repeat(1, sp_size, 1),
+        attention_mask.repeat(1, sp_size, 1, 1),
+        encoder_attention_mask.repeat(1, sp_size),
+    )
 
     return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask
 
 
-
-def  sp_parallel_dataloader_wrapper(dataloader, device, train_batch_size, sp_size, train_sp_batch_size):
-        while True:
-            for data_item in dataloader:
-                latents, cond,attn_mask, cond_mask = data_item    
-                latents = latents.to(device)
-                cond = cond.to(device)
-                attn_mask = attn_mask.to(device)
-                cond_mask = cond_mask.to(device)
-                frame = latents.shape[2]
-                if frame == 1:
-                    yield latents, cond, attn_mask, cond_mask
-                else:
-                    latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data(latents, cond, attn_mask, cond_mask)
-                    assert train_batch_size * sp_size >= train_sp_batch_size, "train_batch_size * sp_size should be greater than train_sp_batch_size"
-                    for iter in range(train_batch_size * sp_size // train_sp_batch_size):
-                        st_idx = iter * train_sp_batch_size
-                        ed_idx = (iter + 1) * train_sp_batch_size
-                        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]
-                        yield latents[st_idx: ed_idx], encoder_hidden_states, attention_mask, encoder_attention_mask
\ No newline at end of file
+def sp_parallel_dataloader_wrapper(
+    dataloader, device, train_batch_size, sp_size, train_sp_batch_size
+):
+    while True:
+        for data_item in dataloader:
+            latents, cond, attn_mask, cond_mask = data_item
+            latents = latents.to(device)
+            cond = cond.to(device)
+            attn_mask = attn_mask.to(device)
+            cond_mask = cond_mask.to(device)
+            frame = latents.shape[2]
+            if frame == 1:
+                yield latents, cond, attn_mask, cond_mask
+            else:
+                latents, cond, attn_mask, cond_mask = prepare_sequence_parallel_data(
+                    latents, cond, attn_mask, cond_mask
+                )
+                assert (
+                    train_batch_size * sp_size >= train_sp_batch_size
+                ), "train_batch_size * sp_size should be greater than train_sp_batch_size"
+                for iter in range(train_batch_size * sp_size // train_sp_batch_size):
+                    st_idx = iter * train_sp_batch_size
+                    ed_idx = (iter + 1) * train_sp_batch_size
+                    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]
+                    yield (
+                        latents[st_idx:ed_idx],
+                        encoder_hidden_states,
+                        attention_mask,
+                        encoder_attention_mask,
+                    )
diff --git a/fastvideo/utils/data_preprocess/finetune_data_T5.py b/fastvideo/utils/data_preprocess/finetune_data_T5.py
index fcd7e4da1..a17b36d34 100644
--- a/fastvideo/utils/data_preprocess/finetune_data_T5.py
+++ b/fastvideo/utils/data_preprocess/finetune_data_T5.py
@@ -6,11 +6,13 @@ from diffusers.utils import export_to_video
 import json
 import os
 import torch.distributed as dist
+
 logger = get_logger(__name__)
 from torch.utils.data import Dataset
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data import DataLoader
 
+
 class T5dataset(Dataset):
     def __init__(
         self,
@@ -21,31 +23,40 @@ class T5dataset(Dataset):
         self.vae_debug = vae_debug
         with open(self.json_path, "r") as f:
             train_dataset = json.load(f)
-            self.train_dataset = sorted(train_dataset, key=lambda x: x['latent_path'])
+            self.train_dataset = sorted(train_dataset, key=lambda x: x["latent_path"])
+
     def __getitem__(self, idx):
-        caption = self.train_dataset[idx]['caption']
-        filename = self.train_dataset[idx]['latent_path'].split('.')[0]
-        length = self.train_dataset[idx]['length']
+        caption = self.train_dataset[idx]["caption"]
+        filename = self.train_dataset[idx]["latent_path"].split(".")[0]
+        length = self.train_dataset[idx]["length"]
         if self.vae_debug:
-            latents = torch.load(os.path.join(args.output_dir, 'latent', self.train_dataset[idx]['latent_path']), map_location="cpu")
+            latents = torch.load(
+                os.path.join(
+                    args.output_dir, "latent", self.train_dataset[idx]["latent_path"]
+                ),
+                map_location="cpu",
+            )
         else:
             latents = []
-        
+
         return dict(caption=caption, latents=latents, filename=filename, length=length)
 
     def __len__(self):
         return len(self.train_dataset)
-            
-def main(args): 
-    local_rank = int(os.getenv('RANK', 0))
-    world_size = int(os.getenv('WORLD_SIZE', 1))
-    print('world_size', world_size, 'local rank', local_rank)
-    
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    print("world_size", world_size, "local rank", local_rank)
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     torch.cuda.set_device(local_rank)
     if not dist.is_initialized():
-        dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
-    
+        dist.init_process_group(
+            backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
+        )
+
     pipe = MochiPipeline.from_pretrained(args.model_path).to(device)
     pipe.vae.enable_tiling()
     os.makedirs(args.output_dir, exist_ok=True)
@@ -53,32 +64,40 @@ def main(args):
     os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
     os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
     os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
-    
+
     latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
     train_dataset = T5dataset(latents_json_path, args.vae_debug)
-    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+    sampler = DistributedSampler(
+        train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True
+    )
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
         batch_size=args.train_batch_size,
         num_workers=args.dataloader_num_workers,
     )
-    
+
     json_data = []
     for _, data in enumerate(train_dataloader):
         with torch.inference_mode():
             with torch.autocast("cuda", dtype=torch.bfloat16):
                 prompt_embeds, prompt_attention_mask, _, _ = pipe.encode_prompt(
-                    prompt=data['caption'],
+                    prompt=data["caption"],
                 )
                 if args.vae_debug:
-                    latents = data['latents']
+                    latents = data["latents"]
                     video = pipe.vae.decode(latents.to(device), return_dict=False)[0]
                     video = pipe.video_processor.postprocess_video(video)
-                for idx, video_name in enumerate(data['filename']):
-                    prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
-                    video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
-                    prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask", video_name + ".pt")
+                for idx, video_name in enumerate(data["filename"]):
+                    prompt_embed_path = os.path.join(
+                        args.output_dir, "prompt_embed", video_name + ".pt"
+                    )
+                    video_path = os.path.join(
+                        args.output_dir, "video", video_name + ".mp4"
+                    )
+                    prompt_attention_mask_path = os.path.join(
+                        args.output_dir, "prompt_attention_mask", video_name + ".pt"
+                    )
                     # save latent
                     torch.save(prompt_embeds[idx], prompt_embed_path)
                     torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
@@ -86,11 +105,11 @@ def main(args):
                     if args.vae_debug:
                         export_to_video(video[idx], video_path, fps=30)
                     item = {}
-                    item['length'] = int(data['length'][idx])
+                    item["length"] = int(data["length"][idx])
                     item["latent_path"] = video_name + ".pt"
                     item["prompt_embed_path"] = video_name + ".pt"
                     item["prompt_attention_mask"] = video_name + ".pt"
-                    item["caption"] = data['caption'][idx]
+                    item["caption"] = data["caption"][idx]
                     json_data.append(item)
     dist.barrier()
     local_data = json_data
@@ -99,19 +118,35 @@ def main(args):
     if local_rank == 0:
         # os.remove(latents_json_path)
         all_json_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.output_dir, "videos2caption.json"), 'w') as f:
+        with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f:
             json.dump(all_json_data, f, indent=4)
-            
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     # dataset & dataloader
     parser.add_argument("--model_path", type=str, default="data/mochi")
     # text encoder & vae & diffusion model
-    parser.add_argument("--dataloader_num_workers", type=int, default=1, 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=1, help="Batch size (per device) for the training dataloader.")
-    parser.add_argument("--text_encoder_name", type=str, default='google/t5-v1_1-xxl')
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
-    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("--vae_debug",action="store_true")
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        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=1,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    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("--vae_debug", action="store_true")
     args = parser.parse_args()
     main(args)
diff --git a/fastvideo/utils/data_preprocess/finetune_data_VAE.py b/fastvideo/utils/data_preprocess/finetune_data_VAE.py
index 2edfb831e..581fae6dd 100644
--- a/fastvideo/utils/data_preprocess/finetune_data_VAE.py
+++ b/fastvideo/utils/data_preprocess/finetune_data_VAE.py
@@ -14,21 +14,30 @@ from torch.utils.data.distributed import DistributedSampler
 
 logger = get_logger(__name__)
 
+
 def main(args):
-    local_rank = int(os.getenv('RANK', 0))
-    world_size = int(os.getenv('WORLD_SIZE', 1))
-    print('world_size', world_size, 'local rank', local_rank)
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    print("world_size", world_size, "local rank", local_rank)
     args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = 4, 8, 8
     args.ae_stride = args.ae_stride_h
     patch_size_t, patch_size_h, patch_size_w = 1, 2, 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
-    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=args.logging_dir)
+    args.patch_size_t, args.patch_size_h, args.patch_size_w = (
+        patch_size_t,
+        patch_size_h,
+        patch_size_w,
+    )
+    accelerator_project_config = ProjectConfiguration(
+        project_dir=args.output_dir, logging_dir=args.logging_dir
+    )
     accelerator = Accelerator(
         project_config=accelerator_project_config,
     )
     train_dataset = getdataset(args)
-    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+    sampler = DistributedSampler(
+        train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True
+    )
     train_dataloader = DataLoader(
         train_dataset,
         sampler=sampler,
@@ -36,29 +45,36 @@ def main(args):
         num_workers=args.dataloader_num_workers,
     )
 
-
     encoder_device = torch.device(f"cuda" if torch.cuda.is_available() else "cpu")
     torch.cuda.set_device(local_rank)
     if not dist.is_initialized():
-        dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
-    vae = AutoencoderKLMochi.from_pretrained(args.model_path, subfolder="vae").to("cuda")
+        dist.init_process_group(
+            backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
+        )
+    vae = AutoencoderKLMochi.from_pretrained(args.model_path, subfolder="vae").to(
+        "cuda"
+    )
     vae.enable_tiling()
     os.makedirs(args.output_dir, exist_ok=True)
     os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
-    
+
     json_data = []
     for _, data in enumerate(train_dataloader):
         with torch.inference_mode():
             with torch.autocast("cuda", dtype=torch.bfloat16):
-                latents = vae.encode(data['pixel_values'].to(encoder_device))['latent_dist'].sample()
-                for idx, video_path in enumerate(data['path']):
+                latents = vae.encode(data["pixel_values"].to(encoder_device))[
+                    "latent_dist"
+                ].sample()
+                for idx, video_path in enumerate(data["path"]):
                     video_name = os.path.basename(video_path).split(".")[0]
-                    latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
+                    latent_path = os.path.join(
+                        args.output_dir, "latent", video_name + ".pt"
+                    )
                     torch.save(latents[idx].to(torch.bfloat16), latent_path)
                     item = {}
                     item["length"] = latents[idx].shape[1]
                     item["latent_path"] = video_name + ".pt"
-                    item["caption"] = data['text'][idx]
+                    item["caption"] = data["text"][idx]
                     json_data.append(item)
                     print(f"{video_name} processed")
     dist.barrier()
@@ -67,40 +83,61 @@ def main(args):
     dist.all_gather_object(gathered_data, local_data)
     if local_rank == 0:
         all_json_data = [item for sublist in gathered_data for item in sublist]
-        with open(os.path.join(args.output_dir, "videos2caption_temp.json"), 'w') as f:
+        with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f:
             json.dump(all_json_data, f, indent=4)
-            
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     # dataset & dataloader
     parser.add_argument("--model_path", type=str, default="data/mochi")
     parser.add_argument("--data_merge_path", type=str, required=True)
     parser.add_argument("--num_frames", type=int, default=163)
-    parser.add_argument("--dataloader_num_workers", type=int, default=1, 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("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        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(
+        "--num_latent_t", type=int, default=28, help="Number of latent timesteps."
+    )
     parser.add_argument("--max_height", type=int, default=480)
     parser.add_argument("--max_width", type=int, default=848)
     parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
-    parser.add_argument("--group_frame", action="store_true") # TODO
-    parser.add_argument("--group_resolution", action="store_true") # TODO
-    parser.add_argument("--dataset", default='t2v')
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)
     parser.add_argument("--speed_factor", type=float, default=1.0)
     parser.add_argument("--drop_short_ratio", type=float, default=1.0)
     # text encoder & vae & diffusion model
-    parser.add_argument("--text_encoder_name", type=str, default='google/t5-v1_1-xxl')
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
-    parser.add_argument('--cfg', type=float, default=0.0)
-    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("--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("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument("--cfg", type=float, default=0.0)
+    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(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
 
     args = parser.parse_args()
     main(args)
diff --git a/fastvideo/utils/dataset_utils.py b/fastvideo/utils/dataset_utils.py
index a28207837..1f801fafa 100644
--- a/fastvideo/utils/dataset_utils.py
+++ b/fastvideo/utils/dataset_utils.py
@@ -13,11 +13,13 @@ from collections import Counter
 import random
 
 
-IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
+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."""
 
@@ -31,17 +33,20 @@ class DecordInit(object):
             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)
+        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})')
+        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:
@@ -50,6 +55,7 @@ def pad_to_multiple(number, ds_stride):
         padding = ds_stride - remainder
         return number + padding
 
+
 class Collate:
     def __init__(self, args):
         self.batch_size = args.train_batch_size
@@ -71,9 +77,9 @@ class Collate:
         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]
+        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):
@@ -81,13 +87,29 @@ class Collate:
 
         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'
+
+        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):
+    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
@@ -98,13 +120,30 @@ class Collate:
             if len(count_dict) != 1:
                 sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
                 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)
+                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)]
+                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]
 
@@ -115,50 +154,61 @@ class Collate:
             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_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_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) 
+            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[0] - 1) // ae_stride_thw[0] + 1),
             max_tube_size[1] // ae_stride_thw[1],
-            max_tube_size[2] // ae_stride_thw[2]
-            ]
+            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[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]
+                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]
+            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.
@@ -184,13 +234,16 @@ def split_to_even_chunks(indices, lengths, num_chunks, batch_size):
         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))]
+                chunk = chunk + [
+                    random.choice(chunk) for _ in range(batch_size - len(chunk))
+                ]
             else:
                 chunk = random.choice(pad_chunks)
-                print(chunks[idx], '->', chunk)
+                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)
@@ -204,48 +257,70 @@ def megabatch_frame_alignment(megabatches, lengths):
         len_each_megabatch = [lengths[i] for i in megabatch]
         idx_length_dict = dict([*zip(megabatch, len_each_megabatch)])
         count_dict = Counter(len_each_megabatch)
-        
+
         # mixed frame length, align megabatch inside
         if len(count_dict) != 1:
             sorted_by_value = sorted(count_dict.items(), key=lambda item: item[1])
             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(idx_length_dict) - len(candidate_batch))]
+            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(idx_length_dict) - len(candidate_batch))
+            ]
             aligned_magabatch = candidate_batch + random_select_batch
             aligned_magabatches.append(aligned_magabatch)
         # already aligned megabatches
         else:
             aligned_magabatches.append(megabatch)
-    
-    return aligned_magabatches
-                
 
-def get_length_grouped_indices(lengths, batch_size, world_size, generator=None, group_frame=False, group_resolution=False, seed=42):
+    return aligned_magabatches
+
+
+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
-    
+        generator = torch.Generator().manual_seed(
+            seed
+        )  # every rank will generate a fixed order but random index
+
     indices = torch.randperm(len(lengths), generator=generator).tolist()
-    
+
     # sort dataset according to frame
     indices = group_frame_fun(indices, lengths)
-    
+
     # chunk dataset to megabatches
     megabatch_size = world_size * batch_size
-    megabatches = [indices[i: i + megabatch_size] for i in range(0, len(lengths), megabatch_size)]
+    megabatches = [
+        indices[i : i + megabatch_size] for i in range(0, len(lengths), megabatch_size)
+    ]
 
     # make sure the length in each magabatch is align with each other
     megabatches = megabatch_frame_alignment(megabatches, lengths)
-    
+
     # aplit aligned megabatch into batches
-    megabatches = [split_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches]
+    megabatches = [
+        split_to_even_chunks(megabatch, lengths, world_size, batch_size)
+        for megabatch in megabatches
+    ]
 
     # random megabatches to do video-image mix training
     indices = torch.randperm(len(megabatches), generator=generator).tolist()
     shuffled_megabatches = [megabatches[i] for i in indices]
 
     # expand indices and return
-    return [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):
@@ -259,9 +334,9 @@ class LengthGroupedSampler(Sampler):
         batch_size: int,
         rank: int,
         world_size: int,
-        lengths: Optional[List[int]] = None, 
-        group_frame=False, 
-        group_resolution=False, 
+        lengths: Optional[List[int]] = None,
+        group_frame=False,
+        group_resolution=False,
         generator=None,
     ):
         if lengths is None:
@@ -279,15 +354,24 @@ class LengthGroupedSampler(Sampler):
         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)
+        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,
+        )
+
         def distributed_sampler(lst, rank, batch_size, world_size):
             result = []
             index = rank * batch_size
             while index < len(lst):
-                result.extend(lst[index:index + batch_size])
+                result.extend(lst[index : index + batch_size])
                 index += batch_size * world_size
             return result
-        
-        indices = distributed_sampler(indices, self.rank, self.batch_size, self.world_size)
+
+        indices = distributed_sampler(
+            indices, self.rank, self.batch_size, self.world_size
+        )
         return iter(indices)
diff --git a/fastvideo/utils/logging.py b/fastvideo/utils/logging.py
index faee5b8dd..ea2fe86ae 100644
--- a/fastvideo/utils/logging.py
+++ b/fastvideo/utils/logging.py
@@ -1,23 +1,24 @@
-        
 import sys
 import pdb
 import os
 
+
 def main_print(content):
-    if int(os.environ['LOCAL_RANK']) <= 0: 
+    if int(os.environ["LOCAL_RANK"]) <= 0:
         print(content)
 
-#ForkedPdb().set_trace()
+
+# ForkedPdb().set_trace()
 class ForkedPdb(pdb.Pdb):
     """A Pdb subclass that may be used
     from a forked multiprocessing child
 
     """
+
     def interaction(self, *args, **kwargs):
         _stdin = sys.stdin
         try:
-            sys.stdin = open('/dev/stdin')
+            sys.stdin = open("/dev/stdin")
             pdb.Pdb.interaction(self, *args, **kwargs)
         finally:
             sys.stdin = _stdin
-            
\ No newline at end of file
diff --git a/fastvideo/utils/optimizer.py b/fastvideo/utils/optimizer.py
index e5aad89f8..f7a70c771 100644
--- a/fastvideo/utils/optimizer.py
+++ b/fastvideo/utils/optimizer.py
@@ -1,7 +1,9 @@
 from accelerate.logging import get_logger
 import torch
+
 logger = get_logger(__name__)
 
+
 def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
     # Optimizer creation
     supported_optimizers = ["adam", "adamw", "prodigy"]
@@ -26,7 +28,9 @@ def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
             )
 
     if args.optimizer.lower() == "adamw":
-        optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW
+        optimizer_class = (
+            bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW
+        )
 
         optimizer = optimizer_class(
             params_to_optimize,
@@ -47,7 +51,9 @@ def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
         try:
             import prodigyopt
         except ImportError:
-            raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+            raise ImportError(
+                "To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`"
+            )
 
         optimizer_class = prodigyopt.Prodigy
 
@@ -68,4 +74,4 @@ def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
             safeguard_warmup=args.prodigy_safeguard_warmup,
         )
 
-    return optimizer
\ No newline at end of file
+    return optimizer
diff --git a/fastvideo/utils/parallel_states.py b/fastvideo/utils/parallel_states.py
index d32210924..1897db207 100644
--- a/fastvideo/utils/parallel_states.py
+++ b/fastvideo/utils/parallel_states.py
@@ -2,6 +2,7 @@ import torch
 import torch.distributed as dist
 import os
 
+
 class COMM_INFO:
     def __init__(self):
         self.group = None
@@ -10,8 +11,11 @@ class COMM_INFO:
         self.rank_within_group = 0
         self.group_id = 0
 
+
 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:
@@ -19,22 +23,29 @@ def initialize_sequence_parallel_state(sequence_parallel_size):
         initialize_sequence_parallel_group(sequence_parallel_size)
     else:
         nccl_info.sp_size = 1
-        nccl_info.global_rank = int(os.getenv('RANK', '0'))
+        nccl_info.global_rank = int(os.getenv("RANK", "0"))
         nccl_info.rank_within_group = 0
-        nccl_info.group_id = int(os.getenv('RANK', '0'))
+        nccl_info.group_id = int(os.getenv("RANK", "0"))
+
 
 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, but got world_size: {}, sequence_parallel_size: {}".format(world_size, sequence_parallel_size)
+    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, but got world_size: {}, sequence_parallel_size: {}".format(
+        world_size, sequence_parallel_size
+    )
     nccl_info.sp_size = sequence_parallel_size
     nccl_info.global_rank = rank
     num_sequence_parallel_groups: int = world_size // sequence_parallel_size
diff --git a/fastvideo/utils/validation.py b/fastvideo/utils/validation.py
index 1204557f1..23ea7615b 100644
--- a/fastvideo/utils/validation.py
+++ b/fastvideo/utils/validation.py
@@ -1,13 +1,14 @@
-
-
-from typing import  Optional, Union, List
+from typing import Optional, Union, List
 import numpy as np
 import torch
 from einops import rearrange
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
-from fastvideo.utils.communications import  all_gather
+from fastvideo.utils.communications import all_gather
 from diffusers.utils.torch_utils import randn_tensor
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule, retrieve_timesteps
+from fastvideo.models.mochi_hf.pipeline_mochi import (
+    linear_quadratic_schedule,
+    retrieve_timesteps,
+)
 from tqdm import tqdm
 from diffusers.video_processor import VideoProcessor
 from diffusers import (
@@ -20,6 +21,8 @@ from diffusers.utils import export_to_video
 import os
 import wandb
 import gc
+
+
 def prepare_latents(
     batch_size,
     num_channels_latents,
@@ -38,10 +41,10 @@ def prepare_latents(
 
     shape = (batch_size, num_channels_latents, num_frames, height, width)
 
-
     latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
     return latents
 
+
 def sample_validation_video(
     transformer,
     vae,
@@ -60,8 +63,8 @@ def sample_validation_video(
     negative_prompt_embeds: Optional[torch.Tensor] = None,
     negative_prompt_attention_mask: Optional[torch.Tensor] = None,
     output_type: Optional[str] = "pil",
-    vae_spatial_scale_factor = 8,
-    vae_temporal_scale_factor = 6, 
+    vae_spatial_scale_factor=8,
+    vae_temporal_scale_factor=6,
 ):
     device = vae.device
 
@@ -70,7 +73,9 @@ def sample_validation_video(
     do_classifier_free_guidance = guidance_scale > 1.0
     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)
+        prompt_attention_mask = torch.cat(
+            [negative_prompt_attention_mask, prompt_attention_mask], dim=0
+        )
 
     # 4. Prepare latent variables
     # TODO: Remove hardcore
@@ -85,13 +90,14 @@ def sample_validation_video(
         device,
         generator,
         vae_spatial_scale_factor,
-        vae_temporal_scale_factor
+        vae_temporal_scale_factor,
     )
     world_size, rank = nccl_info.sp_size, nccl_info.rank_within_group
     if get_sequence_parallel_state():
-        latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
+        latents = rearrange(
+            latents, "b t (n s) h w -> b t n s h w", n=world_size
+        ).contiguous()
         latents = latents[:, :, rank, :, :, :]
-        
 
     # 5. Prepare timestep
     # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
@@ -118,12 +124,16 @@ def sample_validation_video(
     # with self.progress_bar(total=num_inference_steps) as progress_bar:
     # write with tqdm instead
     # only enable if nccl_info.global_rank == 0
-    
-    with tqdm(total=num_inference_steps, disable= nccl_info.rank_within_group != 0, desc="Validation sampling...") as progress_bar:
+
+    with tqdm(
+        total=num_inference_steps,
+        disable=nccl_info.rank_within_group != 0,
+        desc="Validation sampling...",
+    ) 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 = (
+                torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            )
             # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
             timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
             noise_pred = transformer(
@@ -133,16 +143,20 @@ def sample_validation_video(
                 encoder_attention_mask=prompt_attention_mask,
                 return_dict=False,
             )[0]
-            
+
             # Mochi CFG + Sampling runs in FP32
             noise_pred = noise_pred.to(torch.float32)
             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)
+                noise_pred = noise_pred_uncond + guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
 
             # compute the previous noisy sample x_t -> x_t-1
             latents_dtype = latents.dtype
-            latents = scheduler.step(noise_pred, t, latents.to(torch.float32), return_dict=False)[0]
+            latents = scheduler.step(
+                noise_pred, t, latents.to(torch.float32), return_dict=False
+            )[0]
             latents = latents.to(latents_dtype)
 
             if latents.dtype != latents_dtype:
@@ -150,28 +164,35 @@ def sample_validation_video(
                     # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
                     latents = latents.to(latents_dtype)
 
-            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) %  scheduler.order == 0):
+            if i == len(timesteps) - 1 or (
+                (i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0
+            ):
                 progress_bar.update()
 
-
-
     if get_sequence_parallel_state():
         latents = all_gather(latents, dim=2)
 
-
     if output_type == "latent":
         video = latents
     else:
         # unscale/denormalize the latents
         # denormalize with the mean and std if available and not None
-        has_latents_mean = hasattr(vae.config, "latents_mean") and vae.config.latents_mean is not None
-        has_latents_std = hasattr(vae.config, "latents_std") and vae.config.latents_std is not None
+        has_latents_mean = (
+            hasattr(vae.config, "latents_mean") and vae.config.latents_mean is not None
+        )
+        has_latents_std = (
+            hasattr(vae.config, "latents_std") and vae.config.latents_std is not None
+        )
         if has_latents_mean and has_latents_std:
             latents_mean = (
-                torch.tensor(vae.config.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+                torch.tensor(vae.config.latents_mean)
+                .view(1, 12, 1, 1, 1)
+                .to(latents.device, latents.dtype)
             )
             latents_std = (
-                torch.tensor(vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+                torch.tensor(vae.config.latents_std)
+                .view(1, 12, 1, 1, 1)
+                .to(latents.device, latents.dtype)
             )
             latents = latents * latents_std / vae.config.scaling_factor + latents_mean
         else:
@@ -180,87 +201,134 @@ def sample_validation_video(
         video = vae.decode(latents, return_dict=False)[0]
         video_processor = VideoProcessor(vae_scale_factor=vae_spatial_scale_factor)
         video = video_processor.postprocess_video(video, output_type=output_type)
-        
-
 
     return (video,)
 
 
-
 @torch.no_grad()
 @torch.autocast("cuda", dtype=torch.bfloat16)
-def log_validation(args, transformer, device, weight_dtype, global_step,  scheduler_type="euler",shift=1.0, num_euler_timesteps=100,  linear_quadratic_threshold=0.025, linear_range=0.5, ema=False):
-    #TODO
+def log_validation(
+    args,
+    transformer,
+    device,
+    weight_dtype,
+    global_step,
+    scheduler_type="euler",
+    shift=1.0,
+    num_euler_timesteps=100,
+    linear_quadratic_threshold=0.025,
+    linear_range=0.5,
+    ema=False,
+):
+    # TODO
     print(f"Running validation....\n")
-    vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype).to("cuda")
+    vae = AutoencoderKLMochi.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype
+    ).to("cuda")
     vae.enable_tiling()
     if scheduler_type == "euler":
         scheduler = FlowMatchEulerDiscreteScheduler()
     else:
         linear_quadraic = True if scheduler_type == "pcm_linear_quadratic" else False
-        scheduler = PCMFMScheduler(1000, shift, num_euler_timesteps, linear_quadraic, linear_quadratic_threshold, linear_range)
+        scheduler = PCMFMScheduler(
+            1000,
+            shift,
+            num_euler_timesteps,
+            linear_quadraic,
+            linear_quadratic_threshold,
+            linear_range,
+        )
     # args.validation_prompt_dir
-    
+
     validation_guidance_scale_ls = args.validation_guidance_scale.split(",")
-    validation_guidance_scale_ls = [float(scale) for scale in validation_guidance_scale_ls]
+    validation_guidance_scale_ls = [
+        float(scale) for scale in validation_guidance_scale_ls
+    ]
     for validation_sampling_step in args.validation_sampling_steps.split(","):
         validation_sampling_step = int(validation_sampling_step)
         for validation_guidance_scale in validation_guidance_scale_ls:
-
             videos = []
             # prompt_embed are named embed0 to embedN
             # check how many embeds are there
-            num_embeds = len([f for f in os.listdir(args.validation_prompt_dir) if "embed" in f])
+            num_embeds = len(
+                [f for f in os.listdir(args.validation_prompt_dir) if "embed" in f]
+            )
             validation_prompt_ids = list(range(num_embeds))
-            num_sp_groups = int(os.getenv("WORLD_SIZE", '1')) // nccl_info.sp_size
+            num_sp_groups = int(os.getenv("WORLD_SIZE", "1")) // nccl_info.sp_size
             # pad to multiple of groups
             validation_prompt_ids += [0] * (num_sp_groups - num_embeds % num_sp_groups)
             num_embeds_per_group = len(validation_prompt_ids) // num_sp_groups
-            local_prompt_ids = validation_prompt_ids[nccl_info.group_id * num_embeds_per_group: (nccl_info.group_id + 1) * num_embeds_per_group]
-            
+            local_prompt_ids = validation_prompt_ids[
+                nccl_info.group_id * num_embeds_per_group : (nccl_info.group_id + 1)
+                * num_embeds_per_group
+            ]
+
             for i in local_prompt_ids:
-                prompt_embed_path = os.path.join(args.validation_prompt_dir, f"embed{i}.pt")
-                prompt_mask_path = os.path.join(args.validation_prompt_dir, f"mask{i}.pt")
-                prompt_embeds = torch.load(prompt_embed_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
-                prompt_attention_mask = torch.load(prompt_mask_path, map_location="cpu", weights_only=True).to(device).to(weight_dtype).unsqueeze(0)
-                negative_prompt_embeds = torch.zeros(256, 4096).to(device).to(weight_dtype).unsqueeze(0)
-                negative_prompt_attention_mask = torch.zeros(256).bool().to(device).unsqueeze(0)
+                prompt_embed_path = os.path.join(
+                    args.validation_prompt_dir, f"embed{i}.pt"
+                )
+                prompt_mask_path = os.path.join(
+                    args.validation_prompt_dir, f"mask{i}.pt"
+                )
+                prompt_embeds = (
+                    torch.load(prompt_embed_path, map_location="cpu", weights_only=True)
+                    .to(device)
+                    .to(weight_dtype)
+                    .unsqueeze(0)
+                )
+                prompt_attention_mask = (
+                    torch.load(prompt_mask_path, map_location="cpu", weights_only=True)
+                    .to(device)
+                    .to(weight_dtype)
+                    .unsqueeze(0)
+                )
+                negative_prompt_embeds = (
+                    torch.zeros(256, 4096).to(device).to(weight_dtype).unsqueeze(0)
+                )
+                negative_prompt_attention_mask = (
+                    torch.zeros(256).bool().to(device).unsqueeze(0)
+                )
                 generator = torch.Generator(device="cuda").manual_seed(12345)
                 video = sample_validation_video(
-                            transformer,
-                            vae,
-                            scheduler,
-                            scheduler_type=scheduler_type,
-                            num_frames=args.num_frames,
-                            # Peiyuan TODO: remove hardcode
-                            height=480,
-                            width=848,
-                            num_inference_steps=validation_sampling_step,
-                            guidance_scale=validation_guidance_scale,
-                            generator=generator, 
-                            prompt_embeds = prompt_embeds,
-                            prompt_attention_mask = prompt_attention_mask,
-                            negative_prompt_embeds = negative_prompt_embeds,
-                            negative_prompt_attention_mask = negative_prompt_attention_mask,
-                            )[0]
+                    transformer,
+                    vae,
+                    scheduler,
+                    scheduler_type=scheduler_type,
+                    num_frames=args.num_frames,
+                    # Peiyuan TODO: remove hardcode
+                    height=480,
+                    width=848,
+                    num_inference_steps=validation_sampling_step,
+                    guidance_scale=validation_guidance_scale,
+                    generator=generator,
+                    prompt_embeds=prompt_embeds,
+                    prompt_attention_mask=prompt_attention_mask,
+                    negative_prompt_embeds=negative_prompt_embeds,
+                    negative_prompt_attention_mask=negative_prompt_attention_mask,
+                )[0]
                 if nccl_info.rank_within_group == 0:
                     videos.append(video[0])
             # collect videos from all process to process zero
-            
+
             gc.collect()
             torch.cuda.empty_cache()
             # log if main process
             torch.distributed.barrier()
-            all_videos = [None for i in range(int(os.getenv("WORLD_SIZE", '1')))]        # remove padded videos
+            all_videos = [
+                None for i in range(int(os.getenv("WORLD_SIZE", "1")))
+            ]  # remove padded videos
             torch.distributed.all_gather_object(all_videos, videos)
             if nccl_info.global_rank == 0:
                 # remove padding
                 videos = [video for videos in all_videos for video in videos]
                 videos = videos[:num_embeds]
-                # linearize all videos 
+                # linearize all videos
                 video_filenames = []
                 for i, video in enumerate(videos):
-                    filename = os.path.join(args.output_dir, f"validation_step_{global_step}_sample_{validation_sampling_step}_guidance_{validation_guidance_scale}_video_{i}.mp4")
+                    filename = os.path.join(
+                        args.output_dir,
+                        f"validation_step_{global_step}_sample_{validation_sampling_step}_guidance_{validation_guidance_scale}_video_{i}.mp4",
+                    )
                     export_to_video(video, filename, fps=30)
                     video_filenames.append(filename)
 
@@ -271,4 +339,3 @@ def log_validation(args, transformer, device, weight_dtype, global_step,  schedu
                     ]
                 }
                 wandb.log(logs, step=global_step)
-
diff --git a/scripts/download_hf.py b/scripts/download_hf.py
index c661b54d9..cef005c71 100644
--- a/scripts/download_hf.py
+++ b/scripts/download_hf.py
@@ -1,19 +1,39 @@
 from huggingface_hub import snapshot_download, hf_hub_download
 import argparse
-# set args for repo_id, local_dir, repo_type, 
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Download a dataset or model from the Hugging Face Hub')
-    parser.add_argument('--repo_id', type=str, help='The ID of the repository to download')
-    parser.add_argument('--local_dir', type=str, help='The local directory to download the repository to')
-    parser.add_argument('--repo_type', type=str, help='The type of repository to download (dataset or model)')
-    parser.add_argument('--file_name', type=str, help='The file name to download')
+# set args for repo_id, local_dir, repo_type,
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Download a dataset or model from the Hugging Face Hub"
+    )
+    parser.add_argument(
+        "--repo_id", type=str, help="The ID of the repository to download"
+    )
+    parser.add_argument(
+        "--local_dir",
+        type=str,
+        help="The local directory to download the repository to",
+    )
+    parser.add_argument(
+        "--repo_type",
+        type=str,
+        help="The type of repository to download (dataset or model)",
+    )
+    parser.add_argument("--file_name", type=str, help="The file name to download")
     args = parser.parse_args()
     if args.file_name:
-        hf_hub_download(repo_id=args.repo_id, filename=args.file_name, repo_type=args.repo_type, local_dir=args.local_dir)
-    else: 
-        snapshot_download(repo_id=args.repo_id,
-                        local_dir=args.local_dir,
-                        repo_type=args.repo_type,
-                        local_dir_use_symlinks=False,
-                        resume_download=True)
\ No newline at end of file
+        hf_hub_download(
+            repo_id=args.repo_id,
+            filename=args.file_name,
+            repo_type=args.repo_type,
+            local_dir=args.local_dir,
+        )
+    else:
+        snapshot_download(
+            repo_id=args.repo_id,
+            local_dir=args.local_dir,
+            repo_type=args.repo_type,
+            local_dir_use_symlinks=False,
+            resume_download=True,
+        )
diff --git a/scripts/upload_hf.py b/scripts/upload_hf.py
index 653c2af32..4e3888de3 100644
--- a/scripts/upload_hf.py
+++ b/scripts/upload_hf.py
@@ -1,8 +1,9 @@
 from huggingface_hub import HfApi
+
 api = HfApi()
 
 api.upload_folder(
     folder_path="data/Mochi-Synthetic-Data",
     repo_id="Stealths-Video/Mochi-Synthetic-Data",
     repo_type="dataset",
-)
\ No newline at end of file
+)