Add LADD (#45)

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>

README.md

@@ -8,9 +8,7 @@ conda create -n fastvideo python=3.10.12
 conda activate fastvideo
 pip3 install torch==2.5.0 torchvision  --index-url https://download.pytorch.org/whl/cu121
 pip3 install -U xformers==0.0.28.post2 --index-url https://download.pytorch.org/whl/cu121
-pip3 install ray
-cd .. && git clone  https://github.com/huggingface/diffusers
-cd diffusers && pip install -e . && cd ../FastVideo-OSP
+pip install git+https://github.com/huggingface/diffusers.git@76b7d86a9a5c0c2186efa09c4a67b5f5666ac9e3
 ```
 
 ```
@@ -26,7 +24,7 @@ python scripts/download_hf.py --repo_id=Stealths-Video/dummyVid --local_dir=data
 python scripts/download_hf.py --repo_id=Stealths-Video/mochi_diffuser --local_dir=data/mochi --repo_type=model
 python scripts/download_hf.py --repo_id=Stealths-Video/Mochi-Synthetic-Data --local_dir=data/Mochi-Synthetic-Data --repo_type=dataset
 python scripts/download_hf.py --repo_id=Stealths-Video/Encoder_Overfit_Data --local_dir=data/Encoder_Overfit_Data --repo_type=dataset
-python scripts/download_hf.py --repo_id=Stealths-Video/BW-Finetune-Dataset --local_dir=data/BW-Finetune-Dataset --repo_type=dataset
+python scripts/download_hf.py --repo_id=Stealths-Video/Black-Myth-Recaption --local_dir=Stealths-Video/BLACK-MYTH-Finetune-Dataset --repo_type=dataset
 ```
 
 ## How to overfit
@@ -35,10 +33,11 @@ bash scripts/overfit.sh
 ```
 Make sure to edit data/Mochi-Synthetic-Data/videos2caption.json such that this is only one video in the dataset (you can copy multiple annotations of the same video). Also make sure to edit the prompt in scripts/overfit.shto match the prompt in the training data. I observe the overfitting  after 50 steps. 
 
-## TODO
 
-- [X] Delete all npu related stuff.
-- [ ] Remove inpaint. 
-- [ ] Create dummy debug data. 
-- [ ] Add Mochi
-- [ ] Add Mochi VAE
+## Feature for Branding:
+- [] LoRA
+- [] Low mem. Only load DiT during training.
+- [] selective gradient checkpointing 
+- [] Fused kernel
+- [] correct attention mask issue in HF's implementation
+- []

fastvideo/distill.py

@@ -0,0 +1,509 @@
+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.communications import sp_parallel_dataloader_wrapper, broadcast
+from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
+
+from fastvideo.model.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 (
+    FullyShardedDataParallel as FSDP,
+    StateDictType,
+    FullStateDictConfig,
+)
+import json
+from torch.utils.data.distributed import DistributedSampler
+import wandb
+from accelerate.utils import set_seed
+from tqdm.auto import tqdm
+from fastvideo.fsdp_util import get_dit_fsdp_kwargs, get_discriminator_fsdp_kwargs, apply_fsdp_checkpointing
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+)
+import torch.nn.functional as F
+
+import numpy as np
+from fastvideo.model.modeling_mochi import MochiTransformer3DModel
+from diffusers.utils import check_min_version
+from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function
+from fastvideo.gan.discriminator import LADDDiscriminator
+import torch.distributed as dist
+from safetensors.torch import save_file
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.31.0")
+
+
+import sys
+import pdb
+
+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')
+            pdb.Pdb.interaction(self, *args, **kwargs)
+        finally:
+            sys.stdin = _stdin
+            
+
+
+def main_print(content):
+    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)
+    ):
+        cpu_state = transformer.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 
+        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
+        save_file(cpu_state, weight_path)
+        config_dict = dict(transformer.config)
+        config_path = os.path.join(save_dir, "config.json")
+        # 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}")
+    
+                
+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)
+    timesteps = timesteps.to(device)
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+    sigma = sigmas[step_indices].flatten()
+    while len(sigma.shape) < n_dim:
+        sigma = sigma.unsqueeze(-1)
+    return sigma
+
+
+def student_forward(transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size):
+    batch_size = latents.shape[0]
+    noise = torch.randn_like(latents)
+    # uniformly sample from [1.0, 0.75, 0.5, 0.25]
+    p = [0.25, 0.25, 0.25, 0.25]
+    sigmas = np.random.choice([1.0, 0.75, 0.5, 0.25], batch_size, p=p)
+    sigmas = torch.tensor(sigmas, device=latents.device, dtype=latents.dtype)
+    student_timesteps = (sigmas * 1000).long()
+    if sp_size > 1:
+        # Make sure that the timesteps are the same across all sp processes.
+        broadcast(student_timesteps)
+
+    while len(sigmas.shape) < latents.ndim:
+        sigmas = sigmas.unsqueeze(-1)
+    noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
+
+    model_pred = transformer(
+        noisy_model_input,
+        encoder_hidden_states,
+        1000 - student_timesteps,
+        encoder_attention_mask, # B, L
+        return_dict= False,
+    )[0]
+    # calculate X1
+    model_pred = model_pred * sigmas + noisy_model_input
+    return model_pred
+        
+        
+def teacher_forward(transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size, noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale):
+    batch_size = latents.shape[0]
+    noise = torch.randn_like(latents)
+    u = compute_density_for_timestep_sampling(
+        weighting_scheme=weighting_scheme,
+        batch_size=batch_size,
+        logit_mean=logit_mean,
+        logit_std=logit_std,
+        mode_scale=mode_scale,
+    )
+    indices = (u * noise_scheduler.config.num_train_timesteps).long()
+    timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
+    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)
+    noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
+    teacher_features = transformer(
+            noisy_model_input,
+            encoder_hidden_states,
+            noise_scheduler.config.num_train_timesteps - timesteps,
+            encoder_attention_mask, # B, L
+            return_dict= False,
+            output_attn=True
+    )[1]
+    return teacher_features
+    
+def train_one_step_generator(student_transformer, teacher_transformer, discriminator, student_optimizer, loader,noise_scheduler, gradient_accumulation_steps, sp_size,  max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale):
+    total_loss = 0.0
+    teacher_transformer.requires_grad_(False)
+    discriminator.requires_grad_(False)
+    student_transformer.requires_grad_(True)
+    student_optimizer.zero_grad()
+    for _ in range(gradient_accumulation_steps):
+        latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
+        latents = normalize_mochi_dit_input(latents)
+
+        student_pred = student_forward(student_transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size)
+
+        teacher_features = teacher_forward(teacher_transformer, student_pred, encoder_hidden_states, encoder_attention_mask, sp_size,  noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale)
+
+        logits = discriminator(teacher_features, encoder_hidden_states)
+        
+        loss = (F.relu(torch.ones_like(logits) - logits)).mean()  / gradient_accumulation_steps
+        loss.backward()
+        
+        avg_loss = loss.detach().clone()
+        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        total_loss += avg_loss.item() 
+        
+
+    grad_norm = student_transformer.clip_grad_norm_(max_grad_norm)
+    student_optimizer.step()
+    return total_loss, grad_norm.item()
+
+def train_one_step_discriminator(student_transformer, teacher_transformer, discriminator, discriminator_optimizer, loader,noise_scheduler, gradient_accumulation_steps, sp_size,  max_grad_norm, weighting_scheme, logit_mean, logit_std, mode_scale):
+    total_loss = 0.0
+    teacher_transformer.requires_grad_(False)
+    student_transformer.requires_grad_(False)
+    discriminator.requires_grad_(True)
+    discriminator_optimizer.zero_grad()
+    for _ in range(gradient_accumulation_steps):
+        latents, encoder_hidden_states, latents_attention_mask, encoder_attention_mask = next(loader)
+        latents = normalize_mochi_dit_input(latents)
+        with torch.no_grad():
+            student_pred = student_forward(student_transformer, latents, encoder_hidden_states, encoder_attention_mask, sp_size)
+            teacher_features_fake = teacher_forward(teacher_transformer, student_pred, encoder_hidden_states, encoder_attention_mask, sp_size,  noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale)
+            teacher_features_real = teacher_forward(teacher_transformer, latents,      encoder_hidden_states, encoder_attention_mask, sp_size,  noise_scheduler, weighting_scheme, logit_mean, logit_std, mode_scale)
+        fake_logits = discriminator(teacher_features_fake, encoder_hidden_states)
+        real_logits = discriminator(teacher_features_real, encoder_hidden_states)
+        fake_loss = (F.relu(torch.ones_like(fake_logits) + fake_logits)).mean() 
+        real_loss = (F.relu(torch.ones_like(real_logits) - real_logits)).mean()
+
+        loss = (fake_loss + real_loss)/ 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() 
+        
+
+    grad_norm = discriminator.clip_grad_norm_(max_grad_norm)
+    discriminator_optimizer.step()
+    return total_loss, grad_norm.item()
+        
+    
+def main(args):
+    # use LayerNorm, GeLu, SiLu always as fp32 mode
+    # TODO: 
+    if args.enable_stable_fp32:
+        raise NotImplementedError("enable_stable_fp32 is not supported now.")
+    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'])
+    dist.init_process_group("nccl")
+    torch.cuda.set_device(local_rank)
+    device = torch.cuda.current_device()
+    initialize_sequence_parallel_state(args.sp_size)
+
+
+
+
+    # If passed along, set the training seed now. On GPU...
+    if args.seed is not None:
+        # TODO: t within the same seq parallel group should be the same. Noise should be different.
+        set_seed(args.seed + rank)
+
+    # Handle the repository creation
+    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...")
+    student_transformer = MochiTransformer3DModel.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="transformer",
+    )
+    teacher_transformer = MochiTransformer3DModel.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="transformer",
+        load_dtype=torch.bfloat16, # TODO: FP8 or INT8 weight type
+    )
+    # feature_dim, cnn_channel_size, text_cond_dim, num_layers, layers_stride):
+    discriminator = LADDDiscriminator(feature_dim=3072,  text_cond_dim=4096, num_layers=48, layers_stride=4)
+    # get num of discriminator parameters in billion
+    main_print(f"--> discriminator parameters: {sum(p.numel() for p in discriminator.parameters()) / 1e9} B")
+    
+    teacher_transformer.requires_grad_(False)
+    
+    
+    
+    main_print(f"--> Initializing FSDP...")
+    dit_fsdp_kwargs = get_dit_fsdp_kwargs("full")
+    student_transformer = FSDP(
+        student_transformer,
+        **dit_fsdp_kwargs,
+    )
+    teacher_transformer = FSDP(
+        teacher_transformer,
+        **dit_fsdp_kwargs,
+    )
+    discriminator_fsdp_kwargs = get_discriminator_fsdp_kwargs("full")
+    discriminator = FSDP(
+        discriminator,
+        **discriminator_fsdp_kwargs,
+    )
+    
+    
+    main_print(f"--> model loaded")
+    if args.gradient_checkpointing:
+        apply_fsdp_checkpointing(teacher_transformer)
+        apply_fsdp_checkpointing(student_transformer)
+        
+        
+    # Set model as trainable.
+    student_transformer.train()
+
+    noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+
+    student_optimizer = torch.optim.AdamW(
+        student_transformer.parameters(),
+        lr=args.student_learning_rate,
+        betas=(0.9,0.999),
+        weight_decay=0.01,
+        eps=1e-8,
+    )
+    discriminator_optimizer = torch.optim.AdamW(
+        discriminator.parameters(),
+        lr=args.discriminator_learning_rate,
+        betas=(0.9,0.999),
+        weight_decay=0.01,
+        eps=1e-8,
+    )
+    
+    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
+    sampler = DistributedSampler(train_dataset, rank=rank, num_replicas=world_size, shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        collate_fn=latent_collate_function,
+        pin_memory=True,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        drop_last=True, 
+    )
+
+
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if rank <= 0:
+        project = args.tracker_project_name or "fastvideo_distill"
+        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
+    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"  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"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    main_print(f"  Total optimization steps = {args.max_train_steps}")
+    main_print(f"  Total student training parameters per FSDP shard = {sum(p.numel() for p in student_transformer.parameters() if p.requires_grad) / 1e9} B")
+    main_print(f"  Total discriminator training parameters per FSDP shard = {sum(p.numel() for p in discriminator.parameters() if p.requires_grad) / 1e9} B")
+
+    main_print(f"  Master weight dtype: {student_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 
+
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=0,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable= local_rank > 0,
+    )
+
+            
+
+    loader = sp_parallel_dataloader_wrapper(train_dataloader, device, args.train_batch_size, args.sp_size, args.train_sp_batch_size)
+    
+    for step in range(1, args.discriminator_warmup_steps+1):
+        discriminator_loss, discriminator_grad_norm = train_one_step_discriminator(student_transformer, teacher_transformer, discriminator, \
+            discriminator_optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size,  \
+                args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale)
+        progress_bar.set_postfix({"d_loss": discriminator_loss, "d_grad_n": discriminator_grad_norm})
+        progress_bar.update(1)
+        if rank <= 0:
+            wandb.log({"discriminator_loss": discriminator_loss, "discriminator_grad_norm": discriminator_grad_norm}, step=step)
+    
+    for step in range(args.discriminator_warmup_steps+1, args.max_train_steps+1):
+        stage = (step - args.discriminator_warmup_steps) // args.gan_switch_n_steps
+        if stage % 2 == 1:
+            # Train generator
+            generator_loss, student_grad_norm= train_one_step_generator(student_transformer, teacher_transformer, discriminator, \
+                student_optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size,  \
+                    args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale)
+            progress_bar.set_postfix({"g_loss": generator_loss, "g_grad_n": student_grad_norm})
+            if rank <= 0:
+                wandb.log({"generator_loss": generator_loss, "student_grad_norm": student_grad_norm}, step=step)
+            
+        else:
+            # Train discriminator 
+            discriminator_loss, discriminator_grad_norm = train_one_step_discriminator(student_transformer, teacher_transformer, discriminator, \
+                discriminator_optimizer, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size,  \
+                    args.max_grad_norm, args.weighting_scheme, args.logit_mean, args.logit_std, args.mode_scale)
+            progress_bar.set_postfix({"d_loss": discriminator_loss, "d_grad_n": discriminator_grad_norm})
+            if rank <= 0:
+                wandb.log({"discriminator_loss": discriminator_loss, "discriminator_grad_norm": discriminator_grad_norm}, step=step)
+                
+        progress_bar.update(1)
+
+        if args.log_validation and step  % args.validation_steps == 0:
+            log_validation(args, student_transformer, device,
+                            torch.bfloat16, step)
+                    
+        if step  % args.checkpointing_steps == 0:
+            save_checkpoint(student_transformer, rank, args.output_dir, step)
+        dist.barrier()
+
+    if get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # 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
+
+    # 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('--enable_stable_fp32', action='store_true') # TODO
+
+    # 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(
+        "--weighting_scheme",
+        type=str,
+        default="uniform",
+        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."
+    )
+    parser.add_argument(
+        "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+    )
+    parser.add_argument(
+        "--mode_scale",
+        type=float,
+        default=1.29,
+        help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+    )
+    # 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=int, default=64)
+    parser.add_argument('--validation_guidance_scale', type=float, default=4.5)
+    parser.add_argument('--validation_steps', type=float, default=4.5)
+    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("--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("--student_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-4, help="Initial learning rate (after the potential warmup period) to use.")
+    parser.add_argument("--discriminator_warmup_steps", type=int, default=0)
+    parser.add_argument("--gan_switch_n_steps", type=int, default=1)
+    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("--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("--sp_size", type=int, default=1, help="For sequence parallel")
+    parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
+
+    args = parser.parse_args()
+    main(args)

fastvideo/fsdp_util.py

@@ -35,14 +35,33 @@ non_reentrant_wrapper = partial(
 check_fn = lambda submodule: isinstance(submodule, MochiTransformerBlock)
 
 
-def apply_fsdp_checkpointing(model):
+def apply_fsdp_checkpointing(model, p=1):
+    # https://github.com/foundation-model-stack/fms-fsdp/blob/408c7516d69ea9b6bcd4c0f5efab26c0f64b3c2d/fms_fsdp/policies/ac_handler.py#L16
     """apply activation checkpointing to model
     returns None as model is updated directly
     """
     print(f"--> applying fdsp activation checkpointing...")
 
+    block_idx = 0
+    cut_off = 1 / 2
+    # when passing p as a fraction number (e.g. 1/3), it will be interpreted
+    # as a string in argv, thus we need eval("1/3") here for fractions.
+    p = eval(p) if isinstance(p, str) else p
+
+    def selective_checkpointing(submodule):
+        nonlocal block_idx
+        nonlocal cut_off
+
+        if isinstance(submodule, MochiTransformerBlock):
+            block_idx += 1
+            if block_idx * p >= cut_off:
+                cut_off += 1
+                return True
+        return False
+    
+    
     apply_activation_checkpointing(
-        model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
+        model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=selective_checkpointing
     )
 
 bf16_mix = MixedPrecision(
@@ -57,7 +76,7 @@ bf16_mix = MixedPrecision(
 
 
 
-def get_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False):
+def get_dit_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False):
     if use_lora:
         auto_wrap_policy = fsdp_auto_wrap_policy
     else:
@@ -74,6 +93,8 @@ def get_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False):
     
     if sharding_strategy == "full":
         sharding_strategy = ShardingStrategy.FULL_SHARD
+    elif sharding_strategy == "hybrid_full":
+        sharding_strategy = ShardingStrategy.HYBRID_SHARD
     elif sharding_strategy == "none":
         sharding_strategy = ShardingStrategy.NO_SHARD
         auto_wrap_policy = None
@@ -101,4 +122,28 @@ def get_fsdp_kwargs(sharding_strategy, use_lora=False,  cpu_offload=False):
     return fsdp_kwargs
     
     
+        
+        
+
+def get_discriminator_fsdp_kwargs(sharding_strategy):
+
+    auto_wrap_policy = None
+
+
+    # Use existing mixed precision settings
+
+    mixed_precision = bf16_mix
+    sharding_strategy  = ShardingStrategy.NO_SHARD
+    device_id = torch.cuda.current_device()
+    fsdp_kwargs = {
+        "auto_wrap_policy": auto_wrap_policy,
+        "mixed_precision": mixed_precision,
+        "sharding_strategy": sharding_strategy,
+        "device_id": device_id,
+        "limit_all_gathers": True,
+    }
+    
+    return fsdp_kwargs
+    
+    
         

fastvideo/gan/discriminator.py

@@ -0,0 +1,123 @@
+
+import torch
+from torch import nn
+import numpy as np
+from torch.nn.utils.parametrizations import spectral_norm
+import os 
+class DummyDiscriminator(nn.Module):
+    def __init__(self, dim_in, num_layers):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.layers.append(nn.Linear(dim_in, 1))
+        
+    def forward(self, features):
+        logits = []
+        for layer, feature in zip(self.layers, features):
+            mean = feature.mean(dim=1)
+            logits.append(layer(mean))
+        return torch.cat(logits, dim=1)
+    
+
+class ResidualBlock(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (self.fn(x) + x) / np.sqrt(2)
+    
+    
+class SpectralConv1d(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.conv = spectral_norm(nn.Conv1d(*args, **kwargs))
+    def forward(self, x):
+        return self.conv(x)
+    
+class BatchNormLocal(nn.Module):
+    def __init__(self, num_features: int, affine: bool = True, virtual_bs: int = 8, eps: float = 1e-5):
+        super().__init__()
+        self.virtual_bs = virtual_bs
+        self.eps = eps
+        self.affine = affine
+
+        if self.affine:
+            self.weight = nn.Parameter(torch.ones(num_features))
+            self.bias = nn.Parameter(torch.zeros(num_features))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.size()
+
+        # Calculate stats.
+        mean = x.mean([0, 2], keepdim=True)
+        var = x.var([0, 2], keepdim=True, unbiased=False)
+        x = (x - mean) / (torch.sqrt(var + self.eps))
+
+        if self.affine:
+            x = x * self.weight[None, :, None] + self.bias[None, :, None]
+
+        return x.view(shape)
+    
+def make_block(channels: int, kernel_size: int) -> nn.Module:
+    return nn.Sequential(
+        SpectralConv1d(
+            channels,
+            channels,
+            kernel_size = kernel_size,
+            padding = kernel_size//2,
+            padding_mode = 'circular',
+        ),
+        BatchNormLocal(channels),
+        nn.LeakyReLU(0.2, True),
+    )
+    
+class DiscHead(nn.Module):
+    def __init__(self, feature_dim: int, text_c_dim: int, cmap_dim: int = 64, cnn_dim=512):
+        super().__init__()
+        self.channels = feature_dim
+        self.text_c_dim = text_c_dim
+        self.cmap_dim = cmap_dim
+        self.down_proj = SpectralConv1d(feature_dim, cnn_dim, kernel_size=1, padding=0)
+        self.main = nn.Sequential(
+            make_block(cnn_dim, kernel_size=1),
+            ResidualBlock(make_block(cnn_dim, kernel_size=9))
+        )
+
+        self.cmapper = nn.Linear(self.text_c_dim, cmap_dim)
+        self.cls = SpectralConv1d(cnn_dim, cmap_dim, kernel_size=1, padding=0)
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        h = self.down_proj(x)
+        h = self.main(h)
+        out = self.cls(h)
+
+        cmap = self.cmapper(c).unsqueeze(-1)
+        out = (out * cmap).sum(1, keepdim=True) * (1 / np.sqrt(self.cmap_dim))
+
+        return out
+
+
+class LADDDiscriminator(nn.Module):
+    def __init__(self, feature_dim, text_cond_dim, num_layers, layers_stride):
+        super().__init__()
+        heads = []
+        for i in range(0, num_layers, layers_stride):
+            heads.append(DiscHead(feature_dim, text_cond_dim))
+        self.heads = nn.ModuleList(heads)
+        self.layers_stride = layers_stride
+        self.num_layers = num_layers
+        
+    def forward(self, features, text_conditions) -> torch.Tensor:
+        text_conditions = text_conditions.mean(1)
+        # layer, B, L, C -> layer, B, C, L
+        features = features.transpose(2, 3)
+        logits = []
+        for i in range(0, self.num_layers, self.layers_stride):
+            head = self.heads[i//self.layers_stride]
+            feat = features[i]
+            logits.append(head(feat, text_conditions).view(feat.size(0), -1))
+        logits = torch.cat(logits, dim=1)
+
+
+        return logits

fastvideo/model/mochi_latents_stat.py

@@ -1,33 +0,0 @@
-class MochiStat:
-  latents_mean = [
-      -0.06730895953510081,
-      -0.038011381506090416,
-      -0.07477820912866141,
-      -0.05565264470995561,
-      0.012767231469026969,
-      -0.04703542746246419,
-      0.043896967884726704,
-      -0.09346305707025976,
-      -0.09918314763016893,
-      -0.008729793427399178,
-      -0.011931556316503654,
-      -0.0321993391887285
-    ]
-  latents_std = [
-      0.9263795028493863,
-      0.9248894543193766,
-      0.9393059390890617,
-      0.959253732819592,
-      0.8244560132752793,
-      0.917259975397747,
-      0.9294154431013696,
-      1.3720942357788521,
-      0.881393668867029,
-      0.9168315692124348,
-      0.9185249279345552,
-      0.9274757570805041
-      ]
-  scaling_factor = 1.0
-  
-mochi_stat = MochiStat()
-   

fastvideo/model/mochi_latents_utils.py

@@ -0,0 +1,39 @@
+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_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
+
+        

fastvideo/model/modeling_mochi.py

@@ -16,21 +16,87 @@ from typing import Any, Dict, Optional, Tuple
 
 import torch
 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.torch_utils import maybe_allow_in_graph
-from diffusers.models.attention import FeedForward
+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.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.modeling_utils import ModelMixin
 from diffusers.models.normalization import AdaLayerNormContinuous, LuminaLayerNormContinuous, MochiRMSNormZero, RMSNorm
+
+
 from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
 from fastvideo.utils.communications import all_gather, all_to_all_4D
 import torch.nn.functional as F
 from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph
+from einops import rearrange
+
+import numbers
+from flash_attn import flash_attn_varlen_qkvpacked_func
+from flash_attn.bert_padding import pad_input, unpad_input
+
+from liger_kernel.ops.swiglu import LigerSiLUMulFunction
+class FeedForward(HF_FeedForward):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+        inner_dim=None,
+        bias: bool = True,
+    ):
+        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):  
+    # 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_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
+    )
+    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
+    
+                
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+import pdb, sys
+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')
+            pdb.Pdb.interaction(self, *args, **kwargs)
+        finally:
+            sys.stdin = _stdin
+            
 
 
 class MochiAttnProcessor2_0:
@@ -45,6 +111,7 @@ class MochiAttnProcessor2_0:
         attn: Attention,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
@@ -113,38 +180,47 @@ class MochiAttnProcessor2_0:
             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),
-            encoder_key.transpose(1, 2),
-            encoder_value.transpose(1, 2),
-        )
-        # [b, h, s, d]
-        sequence_length = query.size(2)
-        encoder_sequence_length = encoder_query.size(2)
+        # 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),
+        #     encoder_key.transpose(1, 2),
+        #     encoder_value.transpose(1, 2),
+        # )
+        # [b, s, h, d]
+        sequence_length = query.size(1)
+        encoder_sequence_length = encoder_query.size(1)
 
