multi_view/DiffSynth-Studio-main/diffsynth/trainers/utils.py

import imageio, os, torch, warnings, torchvision, argparse, json
from peft import LoraConfig, inject_adapter_in_model
from PIL import Image
import pandas as pd
from tqdm import tqdm
from accelerate import Accelerator
from accelerate.utils import DataLoaderConfiguration, set_seed
import wandb
import tempfile
import requests
import random
import decord
import cv2 
import numpy as np
import shutil
import imageio.v3 as iio
import imageio_ffmpeg as ffmpeg
import math
import re
from diffsynth.trainers.timer import get_timers
import time
import glob
from safetensors.torch import save_file, load_file
import math
import random
from typing import List, Dict, Iterable, Iterator, Sequence, Optional
import torch
from torch.utils.data import Dataset, DataLoader, Sampler, BatchSampler, DistributedSampler
import torch.distributed as dist
from tqdm import tqdm
from multiprocessing import Pool, cpu_count
import math
import random
from collections import defaultdict, deque
from typing import Iterable, List, Sequence, Optional, Callable, Dict

import torch
from torch.utils.data import DataLoader, Sampler, BatchSampler, RandomSampler, DistributedSampler
from accelerate.utils import DataLoaderConfiguration, DeepSpeedPlugin
import random
import matplotlib.pyplot as plt

class DiffusionTrainingModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        
        
    def to(self, *args, **kwargs):
        for name, model in self.named_children():
            model.to(*args, **kwargs)
        return self
        
        
    def trainable_modules(self):
        trainable_modules = filter(lambda p: p.requires_grad, self.parameters())
        return trainable_modules
    
    
    def trainable_param_names(self):
        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters()))
        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])
        return trainable_param_names
    
    
    def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None):
        if lora_alpha is None:
            lora_alpha = lora_rank
        lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules)
        model = inject_adapter_in_model(lora_config, model)
        return model
    
    
    def export_trainable_state_dict(self, state_dict, remove_prefix=None):
        trainable_param_names = self.trainable_param_names()
        state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names}
        if remove_prefix is not None:
            state_dict_ = {}
            for name, param in state_dict.items():
                if name.startswith(remove_prefix):
                    name = name[len(remove_prefix):]
                state_dict_[name] = param
            state_dict = state_dict_
        return state_dict



class ModelLogger:
    """
    一个集成了 Accelerate、W&B、TensorBoard 和文件日志的日志记录器。
    现在它也负责跟踪全局训练步数，并能被 Accelerator 保存和恢复。
    """
    def __init__(self, accelerator: Accelerator, output_path: str):
        self.accelerator = accelerator
        self.output_path = output_path
        
        # 关键状态：全局步数。将被 Accelerator 自动管理。
        self.global_step = 0
        
        # 注册自身，以便 `accelerator.save_state` 和 `load_state` 可以管理其状态。
        self.accelerator.register_for_checkpointing(self)

        self.timers = get_timers()
        
        # --- 日志文件设置 ---
        log_dir = os.path.join(self.output_path, "log")
        if self.accelerator.is_main_process:
            os.makedirs(log_dir, exist_ok=True)
        self.accelerator.wait_for_everyone()
        
        log_file_path = os.path.join(log_dir, f"rank_{self.accelerator.process_index}.log")
        # 'a+' 模式允许读写，如果文件不存在则创建
        self.file_logger = open(log_file_path, "a+") 
        self.log_to_file(f"Logger initialized. Current global_step: {self.global_step}")
    
    def state_dict(self):
        """
        返回需要被保存的状态。这是 `register_for_checkpointing` 要求的方法。
        """
        return {"global_step": self.global_step}

    def load_state_dict(self, state_dict):
        """
        从 state_dict 中加载状态。这是 `register_for_checkpointing` 要求的方法。
        """
        self.global_step = state_dict["global_step"]
        # 在加载状态后记录日志，便于调试
        self.log_to_file(f"Logger state restored. Resumed global_step: {self.global_step}")

    def log_to_file(self, message: str):
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        self.file_logger.write(f"[{timestamp}] {message}\n")
        self.file_logger.flush()

    def on_step_end(self, loss: torch.Tensor, optimizer: torch.optim.Optimizer):
        """在每个实际的优化器步骤后调用。"""
        step_time_ms = self.timers.step_time.stop()
        step_time_s = step_time_ms / 1000.0 # 转换为秒

        gathered_losses = self.accelerator.gather_for_metrics(loss.detach())
        avg_loss = torch.mean(gathered_losses).item()
        local_loss = loss.item()
        learning_rate = optimizer.param_groups[0]['lr']

        if self.accelerator.is_main_process:
            log_dict = {
                "train/loss": avg_loss,
                "train/learning_rate": learning_rate,
                "perf/step_time_seconds": step_time_s,
                "progress/epoch": self.current_epoch, # 假设 epoch 信息已设置
                "progress/step": self.global_step,
            }
            self.accelerator.log(log_dict, step=self.global_step)

        log_message = (
            f"Epoch: {self.current_epoch} | "
            f"Iter: {self.global_step:06d} | "
            f"Step Time: {step_time_s:.3f}s | "
            f"Local Loss: {local_loss:.4f} | "
            f"Avg Loss: {avg_loss:.4f} | "
            f"LR: {learning_rate:.6f}"
        )
        self.log_to_file(log_message)
        
        self.global_step += 1
        self.step_start_time = time.time()

    def set_epoch(self, epoch: int):
        """由外部训练循环设置当前 epoch。"""
        self.current_epoch = epoch

    def close(self):
        self.log_to_file("Logger closing.")
        self.file_logger.close()


