training/train_controlnet_backup.py

# Copyright 2024 The HuggingFace Team.
# All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import random
import signal
import sys
import threading
import time

import cv2
import yaml

sys.path.append('..')
import argparse
from PIL import Image
import logging
import math
import os
import shutil
from pathlib import Path
from typing import List, Optional, Tuple, Union

import torch
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
from huggingface_hub import create_repo, upload_folder
from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from tqdm.auto import tqdm
import numpy as np
from decord import VideoReader
from transformers import AutoTokenizer, T5EncoderModel, T5Tokenizer

import diffusers
from diffusers import AutoencoderKLCogVideoX, CogVideoXDPMScheduler
from diffusers.models.embeddings import get_3d_rotary_pos_embed
from diffusers.optimization import get_scheduler
from diffusers.pipelines.cogvideo.pipeline_cogvideox import get_resize_crop_region_for_grid
from diffusers.training_utils import (
    cast_training_params,
    free_memory,
)
from diffusers.utils import check_min_version, export_to_video, is_wandb_available
from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
from diffusers.utils.torch_utils import is_compiled_module

from controlnet_datasets import AblationFullMotionBlurDataset, AdobeMotionBlurDataset, FullMotionBlurDataset, GoPro2xMotionBlurDataset, GoProLargeMotionBlurDataset, OutsidePhotosDataset, GoProMotionBlurDataset, BAISTDataset, SimpleBAISTDataset
from controlnet_pipeline import ControlnetCogVideoXPipeline
from cogvideo_transformer import CogVideoXTransformer3DModel
from helpers import random_insert_latent_frame, transform_intervals
import os
import tempfile
from atomicwrites import atomic_write



if is_wandb_available():
    import wandb

# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
check_min_version("0.31.0.dev0")

logger = get_logger(__name__)

def save_frames_as_pngs(video_array,output_dir, 
                        downsample_spatial=1,   # e.g. 2 to halve width & height
                        downsample_temporal=1): # e.g. 2 to keep every 2nd frame
    """
    Save each frame of a (T, H, W, C) numpy array as a PNG with no compression.
    """
    assert video_array.ndim == 4 and video_array.shape[-1] == 3, \
        "Expected (T, H, W, C=3) array"
    assert video_array.dtype == np.uint8, "Expected uint8 array"
    
    os.makedirs(output_dir, exist_ok=True)
    
    # temporal downsample
    frames = video_array[::downsample_temporal]
    
    # compute spatially downsampled size
    T, H, W, _ = frames.shape
    new_size = (W // downsample_spatial, H // downsample_spatial)
    
    # PNG compression param: 0 = no compression
    png_params = [cv2.IMWRITE_PNG_COMPRESSION, 0]
    
    for idx, frame in enumerate(frames):
        # frame is RGB; convert to BGR for OpenCV
        bgr = frame[..., ::-1]
        if downsample_spatial > 1:
            bgr = cv2.resize(bgr, new_size, interpolation=cv2.INTER_NEAREST)
        
        filename = os.path.join(output_dir, "frame_{:05d}.png".format(idx))
        success = cv2.imwrite(filename, bgr, png_params)
        if not success:
            raise RuntimeError("Failed to write frame ")


def get_args():
    parser = argparse.ArgumentParser(description="Training script for CogVideoX using config file.")
    parser.add_argument(
        "--config", 
        type=str, 
        required=True, 
        help="Path to the YAML config file."
    )
    args = parser.parse_args()

    with open(args.config, "r") as f:
        config = yaml.safe_load(f)
    
    args = argparse.Namespace(**config)

    # Convert nested config dict to an argparse.Namespace for easier downstream usage
    return args


# def read_video(video_path, start_index=0, frames_count=49, stride=1):
#     video_reader = VideoReader(video_path)
#     end_index = min(start_index + frames_count * stride, len(video_reader)) - 1
#     batch_index = np.linspace(start_index, end_index, frames_count, dtype=int)
#     numpy_video = video_reader.get_batch(batch_index).asnumpy()
#     return numpy_video
    

def log_validation(
    pipe,
    args,
    accelerator,
    pipeline_args,
    epoch,
    is_final_validation: bool = False,
):
    logger.info(
        f"Running validation... \n Generating {args.num_validation_videos} videos with prompt: {pipeline_args['prompt']}."
    )
    # We train on the simplified learning objective. If we were previously predicting a variance, we need the scheduler to ignore it
    scheduler_args = {}

    if "variance_type" in pipe.scheduler.config:
        variance_type = pipe.scheduler.config.variance_type

        if variance_type in ["learned", "learned_range"]:
            variance_type = "fixed_small"

        scheduler_args["variance_type"] = variance_type

    pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args)
    pipe = pipe.to(accelerator.device)
    # pipe.set_progress_bar_config(disable=True)

    # run inference
    generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None

    videos = []
    for _ in range(args.num_validation_videos):
        video = pipe(**pipeline_args, generator=generator, output_type="np").frames[0]
        videos.append(video)

    free_memory()

    return videos