-        # Hint: please check encoder_query.shape
-        query = torch.cat([query, encoder_query], dim=2)
-        key = torch.cat([key, encoder_key], dim=2)
-        value = torch.cat([value, encoder_value], dim=2)
-        
-        
-                
-        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        # H
+        query = torch.cat([query, encoder_query], dim=1).unsqueeze(2)
+        key = torch.cat([key, encoder_key], dim=1).unsqueeze(2)
+        value = torch.cat([value, encoder_value], dim=1).unsqueeze(2)
+        # B, S, 3, H, D
+        qkv = torch.cat([query, key, value], dim=2)
 
+        attn_mask = encoder_attention_mask[:, :].bool()
+        attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True)
+        # 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(
-                (sequence_length, encoder_sequence_length), dim=2
+                (sequence_length, encoder_sequence_length), dim=1
             )
-            # B, H, S, D
-            hidden_states = all_to_all_4D(hidden_states, scatter_dim=2, gather_dim=1)
-            encoder_hidden_states = all_gather(encoder_hidden_states, dim=1).contiguous()
-            hidden_states = hidden_states.transpose(1,2).flatten(2, 3)
+            # 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()
+            hidden_states = hidden_states.flatten(2, 3)
             hidden_states = hidden_states.to(query.dtype)
-            encoder_hidden_states = encoder_hidden_states.transpose(1, 2).flatten(2, 3)
+            encoder_hidden_states = encoder_hidden_states.flatten(2, 3)
             encoder_hidden_states = encoder_hidden_states.to(query.dtype)
         else:
