updates

training/mochi-1/dataset_mochi.py

@@ -3,11 +3,11 @@ from typing import Any, Dict, Tuple
 
 import numpy as np
 import torch
-from torchvision import transforms
+import torch.nn as nn
 from accelerate.logging import get_logger
+from torchvision import transforms
 from torchvision.transforms import InterpolationMode
 from torchvision.transforms.functional import resize
-import torch.nn as nn
 
 
 # Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error
@@ -16,14 +16,17 @@ import decord  # isort:skip
 
 decord.bridge.set_bridge("torch")
 
-import sys 
+import sys
+
+
 sys.path.append("..")
 
 from dataset import VideoDataset as VDS
 
+
 logger = get_logger(__name__)
 
-# TODO (sayakpaul): probably not all buckets are needed for Mochi-1? 
+# TODO (sayakpaul): probably not all buckets are needed for Mochi-1?
 HEIGHT_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536]
 WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 848, 960, 1024, 1280, 1536]
 FRAME_BUCKETS = [16, 24, 32, 48, 64, 80, 85]
@@ -68,7 +71,7 @@ class VideoDataset(VDS):
             # This is hardcoded for now.
             # Output of the VAE encoding is 2 * output_channels and then it's
             # temporal compression factor is 6. Initially, the VAE encodings will have