def _get_t5_prompt_embeds(
    tokenizer: T5Tokenizer,
    text_encoder: T5EncoderModel,
    prompt: Union[str, List[str]],
    num_videos_per_prompt: int = 1,
    max_sequence_length: int = 226,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    text_input_ids=None,
):
    prompt = [prompt] if isinstance(prompt, str) else prompt
    batch_size = len(prompt)

    if tokenizer is not None:
        text_inputs = tokenizer(
            prompt,
            padding="max_length",
            max_length=max_sequence_length,
            truncation=True,
            add_special_tokens=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids
    else:
        if text_input_ids is None:
            raise ValueError("`text_input_ids` must be provided when the tokenizer is not specified.")

    prompt_embeds = text_encoder(text_input_ids.to(device))[0]
    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)

    # duplicate text embeddings for each generation per prompt, using mps friendly method
    _, seq_len, _ = prompt_embeds.shape
    prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
    prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)

    return prompt_embeds


def encode_prompt(
    tokenizer: T5Tokenizer,
    text_encoder: T5EncoderModel,
    prompt: Union[str, List[str]],
    num_videos_per_prompt: int = 1,
    max_sequence_length: int = 226,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    text_input_ids=None,
):
    prompt = [prompt] if isinstance(prompt, str) else prompt
    prompt_embeds = _get_t5_prompt_embeds(
        tokenizer,
        text_encoder,
        prompt=prompt,
        num_videos_per_prompt=num_videos_per_prompt,
        max_sequence_length=max_sequence_length,
        device=device,
        dtype=dtype,
        text_input_ids=text_input_ids,
    )
    return prompt_embeds


def compute_prompt_embeddings(
    tokenizer, text_encoder, prompt, max_sequence_length, device, dtype, requires_grad: bool = False
):
    if requires_grad:
        prompt_embeds = encode_prompt(
            tokenizer,
            text_encoder,
            prompt,
            num_videos_per_prompt=1,
            max_sequence_length=max_sequence_length,
            device=device,
            dtype=dtype,
        )
    else:
        with torch.no_grad():
            prompt_embeds = encode_prompt(
                tokenizer,
                text_encoder,
                prompt,
                num_videos_per_prompt=1,
                max_sequence_length=max_sequence_length,
                device=device,
                dtype=dtype,
            )
    return prompt_embeds


def prepare_rotary_positional_embeddings(
    height: int,
    width: int,
    num_frames: int,
    vae_scale_factor_spatial: int = 8,
    patch_size: int = 2,
    attention_head_dim: int = 64,
    device: Optional[torch.device] = None,
    base_height: int = 480,
    base_width: int = 720,
) -> Tuple[torch.Tensor, torch.Tensor]:
    grid_height = height // (vae_scale_factor_spatial * patch_size)
    grid_width = width // (vae_scale_factor_spatial * patch_size)
    base_size_width = base_width // (vae_scale_factor_spatial * patch_size)
    base_size_height = base_height // (vae_scale_factor_spatial * patch_size)

    grid_crops_coords = get_resize_crop_region_for_grid((grid_height, grid_width), base_size_width, base_size_height)
    freqs_cos, freqs_sin = get_3d_rotary_pos_embed(
        embed_dim=attention_head_dim,
        crops_coords=grid_crops_coords,
        grid_size=(grid_height, grid_width),
        temporal_size=num_frames,
    )

    freqs_cos = freqs_cos.to(device=device)
    freqs_sin = freqs_sin.to(device=device)
    return freqs_cos, freqs_sin