-            hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
+            hidden_states = hidden_states.flatten(2, 3)
             hidden_states = hidden_states.to(query.dtype)
 
             hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
@@ -257,6 +333,7 @@ class MochiTransformerBlock(nn.Module):
         self,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
         temb: torch.Tensor,
         image_rotary_emb: Optional[torch.Tensor] = None,
         output_attn = False, 
@@ -274,6 +351,7 @@ 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
         )
 
         hidden_states = hidden_states + self.norm2(attn_hidden_states) * torch.tanh(gate_msa).unsqueeze(1)
@@ -498,6 +576,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin):
                     create_custom_forward(block),
                     hidden_states,
                     encoder_hidden_states,
+                    encoder_attention_mask,
                     temb,
                     image_rotary_emb,
                     output_attn,
@@ -507,6 +586,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin):
                 hidden_states, encoder_hidden_states, attn_outputs = block(
                     hidden_states=hidden_states,
                     encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
                     temb=temb,
                     image_rotary_emb=image_rotary_emb,
                     output_attn = output_attn,

fastvideo/model/norm.py

@@ -0,0 +1,168 @@
+import numbers
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from liger_kernel.ops.rms_norm import LigerRMSNormFunction
+
+from liger_kernel.ops.layer_norm import LigerLayerNormFunction
+
+class RMSNorm(nn.Module):
+    # mbedding_dim, eps=eps, elementwise_affine=elementwise_affine)
+    def __init__(self, dim, eps: float, elementwise_affine: bool = True):
+        super().__init__()
+
+        self.eps = eps
+
+        if isinstance(dim, numbers.Integral):
+            dim = (dim,)
+
+        self.dim = torch.Size(dim)
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            weight = torch.ones(dim)
+            self.register_buffer("weight", weight, persistent=False)
+            
+    def forward(self, hidden_states):
+  
+        return LigerRMSNormFunction.apply(
+            hidden_states,
+            self.weight,
+            self.eps,
+            0.0,
+            "gemma",
+            True,
+        )
+        
+class LayerNorm(nn.Module):
+    def __init__(self, dim, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True):
+        super().__init__()
+
+        self.eps = eps
+
+        if isinstance(dim, numbers.Integral):
+            dim = (dim,)
+
+        self.dim = torch.Size(dim)
+
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+            self.bias = nn.Parameter(torch.zeros(dim)) if bias else None
+        else:
+            weight = torch.ones(dim)
+            bias = torch.zeros(dim) 
+            self.register_buffer("weight", weight, persistent=False)
+            self.register_buffer("bias", bias, persistent=False)
+    def forward(self, hidden_states):
+        return LigerLayerNormFunction.apply(
+            hidden_states, self.weight, self.bias, self.eps
+        )
+
+class MochiRMSNormZero(nn.Module):
+    r"""
+    Adaptive RMS Norm used in Mochi.
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+    """
+
+    def __init__(
+        self, embedding_dim: int, hidden_dim: int, eps: float = 1e-5, elementwise_affine: bool = False
+    ) -> None:
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, hidden_dim)
+        # TODO: There is a bug in HF. Remove hardcode
+        self.norm = RMSNorm(hidden_dim//4, eps=eps, elementwise_affine=elementwise_affine)
+
+    def forward(
+        self, hidden_states: torch.Tensor, emb: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        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])
+
+        return hidden_states, gate_msa, scale_mlp, gate_mlp
+
+
+class LuminaLayerNormContinuous(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+        out_dim: Optional[int] = None,
+    ):
+        super().__init__()
+
+        # AdaLN
+        self.silu = nn.SiLU()
+        self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
+
+        if norm_type == "layer_norm":
+            self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+        self.linear_2 = None
+        if out_dim is not None:
+            self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias)
+    def forward(
+        self,
+        x: torch.Tensor,
+        conditioning_embedding: torch.Tensor,
+    ) -> torch.Tensor:
+        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
+        emb = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
+        scale = emb
+        x = self.norm(x) * (1 + scale)[:, None, :]
+
+        if self.linear_2 is not None:
+            x = self.linear_2(x)
+
+        return x
+            
+class AdaLayerNormContinuous(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias)
+        if norm_type == "layer_norm":
+            self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = RMSNorm(embedding_dim, eps, elementwise_affine)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+    def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
+        # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
+        emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x

fastvideo/sample/sample_t2v_mochi.py

@@ -17,19 +17,6 @@ import pdb
 import copy
 from typing import Dict
 