class CheckpointManager:
    """
    使用 Accelerator 管理检查点的保存和加载。
    - 只保存可训练的权重以节省空间。
    - 自动轮换检查点，只保留最新的 N 个。
    """
    def __init__(self, accelerator: Accelerator, output_path: str, save_steps: int, save_epoches: int, max_to_keep: int, state_dict_converter=lambda x:x):
        self.accelerator = accelerator
        self.checkpoints_dir = os.path.join(output_path)
        self.save_steps = save_steps
        self.save_epoches = save_epoches
        self.max_to_keep = max_to_keep
        self.state_dict_converter = state_dict_converter
        
        if self.accelerator.is_main_process:
            os.makedirs(self.checkpoints_dir, exist_ok=True)

    # ------ 获取 DeepSpeed 引擎 ------
    def get_deepspeed_engine(self, model):
        try:
            import deepspeed
            if isinstance(model, deepspeed.DeepSpeedEngine):
                return model
        except Exception:
            pass
        st = getattr(self.accelerator, "state", None)
        if st is not None:
            if getattr(st, "deepspeed_engine", None) is not None:
                return st.deepspeed_engine
            plugin = getattr(st, "deepspeed_plugin", None)
            if plugin is not None:
                if getattr(plugin, "engine", None) is not None:
                    return plugin.engine
                if getattr(plugin, "deepspeed_engine", None) is not None:
                    return plugin.deepspeed_engine
        return None

    def save_checkpoint(self, model: torch.nn.Module, global_step: int, epoch_id: int):
        """
        Zero-2 + Accelerate 场景下的推荐写法：
        - 所有 rank：accelerator.save_state() 保存 deepspeed + optimizer + scheduler + RNG
        - rank0 再额外导出 trainable 权重为 weights.safetensors
        """
        if not (global_step > 0 and global_step % self.save_steps == 0):
            return

        checkpoint_name = f"checkpoint-step-{global_step}-epoch-{epoch_id + 1}"
        save_path = os.path.join(self.checkpoints_dir, checkpoint_name)

        # 1) rank0 创建目录
        if self.accelerator.is_main_process:
            os.makedirs(save_path, exist_ok=True)
        self.accelerator.wait_for_everyone()

    def save_epoch_checkpoint(self, model: torch.nn.Module, global_step: int, epoch_id: int):
        if self.save_epoches <= 0:
            return
        if (epoch_id + 1) % self.save_epoches != 0:
            return

        checkpoint_name = f"checkpoint-step-{global_step}-epoch-{epoch_id + 1}"
        save_path = os.path.join(self.checkpoints_dir, checkpoint_name)

        if self.accelerator.is_main_process:
            os.makedirs(save_path, exist_ok=True)
        self.accelerator.wait_for_everyone()

        self.accelerator.print(f"[CKPT] epoch save_state -> {save_path}")
        self.accelerator.save_state(save_path)

        if self.accelerator.is_main_process:
            with open(os.path.join(save_path, "trainer_state.json"), "w") as f:
                json.dump({"global_step": global_step}, f)
            state_dict = self.accelerator.get_state_dict(model)
            state_dict = self.accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix="pipe.dit.")
            state_dict = self.state_dict_converter(state_dict)

            weights_path = os.path.join(save_path, "weights.safetensors")
            self.accelerator.save(state_dict, weights_path, safe_serialization=True)
            self.accelerator.print(f"[CKPT] epoch trainable weights saved -> {weights_path}")
            self._rotate_checkpoints()
        self.accelerator.wait_for_everyone()
        # 2) 让 Accelerate 自己处理 DeepSpeed / 优化器 / scheduler / 随机数
        #    注意：所有 rank 都要调用一次
        self.accelerator.print(f"[CKPT] accelerator.save_state -> {save_path}")
        self.accelerator.save_state(save_path)

        # 3) rank0 额外保存可训练的权重为 safetensors
        if self.accelerator.is_main_process:
            with open(os.path.join(save_path, "trainer_state.json"), "w") as f:
                json.dump({"global_step": global_step}, f)
            state_dict = self.accelerator.get_state_dict(model)
            state_dict = self.accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix="pipe.dit.")
            state_dict = self.state_dict_converter(state_dict)

            weights_path = os.path.join(save_path, "weights.safetensors")
            self.accelerator.save(state_dict, weights_path, safe_serialization=True)
            self.accelerator.print(f"[CKPT] trainable weights saved -> {weights_path}")
            # 滚动删除旧 ckpt
            self._rotate_checkpoints()
        self.accelerator.wait_for_everyone()

    def load_checkpoint(self, model, resume_from_checkpoint):
        if not resume_from_checkpoint:
            return None

        # 1. 找到 checkpoint 路径
        if resume_from_checkpoint == "latest":
            load_path = self._get_load_path(resume_from_checkpoint)
            if not load_path:
                self.accelerator.print("No checkpoint found to resume from. Starting from scratch.")
                return
        else:
            load_path = resume_from_checkpoint

        if load_path is None:
            self.accelerator.print("[Resume] No checkpoint found")
            return None

        self.accelerator.print(f"[Resume] Loading from: {load_path}")

        # 2. 恢复 deepspeed + optimizer + scheduler + RNG + model states
        try:
            self.accelerator.load_state(load_path)
            self.accelerator.print("[Resume] accelerator.load_state() DONE.")
        except Exception as e:
            self.accelerator.print(f"[Resume] ERROR: {e}")

        # 3. 多卡同步
        self.accelerator.wait_for_everyone()
        # self.accelerator.broadcast_parameters(model)

        self.accelerator.print("[Resume] Model fully restored.")
        return load_path


    def _get_load_path(self, resume_from_checkpoint: str | bool) -> str | None:
        """辅助函数，解析 resume_from_checkpoint 并返回有效的加载路径。"""
        if resume_from_checkpoint is True or str(resume_from_checkpoint).lower() == "latest":
            return self.find_latest_checkpoint()
        
        # 如果提供了具体路径，直接返回
        if os.path.isdir(resume_from_checkpoint):
            return resume_from_checkpoint
        return None

    def find_latest_checkpoint(self) -> str | None:
        """在检查点目录中找到最新的检查点。"""
        all_checkpoints = glob.glob(os.path.join(self.checkpoints_dir, "checkpoint-step-*"))
        if not all_checkpoints:
            return None
        
        try:
            latest_checkpoint = max(
                all_checkpoints, 
                key=lambda path: int(re.search(r'checkpoint-step-(\d+)', path).group(1))
            )
            return latest_checkpoint
        except (ValueError, AttributeError):
            # 如果文件夹命名不规范，无法解析出数字，则返回 None
            return None

    def _rotate_checkpoints(self):
        """删除旧的检查点，只保留 `max_to_keep` 个。"""
        if self.max_to_keep <= 0:
            return

        all_checkpoints = glob.glob(os.path.join(self.checkpoints_dir, "checkpoint-step-*"))
        
        # 解析步数并排序，最旧的在前
        try:
            sorted_checkpoints = sorted(
                all_checkpoints,
                key=lambda path: int(re.search(r'checkpoint-step-(\d+)-epoch-\d+', path).group(1))
            )
        except (ValueError, AttributeError):
            self.accelerator.print("[Warning] Could not sort checkpoints for rotation due to naming issues.")
            return

        # 如果检查点数量超过限制，则删除最旧的
        num_to_delete = len(sorted_checkpoints) - self.max_to_keep
        if num_to_delete > 0:
            checkpoints_to_delete = sorted_checkpoints[:num_to_delete]
            for ckpt_path in checkpoints_to_delete:
                self.accelerator.print(f"Deleting old checkpoint: {ckpt_path}")
                shutil.rmtree(ckpt_path)