-            # 24 latent number of frames. So, if we were to train with a 
+            # 24 latent number of frames. So, if we were to train with a
             # max frame size of 85 and frame bucket of [85], we need to have the following logic.
             latent_num_frames = video_latents.size(0)
             num_frames = ((latent_num_frames // 2) * (VAE_TEMPORAL_SCALE_FACTOR + 1) + 1)
@@ -169,13 +172,13 @@ class VideoDatasetWithFlexibleResize(VideoDataset):
             )
             frame_indices = list(
                 range(
-                    0, 
+                    0,
-                    video_num_frames, 
+                    video_num_frames,
                     1 if video_num_frames < nearest_frame_bucket else video_num_frames // nearest_frame_bucket
                 )
             )
             frames = video_reader.get_batch(frame_indices)
-            
+
             # Pad or truncate frames to match the bucket size
             if video_num_frames < nearest_frame_bucket:
                 pad_size = nearest_frame_bucket - video_num_frames
@@ -210,7 +213,7 @@ class VideoDatasetWithFlexibleResize(VideoDataset):
     def _resize_for_rectangle_crop(self, arr: torch.Tensor, image_size: Tuple[int, int]) -> torch.Tensor:
         """
         Resize frames for rectangular cropping.
-        
+
         Args:
             arr (torch.Tensor): The video frames tensor [N, C, H, W].
             image_size (Tuple[int, int]): The target resolution (height, width).
@@ -239,4 +242,4 @@ class VideoDatasetWithFlexibleResize(VideoDataset):
         else:
             raise NotImplementedError(f"Unsupported reshape mode: {self.video_reshape_mode}")
 
-        return transforms.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])
+        return transforms.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])

training/mochi-1/prepare_dataset.py

@@ -8,12 +8,13 @@ import pathlib
 import queue
 import traceback
 import uuid
-from contextlib import nullcontext
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import nullcontext
 from typing import Any, Dict, List, Optional, Union
 
 import torch
 import torch.distributed as dist
+from dataset_mochi import VideoDatasetWithFlexibleResize
 from diffusers import AutoencoderKLMochi
 from diffusers.training_utils import set_seed
 from diffusers.utils import export_to_video, get_logger
@@ -21,7 +22,6 @@ from torch.utils.data import DataLoader
 from torchvision import transforms
 from tqdm import tqdm
 from transformers import T5EncoderModel, T5Tokenizer
-from dataset_mochi import VideoDatasetWithFlexibleResize
 
 
 import decord  # isort:skip
@@ -29,6 +29,8 @@ import decord  # isort:skip
 decord.bridge.set_bridge("torch")
 
 import sys
+
+
 sys.path.append("..")
 
 from dataset import BucketSampler
@@ -489,11 +491,11 @@ def main():
         text_encoder = T5EncoderModel.from_pretrained(
             args.model_id, subfolder="text_encoder", torch_dtype=weight_dtype
         ).to(device)
-        
+
         vae = AutoencoderKLMochi.from_pretrained(
             args.model_id, subfolder="vae", torch_dtype=weight_dtype
         ).to(device)
-        
+
         if args.use_slicing:
             vae.enable_slicing()
         if args.use_tiling:

training/mochi-1/text_to_video_lora.py

@@ -13,7 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import gc
+import json
 import logging
 import math
 import os
@@ -26,7 +28,6 @@ import diffusers
 import torch
 import transformers
 import wandb
-import copy
 from accelerate import Accelerator, DistributedType
 from accelerate.logging import get_logger
 from accelerate.utils import (
@@ -43,7 +44,11 @@ from diffusers import (
 )
 from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.optimization import get_scheduler
-from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
+from diffusers.training_utils import (
+    cast_training_params,
+    compute_density_for_timestep_sampling,
+    compute_loss_weighting_for_sd3,
+)
 from diffusers.utils import convert_unet_state_dict_to_peft, export_to_video
 from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
 from diffusers.utils.torch_utils import is_compiled_module
@@ -53,10 +58,13 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
+
 from args import get_args  # isort:skip
 from dataset_mochi import VideoDatasetWithFlexibleResize  # isort:skip
 
 import sys
+
+
 sys.path.append("..")
 
 from dataset import BucketSampler # isort:skip
@@ -78,7 +86,7 @@ def save_model_card(
     widget_dict = []
     if videos is not None and len(videos) > 0:
         for i, video in enumerate(videos):
-            export_to_video(video, os.path.join(repo_folder, f"final_video_{i}.mp4", fps=fps))
+            export_to_video(video, os.path.join(repo_folder, f"final_video_{i}.mp4"), fps=fps)
             widget_dict.append(
                 {
                     "text": validation_prompt if validation_prompt else " ",
@@ -147,7 +155,8 @@ def log_validation(
         f"Running validation... \n Generating {args.num_validation_videos} videos with prompt: {pipeline_args['prompt']}."
     )
 
-    pipe = pipe.to(accelerator.device)
+    if not args.enable_model_cpu_offload:
+        pipe = pipe.to(accelerator.device)
 
     # run inference
     generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None
@@ -261,7 +270,7 @@ def main(args):
         set_seed(args.seed)
 
     # Handle the repository creation
-    if accelerator.is_main_process:
+    if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
         if args.output_dir is not None:
             os.makedirs(args.output_dir, exist_ok=True)
 
@@ -283,7 +292,7 @@ def main(args):
             subfolder="text_encoder",
             revision=args.revision,
         )
-        
+
         vae = AutoencoderKLMochi.from_pretrained(
             args.pretrained_model_name_or_path,
             subfolder="vae",
@@ -295,10 +304,11 @@ def main(args):
         if args.enable_tiling:
             vae.enable_tiling()
 
+        # keep things in FP32.
         text_encoder.requires_grad_(False)
-        text_encoder.to(accelerator.device, dtype=weight_dtype)
+        text_encoder.to(accelerator.device, dtype=torch.float32)
         vae.requires_grad_(False)
-        vae.to(accelerator.device, dtype=weight_dtype)
+        vae.to(accelerator.device, dtype=torch.float32)
 
     load_dtype = torch.bfloat16 if "5b" in args.pretrained_model_name_or_path.lower() else torch.float16
     transformer = MochiTransformer3DModel.from_pretrained(
@@ -315,7 +325,7 @@ def main(args):
     vae_in_channels = vae_config["latent_channels"]
     has_latents_mean = "latents_mean" in vae_config and vae_config["latents_mean"] is not None
     has_latents_std = "latents_std" in vae_config and vae_config["latents_std"] is not None
-    
+
     VAE_SCALING_FACTOR = vae_config["scaling_factor"]
 
     # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision
@@ -344,10 +354,10 @@ def main(args):
         raise ValueError(
             "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
         )
-    
+
+    # keep the transformer in FP32.
     transformer.requires_grad_(False)
-    # transformer.to(accelerator.device, dtype=weight_dtype)
+    transformer.to(accelerator.device, torch.float32)
-    transformer.to(accelerator.device)
     if args.gradient_checkpointing:
         transformer.enable_gradient_checkpointing()
 
@@ -355,7 +365,7 @@ def main(args):
     transformer_lora_config = LoraConfig(
         r=args.rank,
         lora_alpha=args.lora_alpha,
-        init_lora_weights=True,
+        init_lora_weights="gaussian",
         target_modules=args.target_modules,
     )
     transformer.add_adapter(transformer_lora_config)
@@ -482,6 +492,18 @@ def main(args):
     accelerator.print(f"Using {optimizer.__class__.__name__} optimizer.")
 
     # Dataset and DataLoader
+    if args.load_tensors and args.id_token:
+        with open(os.path.join(args.data_root, "data.jsonl")) as f:
+            contents = [json.loads(jline) for jline in f.read().splitlines()]
+            parsed_id_token = None
+            for content in contents:
+                if "id_token" in content:
+                    parsed_id_token = content["id_token"]
+            if parsed_id_token is not None and parsed_id_token.strip() != args.id_token.strip():
+                raise ValueError(
+                    f"Parsed `id_token` from serialized metadata is {parsed_id_token} and provided `id_token` is {args.id_token}. They should match."
+                )
+
     dataset_init_kwargs = {
         "data_root": args.data_root,
         "video_reshape_mode": args.video_reshape_mode,
@@ -497,7 +519,8 @@ def main(args):
         "random_flip": args.random_flip,
     }
     train_dataset = VideoDatasetWithFlexibleResize(**dataset_init_kwargs)
-    collate_fn = CollateFunction(weight_dtype, args.load_tensors)
+    # keeping things in FP32 for now.
+    collate_fn = CollateFunction(weight_dtype=torch.float32, load_tensors=args.load_tensors)
     train_dataloader = DataLoader(
         train_dataset,
         batch_size=1,
@@ -619,7 +642,6 @@ def main(args):
     if args.load_tensors:
         gc.collect()
         torch.cuda.empty_cache()
-        torch.cuda.synchronize(accelerator.device)
 
     def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
         sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
@@ -650,11 +672,11 @@ def main(args):
                 # Encode videos
                 if not args.load_tensors:
                     videos = videos.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
-                    latent_dist = vae.encode(videos).latent_dist
+                    latent_dist = vae.encode(videos.to(vae.dtype)).latent_dist
                 else:
                     latent_dist = DiagonalGaussianDistribution(videos)
 
-                videos = latent_dist.sample() 
+                videos = latent_dist.sample()
                 if has_latents_mean and has_latents_std:
                     latents_mean = (
                         torch.tensor(vae_config["latents_mean"]).view(1, vae_in_channels, 1, 1, 1).to(videos.device, videos.dtype)
@@ -665,8 +687,9 @@ def main(args):
                     videos = (videos - latents_mean) * VAE_SCALING_FACTOR / latents_std
                 else:
                     videos = videos * VAE_SCALING_FACTOR
-                
+
-                videos = videos.to(memory_format=torch.contiguous_format, dtype=weight_dtype)
+                # keep in FP32 for now.
+                videos = videos.to(memory_format=torch.contiguous_format, dtype=torch.float32)
                 model_input = videos
 
                 # Encode prompts
@@ -677,11 +700,11 @@ def main(args):
                         prompts,
                         model_config.max_text_seq_length,
                         accelerator.device,
-                        weight_dtype,
+                        weight_dtype=weight_dtype,
                         requires_grad=False,
                     )
                 else:
-                    prompt_embeds = prompts.to(dtype=weight_dtype)
+                    prompt_embeds = prompts.to(weight_dtype)
                     prompt_attention_mask = prompt_attention_mask.to(accelerator.device)
 
                 # Sample noise that will be added to the latents
@@ -697,15 +720,25 @@ def main(args):
                     logit_std=args.logit_std,
                     mode_scale=args.mode_scale,
                 )
-                indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
+                # indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
-                timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)
+                # timesteps = noise_scheduler_copy.timesteps[indices].to(device=model_input.device)
+                # revisit.
+                timesteps = (u * noise_scheduler_copy.config.num_train_timesteps)
 
                 # Add noise according to flow matching.
                 # zt = (1 - texp) * x + texp * z1
-                sigmas = get_sigmas(timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+                sigmas = get_sigmas(
-                noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise
+                    timesteps=noise_scheduler_copy.timesteps[timesteps.long()].to(device=model_input.device),
+                    n_dim=model_input.ndim,
+                    dtype=model_input.dtype
+                )
+                noisy_model_input = (1.0 - sigmas) * model_input + sigmas * noise # do we need to revisit this?
+                noisy_model_input = noisy_model_input.to(weight_dtype)
 
                 # Predict the noise residual
+                actual_num_train_timesteps = float(noise_scheduler_copy.config.num_train_timesteps)
+                timesteps = (1 - (timesteps / actual_num_train_timesteps)) * actual_num_train_timesteps # revisit
+                timesteps = timesteps.to(device=model_input.device)
                 model_pred = transformer(
                     hidden_states=noisy_model_input,
                     encoder_hidden_states=prompt_embeds,
@@ -719,7 +752,8 @@ def main(args):
                 weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
 
                 # flow matching loss
-                target = noise - model_input
+                # target = noise - model_input
+                target = model_input - noise # as discussed with Ajay
 
                 loss = torch.mean(
                     (weighting * (model_pred.float() - target.float()) ** 2).reshape(batch_size, -1),
@@ -787,11 +821,13 @@ def main(args):
             if global_step >= args.max_train_steps:
                 break
 
-        if accelerator.is_main_process:
+        if global_step >= args.max_train_steps:
+                break
+
+        if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
             if args.validation_prompt is not None and (epoch + 1) % args.validation_epochs == 0:
                 accelerator.print("===== Memory before validation =====")
                 print_memory(accelerator.device)
-                torch.cuda.synchronize(accelerator.device)
 
                 pipe = MochiPipeline.from_pretrained(
                     args.pretrained_model_name_or_path,
@@ -799,7 +835,6 @@ def main(args):
                     scheduler=scheduler,
                     revision=args.revision,
                     variant=args.variant,
-                    # torch_dtype=weight_dtype,
                 )
 
                 if args.enable_slicing:
@@ -831,25 +866,18 @@ def main(args):
                 print_memory(accelerator.device)
                 reset_memory(accelerator.device)
 
+
+                del pipe.text_encoder
+                del pipe.vae
                 del pipe
                 gc.collect()
                 torch.cuda.empty_cache()
-                torch.cuda.synchronize(accelerator.device)
 
     accelerator.wait_for_everyone()
 
-    if accelerator.is_main_process:
+    if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
         transformer = unwrap_model(transformer)
-        dtype = (
-            torch.float16
-            if args.mixed_precision == "fp16"
-            else torch.bfloat16
-            if args.mixed_precision == "bf16"
-            else torch.float32
-        )
-        # transformer = transformer.to(dtype)
         transformer_lora_layers = get_peft_model_state_dict(transformer)
-
         MochiPipeline.save_lora_weights(
             save_directory=args.output_dir,
             transformer_lora_layers=transformer_lora_layers,
@@ -863,7 +891,6 @@ def main(args):
 
         gc.collect()
         torch.cuda.empty_cache()
-        torch.cuda.synchronize(accelerator.device)
 
         # Final test inference
         validation_outputs = []
@@ -871,14 +898,12 @@ def main(args):
             accelerator.print("===== Memory before testing =====")
             print_memory(accelerator.device)
             reset_memory(accelerator.device)
-            
+
             pipe = MochiPipeline.from_pretrained(
                 args.pretrained_model_name_or_path,
                 revision=args.revision,
                 variant=args.variant,
-                # torch_dtype=weight_dtype,
             )
-            pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
 
             if args.enable_slicing:
                 pipe.vae.enable_slicing()
@@ -898,7 +923,6 @@ def main(args):
                 pipeline_args = {
                     "prompt": validation_prompt,
                     "guidance_scale": 4.5,
-                    "use_dynamic_cfg": args.use_dynamic_cfg,
                     "height": args.height,
                     "width": args.width,
                 }