-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')
-            pdb.Pdb.interaction(self, *args, **kwargs)
-        finally:
-            sys.stdin = _stdin
-
 def initialize_distributed():
     local_rank = int(os.getenv('RANK', 0))
     world_size = int(os.getenv('WORLD_SIZE', 1))
@@ -121,7 +108,6 @@ def main(args):
     
     if args.transformer_path is not None:
         transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path, torch_dtype=torch.bfloat16)
-        
     else:
         transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder = 'transformer/', torch_dtype=torch.bfloat16)
     if args.lora_checkpoint_dir is not None:
@@ -163,7 +149,7 @@ def main(args):
         if prompts is not None:
             for video, prompt in zip(videos, prompts):
                 suffix = prompt.split(".")[0]
-                export_to_video(video, args.output_path + f"_{suffix}.mp4", fps=30)
+                export_to_video(video, args.output_path + f"{suffix}.mp4", fps=30)
         else:
             export_to_video(videos[0], args.output_path + ".mp4", fps=30)
 

fastvideo/train.py

@@ -9,7 +9,7 @@ from diffusers.training_utils import cast_training_params, compute_density_for_t
 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.model.mochi_latents_stat import mochi_stat
+from fastvideo.model.mochi_latents_utils import normalize_mochi_dit_input
 from fastvideo.utils.validation import log_validation
 import time
 from torch.utils.data import DataLoader
@@ -24,7 +24,7 @@ from torch.utils.data.distributed import DistributedSampler
 import wandb
 from accelerate.utils import set_seed
 from tqdm.auto import tqdm
-from fastvideo.fsdp_util import get_fsdp_kwargs, apply_fsdp_checkpointing
+from fastvideo.fsdp_util import get_dit_fsdp_kwargs, apply_fsdp_checkpointing
 import diffusers
 from diffusers import (
     FlowMatchEulerDiscreteScheduler,
@@ -102,25 +102,22 @@ 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, gradient_accumulation_steps, sp_size, precondition_outputs, max_grad_norm):
+def train_one_step_mochi(transformer, optimizer, lr_scheduler, loader,noise_scheduler, 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 = latents.to(dtype=torch.bfloat16)
-        encoder_hidden_states = encoder_hidden_states.to(dtype=torch.bfloat16)
-        latents_mean = (
-            torch.tensor(mochi_stat.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
-        )
-        latents_std = (
-            torch.tensor(mochi_stat.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
-        )
-        latents = (latents - latents_mean) / latents_std
+        latents = normalize_mochi_dit_input(latents)
         
         batch_size = latents.shape[0]
         noise = torch.randn_like(latents)
-        u = torch.rand(size=(batch_size,), device="cpu")
+        u =   compute_density_for_timestep_sampling(
+            weighting_scheme=weighting_scheme,
+            batch_size=batch_size,
+            logit_mean=logit_mean,
+            logit_std=logit_std,
+            mode_scale=mode_scale,
+        )
         indices = (u * noise_scheduler.config.num_train_timesteps).long()
         timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
         if sp_size > 1:
@@ -145,21 +142,19 @@ def train_one_step_mochi(transformer, optimizer, lr_scheduler, loader,noise_sche
         else:
             target = latents - noise
 
-        loss = torch.mean((model_pred.float() - target.float()) ** 2)
-
+        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() / gradient_accumulation_steps
+        total_loss += avg_loss.item() 
         
 
-    transformer.clip_grad_norm_(max_grad_norm)
-    
+    grad_norm = transformer.clip_grad_norm_(max_grad_norm)
     optimizer.step()
     lr_scheduler.step()
-    
-    return total_loss
+    return total_loss, grad_norm.item()
         
 def get_lora_model(transformer, lora_config):
     transformer.requires_grad_(False)
@@ -286,8 +281,9 @@ def main(args):
         transformer = get_lora_model(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: full")
-    fsdp_kwargs = get_fsdp_kwargs("full", args.use_lora, args.use_cpu_offload)
+    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)
+    
     
     if args.use_lora:
         transformer.config.lora_rank = args.lora_rank
@@ -304,7 +300,7 @@ def main(args):
     main_print(f"--> model loaded")
 
     if args.gradient_checkpointing:
-        apply_fsdp_checkpointing(transformer)
+        apply_fsdp_checkpointing(transformer, args.selective_checkpointing)
 
     # Set model as trainable.
     transformer.train()
@@ -312,8 +308,8 @@ def main(args):
     noise_scheduler = FlowMatchEulerDiscreteScheduler()
 
     params_to_optimize = transformer.parameters()
-    if args.use_lora:
-        params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize))
+    params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize))
+
     optimizer = torch.optim.AdamW(
         params_to_optimize,
         lr=args.learning_rate,
@@ -352,13 +348,10 @@ def main(args):
         drop_last=True, 
     )
 
-    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
-    # Afterwards we recalculate our number of training epochs
     args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
 
-    # We need to initialize the trackers we use, and also store our configuration.
-    # The trackers initializes automatically on the main process.
+
     if rank <= 0:
         project = args.tracker_project_name or "fastvideo"
         wandb.init(project=project, config=args)
@@ -402,7 +395,7 @@ def main(args):
         next(loader)
     for step in range(init_steps + 1, args.max_train_steps+1):
         start_time = time.time()
-        loss = train_one_step_mochi(transformer, optimizer, lr_scheduler, loader, noise_scheduler, args.gradient_accumulation_steps, args.sp_size, args.precondition_outputs, args.max_grad_norm)
+        loss, grad_norm= train_one_step_mochi(transformer, optimizer, lr_scheduler, loader, noise_scheduler, 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)
@@ -410,7 +403,8 @@ def main(args):
 
         progress_bar.set_postfix({
         "loss": f"{loss:.4f}", 
-        "step_time": f"{step_time:.2f}s"
+        "step_time": f"{step_time:.2f}s",
+        "grad_norm": grad_norm
     })
         progress_bar.update(1)
         if rank <= 0:
@@ -418,20 +412,25 @@ def main(args):
             "train_loss": loss,
             "learning_rate": lr_scheduler.get_last_lr()[0],
             "step_time": step_time,
-            "avg_step_time": avg_step_time
+            "avg_step_time": avg_step_time,
+            "grad_norm": grad_norm 
         }, step=step)
         if step  % args.checkpointing_steps == 0:
-            save_checkpoint(transformer, rank, args.output_dir, step)
             if args.use_lora:
                 # Save LoRA weights
                 save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
-            
+            else:
+                # Your existing checkpoint saving code
+                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)
 
     if args.use_lora:
         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)
         
     if get_sequence_parallel_state():
         destroy_sequence_parallel_group()
@@ -506,6 +505,7 @@ if __name__ == "__main__":
     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"