def launch_training_task(
    args,
    dataset: torch.utils.data.Dataset,
    model: DiffusionTrainingModule,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    num_epochs: int = 1,
    gradient_accumulation_steps: int = 1,
    output_path: str = "./models/train",
    save_steps: int = 10, # 新增：每多少步保存一次
    save_epoches: int = 1,
    max_checkpoints_to_keep=5, # 最多只保留 5 个最新的检查点
    resume_from_checkpoint: str | bool = "latest", # 新增：从何处恢复
    seed: int = 42,
    visual_log_project_name: str=None,
):
    def collate_fn_identity(batch):
        return batch 
    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size = args.batch_size, collate_fn = collate_fn_identity, num_workers = 8)
    # dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0])
    # dispatch_batches=False 可以绕过accelerator自动的组建batch，组batch 操作放在./dataset/collate_fns/下自定义
    # --- 1. 初始化 Accelerator 和组件 ---
    accelerator = Accelerator(
        gradient_accumulation_steps=gradient_accumulation_steps,
        # log_with=["wandb"],
        project_dir=output_path,
    )
    set_seed(seed, device_specific=True)
    # if visual_log_project_name is None:
    #     visual_log_project_name = os.path.join(output_path, "wandb")
    # 在主进程上初始化 tracker
    if accelerator.is_main_process:
        accelerator.init_trackers(project_name=visual_log_project_name)
    model_logger = ModelLogger(accelerator, output_path)
    checkpoint_manager = CheckpointManager(accelerator, output_path, save_steps, save_epoches, max_checkpoints_to_keep)

    # --- 3. 准备所有组件 ---
    # 注意：顺序很重要，先加载状态，再调用 prepare
    model, optimizer, dataloader, scheduler = accelerator.prepare(
        model, optimizer, dataloader, scheduler
    )

    # --- 2. 加载状态 ---
    if resume_from_checkpoint:
        load_path = checkpoint_manager.load_checkpoint(model, "latest")
        if load_path:
            state_path = os.path.join(load_path, "trainer_state.json")
            if os.path.exists(state_path):
                state = json.load(open(state_path))
                global_step = state.get("global_step", 0)
                model_logger.global_step = global_step
            else:
                accelerator.print(f"[Resume] trainer_state.json not found in {load_path}, start from step 0.")
                model_logger.global_step = 0
        else:
            accelerator.print("[Resume] No valid checkpoint path, start from scratch.")
            model_logger.global_step = 0
    
    ##TODO 加载学习率不对  需要修改！！！！！！！！！！！目前先手动写死了
    ## --- 4. 手动设置学习率（如需） ---


    # --- 4. 计算恢复后的起始点 ---
    # model_logger.global_step 已经被 load_state 恢复了
    total_steps_per_epoch = len(dataloader)
    resume_step = model_logger.global_step % total_steps_per_epoch
    starting_epoch = model_logger.global_step // total_steps_per_epoch
    
    accelerator.print("--- Starting Training ---")
    accelerator.print(f"Num Epochs: {num_epochs}")
    accelerator.print(f"Total steps per epoch: {total_steps_per_epoch}")
    accelerator.print(f"Resuming from Epoch: {starting_epoch}, Step: {resume_step}")
    accelerator.print(f"Total number of data: {len(dataloader)}")

    timers = get_timers()
    if accelerator.is_main_process:
        save_loss_path = os.path.join(output_path, "training_loss_plots", args.visual_log_project_name)
        loss_means = []
        os.makedirs(save_loss_path, exist_ok=True)

    def _save_video_frames(frames, path, fps=16):
        if not frames:
            return
        writer = imageio.get_writer(path, fps=fps)
        try:
            for frame in frames:
                if isinstance(frame, torch.Tensor):
                    frame = frame.detach().cpu().numpy()
                if hasattr(frame, "convert"):
                    frame = np.asarray(frame.convert("RGB"))
                frame = np.asarray(frame)
                if frame.dtype != np.uint8:
                    frame = np.clip(frame, 0, 255).astype(np.uint8)
                writer.append_data(frame)
        finally:
            writer.close()

    def _run_debug_infer(epoch_id):
        if not getattr(args, "debug_infer", False):
            return
        interval = int(getattr(args, "debug_infer_interval", 1))
        if interval <= 0 or epoch_id % interval != 0:
            return
        if not accelerator.is_main_process:
            return

        indices = getattr(args, "debug_infer_indices", [0])
        if isinstance(indices, int):
            indices = [indices]
        if not isinstance(indices, (list, tuple)):
            indices = [0]

        debug_dir = os.path.join(output_path, "debug_infer", f"epoch_{epoch_id}")
        os.makedirs(debug_dir, exist_ok=True)
        debug_log = os.path.join(debug_dir, "debug.log")

        def log(msg):
            print(msg)
            with open(debug_log, "a", encoding="utf-8") as f:
                f.write(msg + "\n")

        model_unwrapped = accelerator.unwrap_model(model)
        model_unwrapped.eval()
        pipe = model_unwrapped.pipe
        orig_pipe_device = getattr(pipe, "device", None)
        orig_pipe_dtype = getattr(pipe, "torch_dtype", None)
        orig_scheduler = pipe.scheduler
        orig_timesteps = getattr(orig_scheduler, "timesteps", None)
        orig_sigmas = getattr(orig_scheduler, "sigmas", None)
        orig_training = getattr(orig_scheduler, "training", False)
        orig_linear_weights = getattr(orig_scheduler, "linear_timesteps_weights", None)

        pipe.device = accelerator.device
        pipe.torch_dtype = torch.bfloat16

        cfg_scale = float(getattr(args, "debug_infer_cfg_scale", 5.0))
        cfg_scale_face = float(getattr(args, "debug_infer_cfg_scale_face", 5.0))
        num_steps = int(getattr(args, "debug_infer_steps", 8))
        seed = int(getattr(args, "debug_infer_seed", args.seed))
        tiled = bool(getattr(args, "debug_infer_tiled", True))
        use_input_video = bool(getattr(args, "debug_infer_use_input_video", False))
        negative_prompt = getattr(args, "debug_infer_negative_prompt", "")

        try:
            for idx in indices:
                if idx < 0 or idx >= len(dataset):
                    log(f"[debug_infer] index {idx} out of range.")
                    continue
                sample = dataset[idx]
                input_video = sample.get("video", [])
                ref_images = sample.get("ref_images", [])
                if not input_video or not ref_images:
                    log(f"[debug_infer] sample {idx} missing video/ref, skip.")
                    continue

                prompt = sample.get("pre_shot_caption", ["xxx"])
                if isinstance(prompt, str):
                    prompt = [prompt]

                log("=" * 80)
                log(f"[debug_infer] epoch={epoch_id} index={idx}")
                log(f"video_path={sample.get('video_path')}")
                log(f"num_frames={len(input_video)} ref_num={sample.get('ref_num')} ID_num={sample.get('ID_num')}")
                log(f"prompt={prompt}")

                ref_dir = os.path.join(debug_dir, f"ref_{idx}")
                os.makedirs(ref_dir, exist_ok=True)
                for id_i, ref_group in enumerate(ref_images):
                    for img_i, img in enumerate(ref_group):
                        img.save(os.path.join(ref_dir, f"id{id_i}_img{img_i}.png"))
            if use_input_video:
                _save_video_frames(input_video, os.path.join(debug_dir, f"input_{idx}.mp4"), fps=16)

                with torch.no_grad():
                    video, _ = pipe(
                        args=args,
                        prompt=[prompt],
                        negative_prompt=[negative_prompt],
                        input_video=[input_video] if use_input_video else None,
                        ref_images=[ref_images],
                        seed=seed,
                        tiled=tiled,
                        height=input_video[0].size[1],
                        width=input_video[0].size[0],
                        num_frames=len(input_video),
                        cfg_scale=cfg_scale,
                        cfg_scale_face=cfg_scale_face,
                        num_inference_steps=num_steps,
                        num_ref_images=sample.get("ref_num"),
                    )
                _save_video_frames(video, os.path.join(debug_dir, f"output_{idx}.mp4"), fps=16)
        finally:
            if orig_pipe_device is not None:
                pipe.device = orig_pipe_device
            if orig_pipe_dtype is not None:
                pipe.torch_dtype = orig_pipe_dtype
            if orig_timesteps is not None:
                orig_scheduler.timesteps = orig_timesteps
            if orig_sigmas is not None:
                orig_scheduler.sigmas = orig_sigmas
            orig_scheduler.training = orig_training
            if orig_linear_weights is not None:
                orig_scheduler.linear_timesteps_weights = orig_linear_weights
            model_unwrapped.train()
    for epoch_id in range(starting_epoch, num_epochs):
        model.train()
        model_logger.set_epoch(epoch_id)
        
        # 如果是恢复的第一个 epoch，需要跳过已经训练过的 steps
        if epoch_id == starting_epoch and resume_step > 0:
            # 使用 enumerate 和 islice 的组合来跳过
            active_dataloader = enumerate(dataloader)
            for _ in range(resume_step):
                next(active_dataloader)
        else:
            active_dataloader = enumerate(dataloader)

        pbar = tqdm(
            active_dataloader, 
            initial=resume_step if epoch_id == starting_epoch else 0,
            total=total_steps_per_epoch,
            disable=not accelerator.is_main_process, 
            desc=f"Epoch {epoch_id}"
        )

        for step, data in pbar:
            timers.step_time.start()
            with accelerator.accumulate(model):
                optimizer.zero_grad()
                loss = model(data, args)
                accelerator.backward(loss)
                optimizer.step()
                scheduler.step()
                # 只在真正的优化步记录（累积步结束）
                if accelerator.sync_gradients:
                    model_logger.on_step_end(loss, optimizer)
                    checkpoint_manager.save_checkpoint(model, model_logger.global_step, epoch_id)
                    # 跨卡聚合 loss（每卡一个标量），取全局平均
                    loss_mean = accelerator.gather(loss.detach()).mean().item()
                if model_logger.global_step % 10 == 0:
                    if accelerator.is_main_process:
                        loss_means.append(loss_mean)
                        plt.figure(figsize=(8, 5))
                        plt.plot(loss_means, marker="o", linestyle="-", label="Training Loss")
                        plt.xlabel("X (every 10 steps)")
                        plt.ylabel("Loss")
                        plt.title(f"Loss Curve up to step {model_logger.global_step}")
                        plt.grid(True)
                        plt.legend()
                        plt.savefig(f"{save_loss_path}/loss_mean.png")
                        plt.close()

        checkpoint_manager.save_epoch_checkpoint(model, model_logger.global_step, epoch_id)
        _run_debug_infer(epoch_id)

    # 重置 resume_step，因为后续 epoch 都是从头开始   
    resume_step = 0

    # --- 6. 训练结束 ---
    accelerator.wait_for_everyone()
    accelerator.print("--- Training Finished ---")
    
    # 保存最终的模型
    if accelerator.is_main_process:
        accelerator.print("Saving final model...")
        final_save_path = os.path.join(checkpoint_manager.checkpoints_dir, "final_model")
        accelerator.save_state(final_save_path)
    
    accelerator.end_training()
    model_logger.close()

    # if accelerator.is_main_process:
    #     wandb.finish()