def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
    # Use DeepSpeed optimzer
    if use_deepspeed:
        from accelerate.utils import DummyOptim


        return DummyOptim(
            params_to_optimize,
            lr=args.learning_rate,
            betas=(args.adam_beta1, args.adam_beta2),
            eps=args.adam_epsilon,
            weight_decay=args.adam_weight_decay,
        )

    # Optimizer creation
    supported_optimizers = ["adam", "adamw", "prodigy"]
    if args.optimizer not in supported_optimizers:
        logger.warning(
            f"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW"
        )
        args.optimizer = "adamw"

    if args.use_8bit_adam and not (args.optimizer.lower() not in ["adam", "adamw"]):
        logger.warning(
            f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was "
            f"set to {args.optimizer.lower()}"
        )

    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`."
            )

    if args.optimizer.lower() == "adamw":
        optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW

        optimizer = optimizer_class(
            params_to_optimize,
            betas=(args.adam_beta1, args.adam_beta2),
            eps=args.adam_epsilon,
            weight_decay=args.adam_weight_decay,
        )
    elif args.optimizer.lower() == "adam":
        optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam


        optimizer = optimizer_class(
            params_to_optimize,
            betas=(args.adam_beta1, args.adam_beta2),
            eps=args.adam_epsilon,
            weight_decay=args.adam_weight_decay,
        )
    elif 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,
        )

    return optimizer


def main(args):
    global signal_recieved_time
    if args.report_to == "wandb" and args.hub_token is not None:
        raise ValueError(
            "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token."
            " Please use `huggingface-cli login` to authenticate with the Hub."
        )

    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )

    logging_dir = Path(args.output_dir, args.logging_dir)

    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
    kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
        kwargs_handlers=[kwargs],
    )

    # Disable AMP for MPS.
    if torch.backends.mps.is_available():
        accelerator.native_amp = False

    if args.report_to == "wandb":
        if not is_wandb_available():
            raise ImportError("Make sure to install wandb if you want to use it for logging during training.")

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)

        if args.push_to_hub:
            repo_id = create_repo(
                repo_id=args.hub_model_id or Path(args.output_dir).name,
                exist_ok=True,
            ).repo_id

    # Prepare models and scheduler
    tokenizer = AutoTokenizer.from_pretrained(
        os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="tokenizer", revision=args.revision
    )

    text_encoder = T5EncoderModel.from_pretrained(
        os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="text_encoder", revision=args.revision
    )

    # CogVideoX-2b weights are stored in float16
    # CogVideoX-5b and CogVideoX-5b-I2V weights are stored in bfloat16

    ## TRYING NEW CONFIG LOADING
    config = CogVideoXTransformer3DModel.load_config(
    os.path.join(args.base_dir, args.pretrained_model_name_or_path),
    subfolder="transformer",
    revision=args.revision,
    variant=args.variant,
    )
    config["ablation_mode"] = args.ablation_mode if hasattr(args, "ablation_mode") else None 

     ##FINISH TRYING NEW CONFIG LOADING



    load_dtype = torch.bfloat16 if "5b" in os.path.join(args.base_dir, args.pretrained_model_name_or_path).lower() else torch.float16
    transformer = CogVideoXTransformer3DModel.from_pretrained(
        os.path.join(args.base_dir, args.pretrained_model_name_or_path),
        subfolder="transformer",
        torch_dtype=load_dtype,
        ablation_mode=args.ablation_mode if hasattr(args, "ablation_mode") else None,
        revision=args.revision,
        variant=args.variant,
        low_cpu_mem_usage=False,
    )

    vae = AutoencoderKLCogVideoX.from_pretrained(
        os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="vae", revision=args.revision, variant=args.variant
    )




    scheduler = CogVideoXDPMScheduler.from_pretrained(os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="scheduler")

    if args.enable_slicing:
        vae.enable_slicing()
    if args.enable_tiling:
        vae.enable_tiling()

    # We only train the additional adapter controlnet layers
    text_encoder.requires_grad_(False)
    transformer.requires_grad_(True)
    vae.requires_grad_(False)

    # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) 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.state.deepspeed_plugin:
        # DeepSpeed is handling precision, use what's in the DeepSpeed config
        if (
            "fp16" in accelerator.state.deepspeed_plugin.deepspeed_config
            and accelerator.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"]
        ):
            weight_dtype = torch.float16
        if (
            "bf16" in accelerator.state.deepspeed_plugin.deepspeed_config
            and accelerator.state.deepspeed_plugin.deepspeed_config["bf16"]["enabled"]
        ):
            weight_dtype = torch.float16
    else:
        if accelerator.mixed_precision == "fp16":
            weight_dtype = torch.float16
        elif accelerator.mixed_precision == "bf16":
            weight_dtype = torch.bfloat16

    if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16:
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )

    text_encoder.to(accelerator.device, dtype=weight_dtype)
    transformer.to(accelerator.device, dtype=weight_dtype)
    vae.to(accelerator.device, dtype=weight_dtype)

    if args.gradient_checkpointing:
        transformer.enable_gradient_checkpointing()

    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    # 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 and torch.cuda.is_available():
        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
        )

    # Make sure the trainable params are in float32.
    if args.mixed_precision == "fp16":
        # only upcast trainable parameters into fp32
        cast_training_params([transformer], dtype=torch.float32)

    trainable_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters()))

    # Optimization parameters
    trainable_parameters_with_lr = {"params": trainable_parameters, "lr": args.learning_rate}
    params_to_optimize = [trainable_parameters_with_lr]

    use_deepspeed_optimizer = (
        accelerator.state.deepspeed_plugin is not None
        and "optimizer" in accelerator.state.deepspeed_plugin.deepspeed_config
    )
    use_deepspeed_scheduler = (
        accelerator.state.deepspeed_plugin is not None
        and "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    )

    optimizer = get_optimizer(args, params_to_optimize, use_deepspeed=use_deepspeed_optimizer)

    # Dataset and DataLoader
    if args.dataset == "adobe":
        train_dataset = AdobeMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "gopro":
        train_dataset = GoProMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "gopro2x":
        train_dataset = GoPro2xMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "goprolarge":
        train_dataset = GoProLargeMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "full":
        train_dataset = FullMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "fullablation":
        train_dataset = AblationFullMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
            ablation_mode = args.ablation_mode, #this is not called for now
        )
    elif args.dataset == "baist":
        train_dataset = BAISTDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        ) #this is not called for now
    elif args.dataset == "simplebaist":
        train_dataset = SimpleBAISTDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = "train",
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
        

    if args.dataset == "adobe":
        val_dataset = AdobeMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "outsidephotos":
        
        val_dataset = OutsidePhotosDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
        train_dataset = val_dataset #dummy dataset
    elif args.dataset == "gopro":
        val_dataset = GoProMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "gopro2x":
        val_dataset = GoPro2xMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "goprolarge":
        val_dataset = GoProLargeMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "full":
        val_dataset = FullMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "fullablation":
        val_dataset = AblationFullMotionBlurDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
            ablation_mode = args.ablation_mode, #this is not called for now
        )
    elif args.dataset == "baist":
        val_dataset = BAISTDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
    elif args.dataset == "simplebaist":
        val_dataset = SimpleBAISTDataset(
            data_dir=os.path.join(args.base_dir, args.video_root_dir),
            split = args.val_split,
            image_size=(args.height, args.width), 
            stride=(args.stride_min, args.stride_max),
            sample_n_frames=args.max_num_frames,
            hflip_p=args.hflip_p,
        )
        
    def encode_video(video):
        video = video.to(accelerator.device, dtype=vae.dtype)
        video = video.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
        latent_dist = vae.encode(video).latent_dist.sample() * vae.config.scaling_factor
        return latent_dist.permute(0, 2, 1, 3, 4).to(memory_format=torch.contiguous_format)
    
    # def atomic_save(save_path, accelerator):

    #     dir_name = os.path.dirname(save_path)
    #     with tempfile.NamedTemporaryFile(delete=False, dir=dir_name) as tmp_file:
    #         tmp_path = tmp_file.name
    #         # Close the file so that it can be moved later
    #     #delete anything at the tmp_path
    #     if accelerator.is_main_process:
    #         accelerator.save_state(tmp_path) #just a backup incase things go crazy
    #         accelerator.save_state(save_path)
    #     os.remove(tmp_path)  
    #     accelerator.wait_for_everyone()



    def atomic_save(save_path, accelerator):
        parent = os.path.dirname(save_path)
        tmp_dir = tempfile.mkdtemp(dir=parent)
        backup_dir = save_path + "_backup"

        try:
            # Save state into the temp directory
            accelerator.save_state(tmp_dir)

            # Backup existing save_path if it exists
            if os.path.exists(save_path):
                os.rename(save_path, backup_dir)

            # Atomically move temp directory into place
            os.rename(tmp_dir, save_path)

            # Clean up the backup directory
            if os.path.exists(backup_dir):
                shutil.rmtree(backup_dir)

        except Exception as e:
            # Clean up temp directory on failure
            if os.path.exists(tmp_dir):
                shutil.rmtree(tmp_dir)

            # Restore from backup if replacement failed
            if os.path.exists(backup_dir):
                if os.path.exists(save_path):
                    shutil.rmtree(save_path)
                os.rename(backup_dir, save_path)

            raise e

        

    def collate_fn(examples):
        blur_img = [example["blur_img"] for example in examples]
        videos = [example["video"] for example in examples]
        if "high_fps_video" in examples[0]:
            high_fps_videos = [example["high_fps_video"] for example in examples]
            high_fps_videos = torch.stack(high_fps_videos)
            high_fps_videos = high_fps_videos.to(memory_format=torch.contiguous_format).float()
        if "bbx" in examples[0]:
            bbx = [example["bbx"] for example in examples]
            bbx = torch.stack(bbx)
            bbx = bbx.to(memory_format=torch.contiguous_format).float()
        prompts = [example["caption"] for example in examples]
        file_names = [example["file_name"] for example in examples]
        num_frames = [example["num_frames"] for example in examples]
        # if full_file_names in examples[0]:
        #     full_file_names = [example["full_file_name"] for example in examples]
        input_intervals = [example["input_interval"] for example in examples]
        output_intervals = [example["output_interval"] for example in examples]
        ablation_condition = [example["ablation_condition"] for example in examples] if "ablation_condition" in examples[0] else None


        videos = torch.stack(videos)
        videos = videos.to(memory_format=torch.contiguous_format).float()

        blur_img = torch.stack(blur_img)
        blur_img = blur_img.to(memory_format=torch.contiguous_format).float()


        input_intervals = torch.stack(input_intervals)
        if args.dataset == "gopro":
            input_intervals = input_intervals.to(memory_format=torch.contiguous_format).long()  #this is a bug, but I trained it like this on GOPRO (sets intervals all to 0), model doesn't need intervals for this dataset cause its always 7 frames in the same spacing
        else:
            input_intervals = input_intervals.to(memory_format=torch.contiguous_format).float()

        output_intervals = torch.stack(output_intervals)
        if args.dataset == "gopro":
            output_intervals = output_intervals.to(memory_format=torch.contiguous_format).long()  #this is a bug, but I trained it like this on GOPRO (sets intervals all to 0), model doesn't need intervals for this dataset cause its always 7 frames in the same spacing
        else:
            output_intervals = output_intervals.to(memory_format=torch.contiguous_format).float()

        #just used for ablation studies
        ablation_condition = torch.stack(ablation_condition) if ablation_condition is not None else None
        if ablation_condition is not None:
            ablation_condition = ablation_condition.to(memory_format=torch.contiguous_format).float()

        out_dict = {
            "file_names": file_names,
            "blur_img": blur_img,
            "videos": videos,
            "num_frames": num_frames,
            "prompts": prompts,
            "input_intervals": input_intervals,
            "output_intervals": output_intervals,
        }

        if "high_fps_video" in examples[0]:
            out_dict["high_fps_video"] = high_fps_videos
        if "bbx" in examples[0]:
            out_dict["bbx"] = bbx
        if ablation_condition is not None:
            out_dict["ablation_condition"] = ablation_condition
        return out_dict

    train_dataloader = DataLoader(
        train_dataset,
        batch_size=args.train_batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=args.dataloader_num_workers,
    )

    val_dataloader = DataLoader(
        val_dataset,
        batch_size=1,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=args.dataloader_num_workers,
    )

    # 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

    if use_deepspeed_scheduler:
        from accelerate.utils import DummyScheduler

        lr_scheduler = DummyScheduler(
            name=args.lr_scheduler,
            optimizer=optimizer,
            total_num_steps=args.max_train_steps * accelerator.num_processes,
            num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        )
    else:
        lr_scheduler = get_scheduler(
            args.lr_scheduler,
            optimizer=optimizer,
            num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
            num_training_steps=args.max_train_steps * accelerator.num_processes,
            num_cycles=args.lr_num_cycles,
            power=args.lr_power,
        )

    # Prepare everything with our `accelerator`.
    transformer, optimizer, train_dataloader, lr_scheduler, val_dataloader = accelerator.prepare(
        transformer, optimizer, train_dataloader, lr_scheduler, val_dataloader
    )

    # 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:
        tracker_name = args.tracker_name or "cogvideox-controlnet"
        accelerator.init_trackers(tracker_name, config=vars(args))

    
    accelerator.register_for_checkpointing(transformer, optimizer, lr_scheduler)
    save_path = os.path.join(args.output_dir, f"checkpoint")

    #check if the checkpoint already exists
    if os.path.exists(save_path):
        accelerator.load_state(save_path)
        logger.info(f"Loaded state from {save_path}")



    # Train!
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
    num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"])

    logger.info("***** Running training *****")
    logger.info(f"  Num trainable parameters = {num_trainable_parameters}")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
    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}")
    global_step = 0
    first_epoch = 0
    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,
    )
    vae_scale_factor_spatial = 2 ** (len(vae.config.block_out_channels) - 1)

    # For DeepSpeed training
    model_config = transformer.module.config if hasattr(transformer, "module") else transformer.config

    for epoch in range(first_epoch, args.num_train_epochs):
        transformer.train()
        for step, batch in enumerate(train_dataloader):
            if not args.just_validate:
                models_to_accumulate = [transformer]
                with accelerator.accumulate(models_to_accumulate):
                    model_input = encode_video(batch["videos"]).to(dtype=weight_dtype)  # [B, F, C, H, W]
                    prompts = batch["prompts"]
                    image_latent = encode_video(batch["blur_img"]).to(dtype=weight_dtype)  # [B, F, C, H, W]
                    input_intervals = batch["input_intervals"]
                    output_intervals = batch["output_intervals"] 
                    ablation_condition = batch["ablation_condition"] if "ablation_condition" in batch else None

                    batch_size = len(prompts)
                    # True = use real prompt (conditional); False = drop to empty (unconditional)
                    guidance_mask = torch.rand(batch_size, device=accelerator.device) >= 0.2  

                    # build a new prompts list: keep the original where mask True, else blank
                    per_sample_prompts = [
                        prompts[i] if guidance_mask[i] else ""
                        for i in range(batch_size)
                    ]
                    prompts = per_sample_prompts

                    # encode prompts
                    prompt_embeds = compute_prompt_embeddings(
                        tokenizer,
                        text_encoder,
                        prompts,
                        model_config.max_text_seq_length,
                        accelerator.device,
                        weight_dtype,
                        requires_grad=False,
                    )

                    # Sample noise that will be added to the latents
                    noise = torch.randn_like(model_input)
                    batch_size, num_frames, num_channels, height, width = model_input.shape

                    # Sample a random timestep for each image
                    timesteps = torch.randint(
                        0, scheduler.config.num_train_timesteps, (batch_size,), device=model_input.device
                    )
                    timesteps = timesteps.long()
            
                    # Prepare rotary embeds
                    image_rotary_emb = (
                        prepare_rotary_positional_embeddings(
                            height=args.height,
                            width=args.width,
                            num_frames=num_frames,
                            vae_scale_factor_spatial=vae_scale_factor_spatial,
                            patch_size=model_config.patch_size,
                            attention_head_dim=model_config.attention_head_dim,
                            device=accelerator.device,
                        )
                        if model_config.use_rotary_positional_embeddings
                        else None
                    )

                    # Add noise to the model input according to the noise magnitude at each timestep (this is the forward diffusion process)
                    noisy_model_input = scheduler.add_noise(model_input, noise, timesteps)

                    input_intervals = transform_intervals(input_intervals, frames_per_latent=4)
                    output_intervals = transform_intervals(output_intervals, frames_per_latent=4)
                    
                    #first interval is always rep
                    noisy_model_input, target, condition_mask, intervals = random_insert_latent_frame(image_latent, noisy_model_input, model_input, input_intervals, output_intervals, special_info=args.special_info)
                    
                    for i in range(batch_size):
                        if not guidance_mask[i]:
                            noisy_model_input[i][condition_mask[i]] = 0

                    # Predict the noise residual
                    model_output = transformer(
                        hidden_states=noisy_model_input,
                        encoder_hidden_states=prompt_embeds,
                        intervals=intervals,
                        condition_mask=condition_mask,
                        timestep=timesteps,
                        image_rotary_emb=image_rotary_emb,
                        return_dict=False,
                        ablation_condition = ablation_condition
                    )[0]

                    #this line below is also scaling the input which is bad - so the model is also learning to scale this input latent somehow
                    #thus, we need to replace the first frame with the original frame later
                    model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps)



                    alphas_cumprod = scheduler.alphas_cumprod[timesteps]
                    weights = 1 / (1 - alphas_cumprod)
                    while len(weights.shape) < len(model_pred.shape):
                        weights = weights.unsqueeze(-1)



                    loss = torch.mean((weights * (model_pred[~condition_mask] - target[~condition_mask]) ** 2).reshape(batch_size, -1), dim=1)
                    loss = loss.mean()
                    accelerator.backward(loss)

                    if accelerator.state.deepspeed_plugin is None:
                        if not args.just_validate:
                            optimizer.step()
                        optimizer.zero_grad()

                    lr_scheduler.step()

            
                    #wait for all processes to finish
                    accelerator.wait_for_everyone()

                    
                    # Checks if the accelerator has performed an optimization step behind the scenes
                    if accelerator.sync_gradients:
                        progress_bar.update(1)
                        global_step += 1

                        if signal_recieved_time != 0:
                            if time.time() - signal_recieved_time > 60:
                                print("Signal received, saving state and exiting")
                                #accelerator.save_state(save_path)
                                atomic_save(save_path, accelerator)
                                signal_recieved_time = 0
                                exit(0)
                            else:
                                exit(0)

                        if accelerator.is_main_process:
                            if global_step % args.checkpointing_steps == 0:
                                #accelerator.save_state(save_path)
                                atomic_save(save_path, accelerator)
                                logger.info(f"Saved state to {save_path}")

                    logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
                    progress_bar.set_postfix(**logs)
                    accelerator.log(logs, step=global_step)

                    if global_step >= args.max_train_steps:
                        break

            print("Step", step)
            accelerator.wait_for_everyone()
            
            if step == 0 or args.validation_prompt is not None and (step + 1) % args.validation_steps == 0:
                # Create pipeline
                pipe = ControlnetCogVideoXPipeline.from_pretrained(
                    os.path.join(args.base_dir, args.pretrained_model_name_or_path),
                    transformer=unwrap_model(transformer),
                    text_encoder=unwrap_model(text_encoder),
                    vae=unwrap_model(vae),
                    scheduler=scheduler,
                    torch_dtype=weight_dtype,
                )

                print("Length of validation dataset: ", len(val_dataloader))
                #create a pipeline per accelerator device (for faster inference)
                with torch.autocast(str(accelerator.device).replace(":0", ""), enabled=accelerator.mixed_precision == "fp16"):
                    for batch in val_dataloader:
                        frame = ((batch["blur_img"][0].permute(0,2,3,1).cpu().numpy() + 1)*127.5).astype(np.uint8)
                        pipeline_args = {
                            "prompt": "",
                            "negative_prompt": "",
                            "image": frame,
                            "input_intervals": batch["input_intervals"][0:1],
                            "output_intervals": batch["output_intervals"][0:1],
                            "ablation_condition": batch["ablation_condition"][0:1] if "ablation_condition" in batch else None,
                            "guidance_scale": args.guidance_scale,
                            "use_dynamic_cfg": args.use_dynamic_cfg,
                            "height": args.height,
                            "width": args.width,
                            "num_frames": args.max_num_frames,
                            "num_inference_steps": args.num_inference_steps,
                        }

                        modified_filenames = []
                        filenames = batch['file_names']
                        for file in filenames:
                            modified_filenames.append(os.path.splitext(file)[0] + ".mp4")

                        num_frames = batch["num_frames"][0]
                        #save the gt_video output
                        if args.dataset not in ["outsidephotos"]:
                            gt_video = batch["videos"][0].permute(0,2,3,1).cpu().numpy()
                            gt_video = ((gt_video + 1) * 127.5)/255
                            gt_video = gt_video[0:num_frames]
                            
                            for file in modified_filenames:
                                #create the directory if it does not exist
                                gt_file_name = os.path.join(args.output_dir, "gt", modified_filenames[0])
                                os.makedirs(os.path.dirname(gt_file_name), exist_ok=True)
                                if args.dataset in ["baist", "simplebaist"]:
                                    bbox = batch["bbx"][0].cpu().numpy().astype(np.int32)
                                    gt_video = gt_video[:, bbox[1]:bbox[3], bbox[0]:bbox[2], :]
                                    gt_video = np.array([cv2.resize(frame, (160, 192)) for frame in gt_video])

                                save_frames_as_pngs((gt_video*255).astype(np.uint8), gt_file_name.replace(".mp4", "").replace("gt", "gt_frames"))
                                export_to_video(gt_video, gt_file_name, fps=20)


                                if "high_fps_video" in batch:
                                    high_fps_video = batch["high_fps_video"][0].permute(0,2,3,1).cpu().numpy()
                                    high_fps_video = ((high_fps_video + 1) * 127.5)/255
                                    gt_file_name = os.path.join(args.output_dir, "gt_highfps", modified_filenames[0])
                                    
                        
                        if args.dataset in ["adobe", "full", "baist", "outsidephotos", "gopro2x", "goprolarge", "simplebaist"]:
                            for file in modified_filenames:
                                #create the directory if it does not exist
                                blurry_file_name = os.path.join(args.output_dir, "blurry", modified_filenames[0].replace(".mp4", ".png"))
                                #save the blurry image
                                os.makedirs(os.path.dirname(blurry_file_name), exist_ok=True)
                                if args.dataset in ["baist", "simplebaist"]:
                                    bbox = batch["bbx"][0].cpu().numpy().astype(np.int32)
                                    frame0 = frame[0][bbox[1]:bbox[3], bbox[0]:bbox[2], :]
                                    #resize to  192x160
                                    frame0 = cv2.resize(frame0, (160, 192)) 
                                    Image.fromarray(frame0).save(blurry_file_name)
                                else:
                                    Image.fromarray(frame[0]).save(blurry_file_name)

                        videos = log_validation(
                            pipe=pipe,
                            args=args,
                            accelerator=accelerator,
                            pipeline_args=pipeline_args,
                            epoch=epoch,
                        )

                        for i, video in enumerate(videos):
                            prompt = (
                                pipeline_args["prompt"][:25]
                                .replace(" ", "_")
                                .replace(" ", "_")
                                .replace("'", "_")
                                .replace('"', "_")
                                .replace("/", "_")
                            )
                            video = video[0:num_frames]
                            filename = os.path.join(args.output_dir, "deblurred", modified_filenames[0])
                            print("Deblurred file name", filename)
                            os.makedirs(os.path.dirname(filename), exist_ok=True)
                            if args.dataset in ["baist", "simplebaist"]:
                                bbox = batch["bbx"][0].cpu().numpy().astype(np.int32)
                                video = video[:, bbox[1]:bbox[3], bbox[0]:bbox[2], :]
                                #resize to  192x160
                                video = np.array([cv2.resize(frame, (160, 192)) for frame in video])
                            save_frames_as_pngs((video*255).astype(np.uint8), filename.replace(".mp4", "").replace("deblurred", "deblurred_frames"))
                            export_to_video(video, filename, fps=20)

                            accelerator.wait_for_everyone()

                if args.just_validate:
                    exit(0)

    accelerator.wait_for_everyone()
    accelerator.end_training()

signal_recieved_time = 0

def handle_signal(signum, frame):
    global signal_recieved_time
    signal_recieved_time = time.time()

    print(f"Signal {signum} received at {time.ctime()}")

    with open("/datasets/sai/gencam/cogvideox/interrupted.txt", "w") as f:
        f.write(f"Training was interrupted at {time.ctime()}")

if __name__ == "__main__":

    args = get_args()

    print("Registering signal handler")
    #Register the signal handler (catch SIGUSR1)
    signal.signal(signal.SIGUSR1, handle_signal)

    main_thread = threading.Thread(target=main, args=(args,))
    main_thread.start()

    print("SIGNAL RECIEVED TIME", signal_recieved_time)
    while signal_recieved_time!= 0:
        time.sleep(1)
    
    #call main with args as a thread