@@ -527,8 +527,28 @@ if __name__ == "__main__":
     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(
+        "--weighting_scheme",
+        type=str,
+        default="uniform",
+        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."
+    )
+    parser.add_argument(
+        "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+    )
+    parser.add_argument(
+        "--mode_scale",
+        type=float,
+        default=1.29,
+        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="cosine_with_restarts",
+    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"]'

fastvideo/utils/encode_vae_t5.py

@@ -1,867 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-"""
-A minimal training script for DiT using PyTorch DDP.
-"""
-import argparse
-import logging
-import math
-import os
-import shutil
-from pathlib import Path
-import gc
-import wandb
-from einops import rearrange
-from tqdm import tqdm
-
-from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
-    destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
-from fastvideo.utils.communications import prepare_parallel_data, broadcast
-import time
-from torch.utils.data import DataLoader
-from copy import deepcopy
-import accelerate
-import torch
-from torch.nn import functional as F
-import transformers
-from accelerate import Accelerator
-from accelerate.logging import get_logger
-from accelerate.utils import DistributedType, ProjectConfiguration, set_seed
-from packaging import version
-from tqdm.auto import tqdm
-
-import diffusers
-from diffusers import (
-    AutoencoderKL,
-    FlowMatchEulerDiscreteScheduler,
-)
-from fastvideo.model.modeling_mochi import MochiTransformer3DModel
-from fastvideo.model.pipeline_mochi import MochiPipeline
-from diffusers.optimization import get_scheduler
-from diffusers.training_utils import compute_snr
-from diffusers.utils import check_min_version, is_wandb_available
-
-from fastvideo.utils.ema import EMAModel
-from fastvideo.dataset import getdataset
-from fastvideo.utils.dataset_utils import Collate, LengthGroupedSampler
-
-
-# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
-check_min_version("0.24.0")
-logger = get_logger(__name__)
-
-
-@torch.inference_mode()
-def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, ema=False):
-    logger.info(f"Running validation....\n")
-    model = accelerator.unwrap_model(model)
-    scheduler = FlowMatchEulerDiscreteScheduler()
-    mochi_pipeline = MochiPipeline(vae=vae,
-                                         text_encoder=text_encoder,
-                                         tokenizer=tokenizer,
-                                         scheduler=scheduler,
-                                         transformer=model).to(device=accelerator.device)
-    videos = []
-    for prompt in args.validaiton_prompt:
-        logger.info('Processing the ({}) prompt'.format(prompt))
-        video = mochi_pipeline(
-                                args.prompt,
-                                num_frames=args.num_frames,
-                                height=args.max_height,
-                                width=args.max_width,
-                                num_inference_steps=args.validation_sampling_steps,
-                                guidance_scale=args.guidance_scale,
-                                ).images
-        videos.append(video[0])
-    # import ipdb;ipdb.set_trace()
-    gc.collect()
-    torch.cuda.empty_cache()
-    videos = torch.stack(videos).numpy()
-    videos = rearrange(videos, 'b t h w c -> b t c h w')
-    for tracker in accelerator.trackers:
-        if videos.shape[1] == 1:
-            images = rearrange(videos, 'b 1 c h w -> (b 1) h w c')
-            logs = {
-                f"{'ema_' if ema else ''}validation": [
-                    wandb.Image(image, caption=f"{i}: {prompt}")
-                    for i, (image, prompt) in enumerate(zip(images, args.validaiton_prompt))
-                ]
-            }
-        else:
-            logs = {
-                f"{'ema_' if ema else ''}validation": [
-                    wandb.Video(video, caption=f"{i}: {prompt}", fps=24)
-                    for i, (video, prompt) in enumerate(zip(videos, args.validaiton_prompt))
-                ]
-            }
-        tracker.log(logs, step=global_step)
-
-    del mochi_pipeline
-    gc.collect()
-    torch.cuda.empty_cache()
-
-
-class ProgressInfo:
-    def __init__(self, global_step, train_loss=0.0):
-        self.global_step = global_step
-        self.train_loss = train_loss
-
-
-#################################################################################
-#                                  Training Loop                                #
-#################################################################################
-
-def main(args):
-    logging_dir = Path(args.output_dir, args.logging_dir)
-
-    # use LayerNorm, GeLu, SiLu always as fp32 mode
-    # TODO: 
-    if args.enable_stable_fp32:
-        replace_with_fp32_forwards()
-
-    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
-
-    accelerator = Accelerator(
-        gradient_accumulation_steps=args.gradient_accumulation_steps,
-        mixed_precision=args.mixed_precision,
-        log_with=args.report_to,
-        project_config=accelerator_project_config,
-        log_with="wandb"
-    )
-
-    if args.num_frames != 1 and args.use_image_num == 0:
-        initialize_sequence_parallel_state(args.sp_size)
-
-    if args.report_to == "wandb":
-        if not is_wandb_available():
-            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
-
-    # Make one log on every process with the configuration for debugging.
-    logging.basicConfig(
-        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
-        datefmt="%m/%d/%Y %H:%M:%S",
-        level=logging.INFO,
-    )
-    logger.info(accelerator.state, main_process_only=False)
-    if accelerator.is_local_main_process:
-        transformers.utils.logging.set_verbosity_warning()
-        diffusers.utils.logging.set_verbosity_info()
-    else:
-        transformers.utils.logging.set_verbosity_error()
-        diffusers.utils.logging.set_verbosity_error()
-
-    # If passed along, set the training seed now. On GPU...
-    if args.seed is not None:
-        set_seed(args.seed)
-
-    # Handle the repository creation
-    if accelerator.is_main_process:
-        if args.output_dir is not None:
-            os.makedirs(args.output_dir, exist_ok=True)
-
-    # For mixed precision training we cast all non-trainable weigths to half-precision
-    # as these weights are only used for inference, keeping weights in full precision is not required.
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-
-    # Create model:
-    kwargs = {}
-    ae = CausalVAEModelWrapper(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval()
-    if args.enable_tiling:
-        ae.vae.enable_tiling()
-        ae.vae.tile_sample_stride = args.tile_sample_stride
-
-    kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True}
-    text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval()
-
-    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
-    ae.vae_scale_factor = (ae_stride_t, ae_stride_h, ae_stride_w)
-    assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})"
-    args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w
-    args.ae_stride = args.ae_stride_h
-    patch_size = args.model[-3:]
-    patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2])
-    args.patch_size = patch_size_h
-    args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w
-    assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})"
-    # assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})."
-    assert args.max_height % ae_stride_h == 0, f"Height must be divisible by ae_stride_h, but found Height ({args.max_height}), ae_stride_h ({ae_stride_h})."
-    assert args.max_width % ae_stride_h == 0, f"Width size must be divisible by ae_stride_h, but found Width ({args.max_width}), ae_stride_h ({ae_stride_h})."
-
-    args.stride_t = ae_stride_t * patch_size_t
-    args.stride = ae_stride_h * patch_size_h
-    latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w)
-    ae.latent_size = latent_size
-
-    if args.num_frames % 2 == 1:
-        args.latent_size_t = latent_size_t = (args.num_frames - 1) // ae_stride_t + 1
-    else:
-        latent_size_t = args.num_frames // ae_stride_t
-    model = Diffusion_models[args.model](
-        in_channels=ae_channel_config[args.ae],
-        out_channels=ae_channel_config[args.ae],
-        # caption_channels=4096,
-        # cross_attention_dim=1152,
-        attention_bias=True,
-        sample_size=latent_size,
-        sample_size_t=latent_size_t,
-        num_vector_embeds=None,
-        activation_fn="gelu-approximate",
-        num_embeds_ada_norm=1000,
-        use_linear_projection=False,
-        only_cross_attention=False,
-        double_self_attention=False,
-        upcast_attention=False,
-        # norm_type="ada_norm_single",
-        norm_elementwise_affine=False,
-        norm_eps=1e-6,
-        attention_type='default',
-        attention_mode=args.attention_mode,
-        interpolation_scale_h=args.interpolation_scale_h,
-        interpolation_scale_w=args.interpolation_scale_w,
-        interpolation_scale_t=args.interpolation_scale_t,
-        downsampler=args.downsampler,
-        use_rope=args.use_rope,
-        # text_max_length=args.text_max_length,
-        use_stable_fp32=args.enable_stable_fp32, 
-    )
-    model.gradient_checkpointing = args.gradient_checkpointing
-
-    # # use pretrained model?
-    if args.pretrained:
-        model_state_dict = model.state_dict()
-        if 'safetensors' in args.pretrained:  # pixart series
-            from safetensors.torch import load_file as safe_load
-            # import ipdb;ipdb.set_trace()
-            pretrained_checkpoint = safe_load(args.pretrained, device="cpu")
-            pretrained_keys = set(list(pretrained_checkpoint.keys()))
-            model_keys = set(list(model_state_dict.keys()))
-            common_keys = list(pretrained_keys & model_keys)
-            checkpoint = {k: pretrained_checkpoint[k] for k in common_keys if model_state_dict[k].numel() == pretrained_checkpoint[k].numel()}
-            # if checkpoint['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and checkpoint['pos_embed.proj.weight'].ndim == 4:
-            #     logger.info(f"Resize pos_embed, {checkpoint['pos_embed.proj.weight'].shape} -> {model.pos_embed.proj.weight.shape}")
-            #     repeat = model.pos_embed.proj.weight.shape[2]
-            #     checkpoint['pos_embed.proj.weight'] = checkpoint['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat)
-                # del checkpoint['proj_out.weight'], checkpoint['proj_out.bias']
-        else:  # latest stage training weight
-            checkpoint = torch.load(args.pretrained, map_location='cpu')
-            if 'model' in checkpoint:
-                checkpoint = checkpoint['model']
-        missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
-        logger.info(f'missing_keys {len(missing_keys)} {missing_keys}, unexpected_keys {len(unexpected_keys)}')
-        logger.info(f'Successfully load {len(model_state_dict) - len(missing_keys)}/{len(model_state_dict)} keys from {args.pretrained}!')
-
-    # Freeze vae and text encoders.
-    ae.vae.requires_grad_(False)
-    text_enc.requires_grad_(False)
-    # Set model as trainable.
-    model.train()
-
-    noise_scheduler = DDPMScheduler()
-    # Move unet, vae and text_encoder to device and cast to weight_dtype
-    # The VAE is in float32 to avoid NaN losses.
-    ae.vae.to(accelerator.device, dtype=torch.float32)
-    # ae.vae.to(accelerator.device, dtype=weight_dtype)
-    text_enc.to(accelerator.device, dtype=weight_dtype)
-
-    # Create EMA for the unet.
-    if args.use_ema:
-        ema_model = deepcopy(model)
-        ema_model = EMAModel(ema_model.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step,
-                             model_cls=Diffusion_models_class[args.model], model_config=ema_model.config)
-
-    # `accelerate` 0.16.0 will have better support for customized saving
-    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
-        # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
-        def save_model_hook(models, weights, output_dir):
-            if accelerator.is_main_process:
-                if args.use_ema:
-                    ema_model.save_pretrained(os.path.join(output_dir, "model_ema"))
-
-                for i, model in enumerate(models):
-                    model.save_pretrained(os.path.join(output_dir, "model"))
-                    if weights:  # Don't pop if empty
-                        # make sure to pop weight so that corresponding model is not saved again
-                        weights.pop()
-
-        def load_model_hook(models, input_dir):
-            if args.use_ema:
-                load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), Diffusion_models_class[args.model])
-                ema_model.load_state_dict(load_model.state_dict())
-                ema_model.to(accelerator.device)
-                del load_model
-
-            for i in range(len(models)):
-                # pop models so that they are not loaded again
-                model = models.pop()
-
-                # load diffusers style into model
-                load_model = Diffusion_models_class[args.model].from_pretrained(input_dir, subfolder="model")
-                model.register_to_config(**load_model.config)
-
-                model.load_state_dict(load_model.state_dict())
-                del load_model
-
-        accelerator.register_save_state_pre_hook(save_model_hook)
-        accelerator.register_load_state_pre_hook(load_model_hook)
-
-    # Enable TF32 for faster training on Ampere GPUs,
-    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
-    if args.allow_tf32:
-        torch.backends.cuda.matmul.allow_tf32 = True
-
-    if args.scale_lr:
-        args.learning_rate = (
-                args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
-        )
-
-    params_to_optimize = model.parameters()
-    # Optimizer creation
-    if not (args.optimizer.lower() == "prodigy" or args.optimizer.lower() == "adamw"):
-        logger.warning(
-            f"Unsupported choice of optimizer: {args.optimizer}.Supported optimizers include [adamW, prodigy]."
-            "Defaulting to adamW"
-        )
-        args.optimizer = "adamw"
-
-    if args.use_8bit_adam and not args.optimizer.lower() == "adamw":
-        logger.warning(
-            f"use_8bit_adam is ignored when optimizer is not set to 'AdamW'. Optimizer was "
-            f"set to {args.optimizer.lower()}"
-        )
-
-    if args.optimizer.lower() == "adamw":
-        if args.use_8bit_adam:
-            try:
-                import bitsandbytes as bnb
-            except ImportError:
-                raise ImportError(
-                    "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
-                )
-
-            optimizer_class = bnb.optim.AdamW8bit
-        else:
-            optimizer_class = torch.optim.AdamW
-
-        optimizer = optimizer_class(
-            params_to_optimize,
-            lr=args.learning_rate,
-            betas=(args.adam_beta1, args.adam_beta2),
-            weight_decay=args.adam_weight_decay,
-            eps=args.adam_epsilon,
-        )
-
-    if args.optimizer.lower() == "prodigy":
-        try:
-            import prodigyopt
-        except ImportError:
-            raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
-
-        optimizer_class = prodigyopt.Prodigy
-
-        if args.learning_rate <= 0.1:
-            logger.warning(
-                "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
-            )
-
-        optimizer = optimizer_class(
-            params_to_optimize,
-            lr=args.learning_rate,
-            betas=(args.adam_beta1, args.adam_beta2),
-            beta3=args.prodigy_beta3,
-            weight_decay=args.adam_weight_decay,
-            eps=args.adam_epsilon,
-            decouple=args.prodigy_decouple,
-            use_bias_correction=args.prodigy_use_bias_correction,
-            safeguard_warmup=args.prodigy_safeguard_warmup,
-        )
-    logger.info(f"optimizer: {optimizer}")
-    
-    # Setup data:
-    train_dataset = getdataset(args)
-    sampler = LengthGroupedSampler(
-                args.train_batch_size,
-                world_size=accelerator.num_processes,
-                lengths=train_dataset.lengths, 
-                group_frame=args.group_frame, 
-                group_resolution=args.group_resolution, 
-            ) if (args.group_frame or args.group_resolution) else None
-    train_dataloader = DataLoader(
-        train_dataset,
-        shuffle=sampler is None,
-        # pin_memory=True,
-        collate_fn=Collate(args),
-        batch_size=args.train_batch_size,
-        num_workers=args.dataloader_num_workers,
-        sampler=sampler if (args.group_frame or args.group_resolution) else None, 
-        drop_last=True, 
-        # prefetch_factor=4
-    )
-    logger.info(f'after train_dataloader')
-
-    # Scheduler and math around the number of training steps.
-    overrode_max_train_steps = False
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-    if args.max_train_steps is None:
-        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
-        overrode_max_train_steps = True
-
-    lr_scheduler = get_scheduler(
-        args.lr_scheduler,
-        optimizer=optimizer,
-        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
-        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
-    )
-
-    # Prepare everything with our `accelerator`.
-    # model.requires_grad_(False)
-    # model.pos_embed.requires_grad_(True)
-    logger.info(f'before accelerator.prepare')
-    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        model, optimizer, train_dataloader, lr_scheduler
-    )
-    logger.info(f'after accelerator.prepare')
-    if args.use_ema:
-        ema_model.to(accelerator.device)
-
-    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
-    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
-    if overrode_max_train_steps:
-        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
-    # Afterwards we recalculate our number of training epochs
-    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
-
-    # We need to initialize the trackers we use, and also store our configuration.
-    # The trackers initializes automatically on the main process.
-    if accelerator.is_main_process:
-        accelerator.init_trackers(os.path.basename(args.output_dir), config=vars(args))
-
-    # Train!
-    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
-    total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size
-    logger.info("***** Running training *****")
-    logger.info(f"  Num examples = {len(train_dataset)}")
-    logger.info(f"  Num Epochs = {args.num_train_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
-    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
-    logger.info(f"  Total optimization steps = {args.max_train_steps}")
-    logger.info(f"  Total training parameters = {sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B")
-    global_step = 0
-    first_epoch = 0
-
-    # Potentially load in the weights and states from a previous save
-    if args.resume_from_checkpoint:
-        if args.resume_from_checkpoint != "latest":
-            path = os.path.basename(args.resume_from_checkpoint)
-        else:
-            # Get the most recent checkpoint
-            dirs = os.listdir(args.output_dir)
-            dirs = [d for d in dirs if d.startswith("checkpoint")]
-            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-            path = dirs[-1] if len(dirs) > 0 else None
-
-        if path is None:
-            accelerator.print(
-                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
-            )
-            args.resume_from_checkpoint = None
-            initial_global_step = 0
-        else:
-            accelerator.print(f"Resuming from checkpoint {path}")
-            accelerator.load_state(os.path.join(args.output_dir, path))
-            global_step = int(path.split("-")[1])
-
-            initial_global_step = global_step
-            first_epoch = global_step // num_update_steps_per_epoch
-
-
-
-    else:
-        initial_global_step = 0
-
-    progress_bar = tqdm(
-        range(0, args.max_train_steps),
-        initial=initial_global_step,
-        desc="Steps",
-        # Only show the progress bar once on each machine.
-        disable=not accelerator.is_local_main_process,
-    )
-    progress_info = ProgressInfo(global_step, train_loss=0.0)
-
-    def sync_gradients_info(loss):
-        # Checks if the accelerator has performed an optimization step behind the scenes
-        if args.use_ema:
-            ema_model.step(model.parameters())
-        progress_bar.update(1)
-        progress_info.global_step += 1
-        end_time = time.time()
-        one_step_duration = end_time - start_time
-        accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step)
-        progress_info.train_loss = 0.0
-
-        # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
-        if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
-            if progress_info.global_step % args.checkpointing_steps == 0:
-                # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
-                if accelerator.is_main_process and args.checkpoints_total_limit is not None:
-                    checkpoints = os.listdir(args.output_dir)
-                    checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
-                    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
-
-                    # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
-                    if len(checkpoints) >= args.checkpoints_total_limit:
-                        num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
-                        removing_checkpoints = checkpoints[0:num_to_remove]
-
-                        logger.info(
-                            f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
-                        )
-                        logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
-
-                        for removing_checkpoint in removing_checkpoints:
-                            removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
-                            shutil.rmtree(removing_checkpoint)
-
-                save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}")
-                accelerator.save_state(save_path)
-                logger.info(f"Saved state to {save_path}")
-
-        logs = {"step_loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
-        progress_bar.set_postfix(**logs)
-
-    def run(model_input, model_kwargs, prof):
-        global start_time
-        start_time = time.time()
-
-        noise = torch.randn_like(model_input)
-        if args.noise_offset:
-            # https://www.crosslabs.org//blog/diffusion-with-offset-noise
-            noise += args.noise_offset * torch.randn((model_input.shape[0], model_input.shape[1], 1, 1, 1),
-                                                     device=model_input.device)
-
-        bsz = model_input.shape[0]
-        current_step_frame = model_input.shape[2]
-        # Sample a random timestep for each image without bias.
-        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device)
-        if current_step_frame != 1 and get_sequence_parallel_state():  # image do not need sp
-            broadcast(timesteps)
-
-        # Add noise to the model input according to the noise magnitude at each timestep
-        # (this is the forward diffusion process)
-
-        noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)
-
-        model_pred = model(
-            noisy_model_input,
-            timesteps,
-            **model_kwargs
-        )[0]
-        # Get the target for loss depending on the prediction type
-        if args.prediction_type is not None:
-            # set prediction_type of scheduler if defined
-            noise_scheduler.register_to_config(prediction_type=args.prediction_type)
-
-        if noise_scheduler.config.prediction_type == "epsilon":
-            target = noise
-        elif noise_scheduler.config.prediction_type == "v_prediction":
-            target = noise_scheduler.get_velocity(model_input, noise, timesteps)
-        elif noise_scheduler.config.prediction_type == "sample":
-            # We set the target to latents here, but the model_pred will return the noise sample prediction.
-            target = model_input
-            # We will have to subtract the noise residual from the prediction to get the target sample.
-            model_pred = model_pred - noise
-        else:
-            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
-
-        mask = model_kwargs.get('attention_mask', None)
-        if torch.all(mask.bool()):
-            mask = None
-        if get_sequence_parallel_state():
-            assert mask is None
-        b, c, _, _, _ = model_pred.shape
-        if mask is not None:
-            mask = mask.unsqueeze(1).repeat(1, c, 1, 1, 1).float()  # b t h w -> b c t h w
-            mask = mask.reshape(b, -1)
-        if args.snr_gamma is None:
-            # model_pred: b c t h w, attention_mask: b t h w
-            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
-            loss = loss.reshape(b, -1)
-            if mask is not None:
-                loss = (loss * mask).sum() / mask.sum()  # mean loss on unpad patches
-            else:
-                loss = loss.mean()
-        else:
-            # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
-            # Since we predict the noise instead of x_0, the original formulation is slightly changed.
-            # This is discussed in Section 4.2 of the same paper.
-            snr = compute_snr(noise_scheduler, timesteps)
-            mse_loss_weights = torch.stack([snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(
-                dim=1
-            )[0]
-            if noise_scheduler.config.prediction_type == "epsilon":
-                mse_loss_weights = mse_loss_weights / snr
-            elif noise_scheduler.config.prediction_type == "v_prediction":
-                mse_loss_weights = mse_loss_weights / (snr + 1)
-            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
-            loss = loss.reshape(b, -1)
-            mse_loss_weights = mse_loss_weights.reshape(b, 1)
-            if mask is not None:
-                loss = (loss * mask * mse_loss_weights).sum() / mask.sum()  # mean loss on unpad patches
-            else:
-                loss = (loss * mse_loss_weights).mean()
-
-        # Gather the losses across all processes for logging (if we use distributed training).
-        avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()
-        progress_info.train_loss += avg_loss.detach().item() / args.gradient_accumulation_steps
-
-        # Backpropagate
-        accelerator.backward(loss)
-        if accelerator.sync_gradients:
-            params_to_clip = model.parameters()
-            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
-        optimizer.step()
-        lr_scheduler.step()
-        optimizer.zero_grad()
-
-        if accelerator.sync_gradients:
-            sync_gradients_info(loss)
-
-        if accelerator.is_main_process:
-
-            if progress_info.global_step % args.checkpointing_steps == 0:
-
-                if args.enable_tracker:
-                    log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator,
-                                   weight_dtype, progress_info.global_step)
-
-                    if args.use_ema:
-                        # Store the UNet parameters temporarily and load the EMA parameters to perform inference.
-                        ema_model.store(model.parameters())
-                        ema_model.copy_to(model.parameters())
-                        log_validation(args, model, ae, text_enc.text_enc, train_dataset.tokenizer, accelerator,
-                                    weight_dtype, progress_info.global_step, ema=True)
-                        # Switch back to the original UNet parameters.
-                        ema_model.restore(model.parameters())
-
-        if prof is not None:
-            prof.step()
-
-
-        return loss
-
-    def train_one_step(step_, data_item_, prof_=None):
-        train_loss = 0.0
-        x, attn_mask, input_ids, cond_mask = data_item_
-        if args.group_frame or args.group_resolution:
-            if not args.group_frame:
-                each_latent_frame = torch.any(attn_mask.flatten(-2), dim=-1).int().sum(-1).tolist()
-                # logger.info(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask '
-                #             f'each_latent_frame: {each_latent_frame}')
-                print(f'rank: {accelerator.process_index}, step {step_}, special batch has attention_mask '
-                            f'each_latent_frame: {each_latent_frame}')
-        assert not torch.any(torch.isnan(x)), 'torch.any(torch.isnan(x))'
-        x = x.to(accelerator.device, dtype=ae.vae.dtype)  # B C T+num_images H W, 16 + 4
-
-        attn_mask = attn_mask.to(accelerator.device)  # B T+num_images H W
-        input_ids = input_ids.to(accelerator.device)  # B 1+num_images L
-        cond_mask = cond_mask.to(accelerator.device)  # B 1+num_images L
-        # if accelerator.process_index == 0:
-        #     logger.info(f'rank: {accelerator.process_index}, x: {x.shape}, attn_mask: {attn_mask.shape}')
-
-        with torch.no_grad():
-            B, N, L = input_ids.shape  # B 1+num_images L
-            input_ids_ = input_ids.reshape(-1, L)
-            cond_mask_ = cond_mask.reshape(-1, L)
-            cond = text_enc(input_ids_, cond_mask_)  # B 1+num_images L D
-            cond = cond.reshape(B, N, L, -1)
-            x = ae.encode(x)  # B C T H W
-        current_step_frame = x.shape[2]
-        current_step_sp_state = get_sequence_parallel_state()
-        if args.sp_size != 1:  # enable sp
-            if current_step_frame == 1:  # but image do not need sp
-                set_sequence_parallel_state(False)
-            else:
-                set_sequence_parallel_state(True)
-        if get_sequence_parallel_state():
-            x, cond, attn_mask, cond_mask, use_image_num = prepare_parallel_data(x, cond, attn_mask, cond_mask,
-                                                                                 args.use_image_num)
-            for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size):
-                with accelerator.accumulate(model):
-                    st_idx = iter * args.train_sp_batch_size
-                    ed_idx = (iter + 1) * args.train_sp_batch_size
-                    model_kwargs = dict(encoder_hidden_states=cond[st_idx: ed_idx],
-                                        attention_mask=attn_mask[st_idx: ed_idx],
-                                        encoder_attention_mask=cond_mask[st_idx: ed_idx], use_image_num=use_image_num)
-                    run(x[st_idx: ed_idx], model_kwargs, prof_)
-
-        else:
-            with accelerator.accumulate(model):
-                assert not torch.any(torch.isnan(x)), 'after vae'
-                x = x.to(weight_dtype)
-                model_kwargs = dict(encoder_hidden_states=cond, attention_mask=attn_mask,
-                                    encoder_attention_mask=cond_mask, use_image_num=args.use_image_num)
-                run(x, model_kwargs, prof_)
-
-        set_sequence_parallel_state(current_step_sp_state)  # in case the next step use sp, which need broadcast(timesteps)
-
-        if progress_info.global_step >= args.max_train_steps:
-            return True
-
-        return False
-
-    def train_all_epoch(prof_=None):
-        for epoch in range(first_epoch, args.num_train_epochs):
-            progress_info.train_loss = 0.0
-            if progress_info.global_step >= args.max_train_steps:
-                return True
-
-            for step, data_item in enumerate(train_dataloader):
-                if train_one_step(step, data_item, prof_):
-                    break
-
-
-    train_all_epoch()
-    accelerator.wait_for_everyone()
-    accelerator.end_training()
-    if get_sequence_parallel_state():
-        destroy_sequence_parallel_group()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    # dataset & dataloader
-    parser.add_argument("--dataset", type=str, required=True)
-    parser.add_argument("--data", type=str, required='')
-    parser.add_argument("--cache_vae_latent", action="store_true")
-    parser.add_argument("--train_fps", type=int, default=24)
-    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
-    parser.add_argument("--speed_factor", type=float, default=1.0)
-    parser.add_argument("--num_frames", type=int, default=65)
-    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
-    parser.add_argument("--max_height", type=int, default=320)
-    parser.add_argument("--max_width", type=int, default=240)
-    parser.add_argument("--use_image_num", type=int, default=0)
-    parser.add_argument("--text_max_length", type=int, default=512)
-    parser.add_argument('--cfg', type=float, default=0.1)
-    parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.")
-    parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.")
-    parser.add_argument("--group_frame", action="store_true")
-    parser.add_argument("--group_resolution", action="store_true")
-
-    # text encoder & vae & diffusion model
-    parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122")
-    parser.add_argument('--enable_8bit_t5', action='store_true')
-    parser.add_argument('--tile_sample_stride', type=float, default=0.125)
-    parser.add_argument('--enable_tiling', action='store_true')
-    parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers")
-    parser.add_argument('--use_rope', action='store_true')
-    parser.add_argument('--interpolation_scale_h', type=float, default=1.0)
-    parser.add_argument('--interpolation_scale_w', type=float, default=1.0)
-    parser.add_argument('--interpolation_scale_t', type=float, default=1.0)
-    parser.add_argument("--downsampler", type=str, default=None)
-    parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse")
-    parser.add_argument("--ae_path", type=str, default="stabilityai/sd-vae-ft-mse")
-    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
-    parser.add_argument("--pretrained", type=str, default=None)
-    parser.add_argument('--enable_stable_fp32', action='store_true')
-    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
-
-    # diffusion setting
-    parser.add_argument("--snr_gamma", type=float, default=None, help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. More details here: https://arxiv.org/abs/2303.09556.")
-    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
-    parser.add_argument("--ema_decay", type=float, default=0.999)
-    parser.add_argument("--ema_start_step", type=int, default=0)
-    parser.add_argument("--noise_offset", type=float, default=0.02, help="The scale of noise offset.")
-    parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
-
-    # validation & logs
-    parser.add_argument("--validaiton_prompt", nargs='+', help="List of prompts to use for validation.")
-    parser.add_argument("--validation_sampling_steps", type=int, default=64)
-    parser.add_argument('--guidance_scale', type=float, default=4.5)
-    parser.add_argument("--enable_tracker", action="store_true")
-    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
-    parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.")
-    parser.add_argument("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store."))
-    parser.add_argument("--checkpointing_steps", type=int, default=500,
-                        help=(
-                            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
-                            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
-                            " training using `--resume_from_checkpoint`."
-                        ),
-                        )
-    parser.add_argument("--resume_from_checkpoint", type=str, default=None,
-                        help=(
-                            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
-                            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
-                        ),
-                        )
-    parser.add_argument("--logging_dir", type=str, default="logs",
-                        help=(
-                            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
-                            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
-                        ),
-                        )
-    parser.add_argument("--report_to", type=str, default="tensorboard",
-                        help=(
-                            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
-                            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
-                        ),
-                        )
-    # optimizer & scheduler
-    parser.add_argument("--num_train_epochs", type=int, default=100)
-    parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform.  If provided, overrides num_train_epochs.")
-    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")
-    parser.add_argument("--optimizer", type=str, default="adamW", help='The optimizer type to use. Choose between ["AdamW", "prodigy"]')
-    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")
-    parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
-    parser.add_argument("--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler.")
-    parser.add_argument("--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW")
-    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers.")
-    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers.")
-    parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
-    parser.add_argument("--adam_weight_decay", type=float, default=1e-02, help="Weight decay to use for unet params")
-    parser.add_argument("--adam_weight_decay_text_encoder", type=float, default=None, help="Weight decay to use for text_encoder")
-    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer and Prodigy optimizers.")
-    parser.add_argument("--prodigy_use_bias_correction", type=bool, default=True, help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW")
-    parser.add_argument("--prodigy_safeguard_warmup", type=bool, default=True, help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. Ignored if optimizer is adamW")
-    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
-    parser.add_argument("--prodigy_beta3", type=float, default=None,
-                        help="coefficients for computing the Prodidy stepsize using running averages. If set to None, "
-                             "uses the value of square root of beta2. Ignored if optimizer is adamW",
-                        )
-    parser.add_argument("--lr_scheduler", type=str, default="constant",
-                        help=(
-                            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
-                            ' "constant", "constant_with_warmup"]'
-                        ),
-                        )
-    parser.add_argument("--allow_tf32", action="store_true",
-                        help=(
-                            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
-                            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
-                        ),
-                        )
-    parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"],
-                        help=(
-                            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
-                            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
-                            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
-                        ),
-                        )
-
-    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
-    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
-    parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
-
-    args = parser.parse_args()
-    main(args)

inference_mochi.sh

@@ -1,24 +1,87 @@
-#!/bin/bash
+num_gpus=2
+prompts=(
+    "In this animated scene, Tom the cat is peeking out from behind a tall stack of books in a cozy library. His eyes are wide with curiosity as he spots Jerry, the tiny brown mouse, sitting on a book cover at the top of the pile. Jerry is holding a small, crumpled piece of cheese in his paws, looking down at Tom with a mischievous grin. The style of the animation brings warmth to the scene with rich, textured colors that highlight the humorous interaction between the characters."
+    "In this playful scene, Tom, the gray-blue cat, is dressed as a chef, complete with a white hat and apron. He’s standing over a steaming pot on a stove, stirring with a large spoon. Meanwhile, Jerry, the brown mouse, is perched on the edge of the counter, tossing in tiny sprigs of herbs. Their expressions suggest a lighthearted teamwork moment. The animation style emphasizes vibrant colors and exaggerated movements that enhance the lively interaction between the characters."
+)
+
+for prompt in "${prompts[@]}"; do
+    torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+        fastvideo/sample/sample_t2v_mochi.py \
+        --model_path data/mochi \
+        --prompts "$prompt" \
+        --num_frames 79 \
+        --height 480 \
+        --width 848 \
+        --num_inference_steps 64 \
+        --guidance_scale 4.5 \
+        --output_path "outputs_video/T_J_original" \
+        --seed 12345
+done
+
+
+
+
+
+
+
+
+num_gpus=2
+prompts=(
+    "In this animated scene, Tom the cat is peeking out from behind a tall stack of books in a cozy library. His eyes are wide with curiosity as he spots Jerry, the tiny brown mouse, sitting on a book cover at the top of the pile. Jerry is holding a small, crumpled piece of cheese in his paws, looking down at Tom with a mischievous grin. The style of the animation brings warmth to the scene with rich, textured colors that highlight the humorous interaction between the characters."
+    "In this playful scene, Tom, the gray-blue cat, is dressed as a chef, complete with a white hat and apron. He’s standing over a steaming pot on a stove, stirring with a large spoon. Meanwhile, Jerry, the brown mouse, is perched on the edge of the counter, tossing in tiny sprigs of herbs. Their expressions suggest a lighthearted teamwork moment. The animation style emphasizes vibrant colors and exaggerated movements that enhance the lively interaction between the characters."
+)
+
+for prompt in "${prompts[@]}"; do
+    torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+        fastvideo/sample/sample_t2v_mochi.py \
+        --model_path data/mochi \
+        --prompts "$prompt" \
+        --num_frames 79 \
+        --height 480 \
+        --width 848 \
+        --num_inference_steps 64 \
+        --guidance_scale 4.5 \
+        --output_path "outputs_video/T_J_FT_normal_79_1000" \
+        --transformer_path data/outputs/T_J_FT/checkpoint-1000 \
+        --seed 12345
+done
+
+
+
+
+
+
+
+black myth: a character with a pig-like face, dressed in a green robe and holding a string of black beads. the character is riding a white horse outdoors, with a mountainous landscape visible in the background. The horse is a strong, muscular white steed with a slightly wild, untamed mane flowing in the wind, suggesting speed and vigor. Its eyes are alert, with a slight intensity that complements the adventurous aura of the character. 
+black myth: a character dressed in elaborate, ornate armor that is richly decorated with intricate designs and patterns. the armor appears to be of a high-quality material, possibly metal, and is adorned with various embellishments such as jewels and metallic accents. the character is riding a white horse outdoors, with a mountainous landscape visible in the background. the lighting suggests it might be either dawn or dusk, casting a soft glow on the scene. the character's pose is dynamic, suggesting movement or action.
+
+
+
 
 num_gpus=4
-
-torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 \
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/sample_t2v_mochi.py \
     --model_path data/mochi \
-    --prompts "A cute fluffy panda eating Chinese food in a restaurant" \
-    --num_frames 163 \
+    --prompts  "A hyper-realistic sci-fi scene featuring a brave monkey astronaut, floating in zero gravity aboard a futuristic spacecraft. The monkey wears a sleek, advanced spacesuit with glowing blue accents, gazing through the spacecraft window at a mesmerizing nebula and distant planets. The spacecraft interior is filled with detailed control panels and holographic displays, creating an immersive, high-tech atmosphere." \
+    --num_frames 91 \
     --height 480 \
     --width 848 \
     --num_inference_steps 64 \
     --guidance_scale 4.5 \
-    --output_path outputs_video/debug \
+    --output_path outputs_video/\
     --seed 12345
 
 
 
 
-num_gpus=4
-torchrun --nnodes=1 --nproc_per_node=4 --master_port 29503 \
+
+
+
+
+
+
+num_gpus=2
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
     fastvideo/sample/sample_t2v_mochi.py \
     --model_path data/mochi/ \
     --prompt_embed_path data/Encoder_Overfit_Data/prompt_embed/0.pt \

pyproject.toml

@@ -13,15 +13,15 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License",
 ]
 dependencies = [
-"transformers==4.46.1", "accelerate==1.0.1", "tokenizers==0.20.1", "diffusers==0.31.0", "albumentations==1.4.20", "av==13.1.0", 
+"transformers==4.46.1", "accelerate==1.0.1", "tokenizers==0.20.1", "albumentations==1.4.20", "av==13.1.0", 
 "decord==0.6.0", "einops==0.8.0", "fastapi==0.115.3", "gdown==5.2.0", "h5py==3.12.1", "idna==3.6", "imageio==2.36.0", 
 "matplotlib==3.9.2", "numpy==1.26.3", "omegaconf==2.3.0", "opencv-python==4.10.0.84", "opencv-python-headless==4.10.0.84", 
 "pandas==2.2.3", "pillow==10.2.0", "pydub==0.25.1", "pytorch-lightning==2.4.0", "pytorchvideo==0.1.5", "PyYAML==6.0.1", 
 "regex==2024.9.11", "requests==2.31.0", "scikit-learn==1.5.2", "scipy==1.14.1", "six==1.16.0", "test-tube==0.7.5", 
 "timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0", 
 "scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0", 
-"moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", "peft==0.13.2"
-"watch", "gpustat"
+"moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3", 
+"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1"
 ]
 
 [project.optional-dependencies]