def wan_parser():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument("--dataset_base_path", type=str, default="", required=False, help="Base path of the dataset.")
    parser.add_argument("--dataset_metadata_path", type=str, default=None, help="Path to the metadata file of the dataset.")
    parser.add_argument("--max_pixels", type=int, default=1280*720, help="Maximum number of pixels per frame, used for dynamic resolution..")
    parser.add_argument("--height", type=int, default=None, help="Height of images or videos. Leave `height` and `width` empty to enable dynamic resolution.")
    parser.add_argument("--width", type=int, default=None, help="Width of images or videos. Leave `height` and `width` empty to enable dynamic resolution.")
    parser.add_argument("--num_frames", type=int, default=81, help="Number of frames per video. Frames are sampled from the video prefix.")
    parser.add_argument("--data_file_keys", type=str, default="image,video", help="Data file keys in the metadata. Comma-separated.")
    parser.add_argument("--dataset_repeat", type=int, default=1, help="Number of times to repeat the dataset per epoch.")
    parser.add_argument("--model_paths", type=str, default=None, help="Paths to load models. In JSON format.")
    parser.add_argument("--model_id_with_origin_paths", type=str, default=None, help="Model ID with origin paths, e.g., Wan-AI/Wan2.1-T2V-1.3B:diffusion_pytorch_model*.safetensors. Comma-separated.")
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate.")
    parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs.")
    parser.add_argument("--output_path", type=str, default="./checkpoints", help="Output save path.")
    parser.add_argument("--remove_prefix_in_ckpt", type=str, default="pipe.dit.", help="Remove prefix in ckpt.")
    parser.add_argument("--trainable_models", type=str, default=None, help="Models to train, e.g., dit, vae, text_encoder.")
    parser.add_argument("--lora_base_model", type=str, default=None, help="Which model LoRA is added to.")
    parser.add_argument("--lora_target_modules", type=str, default="q,k,v,o,ffn.0,ffn.2", help="Which layers LoRA is added to.")
    parser.add_argument("--lora_rank", type=int, default=32, help="Rank of LoRA.")
    parser.add_argument("--extra_inputs", default=None, help="Additional model inputs, comma-separated.")
    parser.add_argument("--use_gradient_checkpointing_offload", default=False, action="store_true", help="Whether to offload gradient checkpointing to CPU memory.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Gradient accumulation steps.")
    parser.add_argument("--max_timestep_boundary", type=float, default=1.0, help="Max timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
    parser.add_argument("--min_timestep_boundary", type=float, default=0.0, help="Min timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
    parser.add_argument("--save_steps", type=int, default=10, help="The epoch to save.")
    parser.add_argument("--train_yaml", type=str, default="../../../../conf/config.yaml", help="The train yaml file.")
    parser.add_argument("--max_checkpoints_to_keep", type=int, default=5, help="max_checkpoints_to_keep")
    parser.add_argument("--resume_from_checkpoint", type=bool, default=False, help="The dataset yaml file")
    parser.add_argument("--seed", type=int, default=42, help="The random seed")
    parser.add_argument("--visual_log_project_name", type=str, default=None, help="Project name")
    parser.add_argument("--max_frames_per_batch", type=int, default=136, help="Max frames num for each batch")
    parser.add_argument("--prompt_index", type=int, default=-1, help="Prompt index for generation")
    parser.add_argument("--ref_num", type=int, default=3, help="The number of reference images")
    parser.add_argument("--local_model_path", type=str, default="", help="The default root path of the Wan weights")
    parser.add_argument("--batch_size", type=int, default=1, help="The default batch size")
    parser.add_argument("--save_epoches", type=int, default=1, help="The default saving epoch")
    parser.add_argument("--split_rope", type=bool, default=False, help="Whether apply different rope into reference images ")
    parser.add_argument("--split1", type=bool, default=False, help=" ")
    parser.add_argument("--split2", type=bool, default=False, help=" ")
    parser.add_argument("--split3", type=bool, default=False, help=" ")
    parser.add_argument("--shot_rope", type=bool, default=False, help="Whether apply shot rope for multi-shot video")

    return parser


if __name__ == '__main__':

    dataset = VideoDataset(use_history = True, base_path = "/user/kg-aigc/rd_dev/qizipeng/luo_data_sorted_full.json")

    train_dataset = dataset.data
    for d in tqdm(train_dataset[7772:7773]):
        # video = dataset.get_video_from_path(d["video_path"], d["video_time"], is_history=False)
        # if len(video) != 81 :
        #     print(d)

        video =  dataset.get_video_from_path(
                            d["videoHis_path"], d["videoHis_time"], is_history=True
                        )
        if video == None:
            print(d)


    import pdb; pdb.set_trace()