scripts/data_preprocess/finetune_data_gen.sh

@@ -4,7 +4,7 @@ MODEL_PATH="/ephemeral/hao.zhang/outputfolder/ckptfolder/mochi_diffuser"
 DATA_MERGE_PATH="/ephemeral/hao.zhang/resourcefolder/Mochi-Synthetic-Data-BW-Finetune/merge.txt"
 OUTPUT_DIR="./data/BW-Finetune-Synthetic-Data_test"
 
-torchrun --nproc_per_node=$GPU_NUM \
+rchrun --nproc_per_node=$GPU_NUM \
     ./fastvideo/utils/data_preprocess/finetune_data_VAE.py \
     --model_path $MODEL_PATH \
     --data_merge_path $DATA_MERGE_PATH \
@@ -15,6 +15,10 @@ torchrun --nproc_per_node=$GPU_NUM \
     --dataloader_num_workers 1 \
     --output_dir=$OUTPUT_DIR
 
+to
+
+
+
 torchrun --nproc_per_node=$GPU_NUM \
     ./fastvideo/utils/data_preprocess/finetune_data_T5.py \
     --model_path $MODEL_PATH \

scripts/finetune_T_J.sh

@@ -0,0 +1,117 @@
+torchrun --nnodes 4 --nproc_per_node 4 \
+    --node_rank=3 \
+    --rdzv_id=456 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint=172.23.30.15:29500 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --cache_dir "data/.cache" \
+    --data_json_path "data/BLACK-MYTH-Finetune-Dataset/videos2caption.json" \
+    --validation_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/validation_prompt_embed_mask" \
+    --uncond_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/uncond_prompt_embed_mask" \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 16 \
+    --sp_size 4 \
+    --train_sp_batch_size 4 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=4000 \
+    --learning_rate=2e-5 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=500 \
+    --validation_steps=100 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir="data/outputs/black_myth_correct_video_mask" \
+    --weighting_scheme "uniform" \
+    --num_frames 91 \
+    --selective_checkpointing 1.0
+
+
+torchrun --nnodes 4 --nproc_per_node 4 \
+    --node_rank=3 \
+    --rdzv_id=456 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint=172.23.30.15:29500 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --cache_dir "data/.cache" \
+    --data_json_path "data/BLACK-MYTH-DREAM-BOOTH/videos2caption.json" \
+    --validation_prompt_dir "data/BLACK-MYTH-DREAM-BOOTH/validation_prompt_embed_mask" \
+    --uncond_prompt_dir "data/BLACK-MYTH-DREAM-BOOTH/uncond_prompt_embed_mask" \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 16 \
+    --sp_size 4 \
+    --train_sp_batch_size 4 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=4000 \
+    --learning_rate=2e-5 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=100 \
+    --validation_steps=50 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.1 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir="data/outputs/black_myth_correct_video_mask" \
+    --weighting_scheme "uniform" \
+    --num_frames 91 \
+    --selective_checkpointing 1.0
+
+
+
+
+torchrun --nnodes 1 --nproc_per_node 4 \
+    fastvideo/train.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --cache_dir "data/.cache" \
+    --data_json_path "data/BLACK-MYTH-Finetune-Dataset/videos2caption.json" \
+    --validation_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/validation_prompt_embed_mask" \
+    --uncond_prompt_dir "data/BLACK-MYTH-Finetune-Dataset/uncond_prompt_embed_mask" \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 16 \
+    --sp_size 4 \
+    --train_sp_batch_size 4 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=2000 \
+    --learning_rate=2e-5 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=100 \
+    --validation_steps=100 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir="data/outputs/black_myth_dream" \
+    --weighting_scheme "uniform" \
+    --num_frames 91 \
+
+
+| 4/2000 [01:23<10:59:33, 19.83s/it, loss=0.74, grad_norm=2.6]
+
+
+| 2/2000 [00:44<11:59:27, 21.61s/it, loss=0.873, grad_norm=1.57]
+| 3/2000 [01:03<11:18:51, 20.40s/it, loss=0.813, grad_norm=1.48]
+
+
+# FSDP Hybrid: 
+

t2v_debug_multi.sh

@@ -1,3 +1,4 @@
+export WANDB_MODE=online
 torchrun --nnodes 1 --nproc_per_node 4 \
     fastvideo/train.py \
     --seed 42 \
@@ -5,72 +6,26 @@ torchrun --nnodes 1 --nproc_per_node 4 \
     --cache_dir "data/.cache" \
     --data_json_path "data/Encoder_Overfit_Data/videos2caption.json" \
     --validation_prompt_dir "data/Encoder_Overfit_Data/validation_prompt_embed_mask" \
-    --uncond_prompt_dir "data/Encoder_Overfit_Data/uncond_prompt_embed_mask" \
+    --uncond_prompt_dir "data/Encoder_Overfit_Dataa/uncond_prompt_embed_mask" \
     --gradient_checkpointing \
     --train_batch_size=1 \
-    --num_latent_t 2 \
-    --sp_size 2 \
-    --train_sp_batch_size 1 \
-    --dataloader_num_workers 1 \
-    --gradient_accumulation_steps 4 \
-    --max_train_steps 200 \
-    --learning_rate 1e-4 \
-    --mixed_precision="bf16" \
-    --checkpointing_steps 6 \
-    --validation_steps 40 \
-    --validation_sampling_steps 2 \
-    --checkpoints_total_limit 3 \
-    --allow_tf32 \
-    --ema_start_step 0 \
-    --cfg 0.1 \
-    --ema_decay 0.999 \
-    --log_validation \
-    --output_dir="data/outputs/BW_Testrun_Test" \
-    --use_lora \
-    --resume_from_lora_checkpoint data/outputs/BW_Testrun/lora-checkpoint-5 \
-
-torchrun --nnodes 1 --nproc_per_node 4 \
-    fastvideo/train.py \
-    --seed 42 \
-    --pretrained_model_name_or_path data/mochi \
-    --cache_dir data/.cache \
-    --data_json_path data/Encoder_Overfit_Data/videos2caption.json \
-    --validation_prompt_dir data/Encoder_Overfit_Data/validation_prompt_embed_mask \
-    --uncond_prompt_dir data/Encoder_Overfit_Data/uncond_prompt_embed_mask \
-    --gradient_checkpointing \
-    --train_batch_size 1 \
-    --num_latent_t 12 \
+    --num_latent_t 14 \
     --sp_size 2 \
     --train_sp_batch_size 2 \
-    --dataloader_num_workers 1 \
-    --gradient_accumulation_steps 1 \
-    --max_train_steps 400 \
-    --learning_rate 1e-4 \
-    --mixed_precision bf16 \
-    --checkpointing_steps 200 \
-    --validation_steps 50 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=2 \
+    --max_train_steps=2000 \
+    --learning_rate=1e-5 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=250 \
+    --validation_steps 250 \
     --validation_sampling_steps 64 \
     --checkpoints_total_limit 3 \
     --allow_tf32 \
     --ema_start_step 0 \
-    --cfg 0.1 \
+    --cfg 0.0 \
     --ema_decay 0.999 \
     --log_validation \
-    --output_dir data/outputs/BW_Testrun_Test/ \
-    --lora_rank 128 \
-    --lora_alpha 256 \
-    --use_lora \
-    --use_cpu_offload \
-    --resume_from_lora_checkpoint data/outputs/BW_Testrun_2/lora-checkpoint-100 \
-
-
-
-
-# SP single node: 13
-# SP multi node: 20
-# SP multi node with larger batch size: 15
-# No SP Single node: 10.5
-# No SP Multe node: 11
-
-
+    --output_dir="data/outputs/T_J_FT" \
+    --weighting_scheme "uniform"