scripts/pavrm/inference_pavrm.py

import argparse
import json
import logging
import os
import time
import itertools
from copy import deepcopy
from collections import deque
from easydict import EasyDict

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import random

import torch.amp as amp
from diffusers.optimization import get_scheduler
from einops import rearrange
from omegaconf import OmegaConf
from peft import LoraConfig, get_peft_model
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score

from diffusers_lite.constants import PRECISION_TO_TYPE
from diffusers_lite.datasets.image2video_dataset import Image2VideoTrainDataset
from diffusers_lite.schedulers.scheduling_flow_match_discrete import (
    FlowMatchDiscreteScheduler,
)
from diffusers_lite.wan.modules.model import WanModel
from diffusers_lite.wan.modules.t5 import T5EncoderModel
from diffusers_lite.wan.modules.vae import WanVAE
from diffusers_lite.wan.modules.clip import CLIPModel
from diffusers_lite.utils.communication import (
    broadcast,
    sp_parallel_dataloader_wrapper_wanx,
)
from diffusers_lite.utils.data_utils import (
    LengthGroupedSampler,
    save_videos_grid,
    crop_tensor,
    BlockDistributedSampler,
    VideoImageBatchIterator
)
from diffusers_lite.utils.fsdp_utils import (
    apply_fsdp_checkpointing,
    get_dit_fsdp_kwargs,
)
from diffusers_lite.utils.parallel_states import initialize_sequence_parallel_state,nccl_info,get_sequence_parallel_state
from diffusers_lite.utils.torch_utils import set_manual_seed, free_memory, set_logging, set_worker_seed_builder
from diffusers_lite.utils.diffusion_utils import (
    batch2list,
    list2batch,
    vae_encode,
    vae_decode,
    image_encode,
    prompt2states,
    load_lora_state_dict,
    transformer_zero_init,
    prepare_video_condition_wanx,
    stable_mse_loss,
)
from diffusers_lite.utils.model_utils import (
    save_lora_checkpoint,
    save_checkpoint,
    load_state_dict,
    print_parameters_information,
    update_ema_model,
)

from diffusers_lite.utils.network import MLP, QueryAttention, forward_siamese, forward_mlp, train_model, save_model

NAME_MAPPING = {
    "t2v-1.3b": "Wan2.1-T2V-1.3B",
    "t2v-14b": "Wan2.1-T2V-14B",
    "i2v-1.3b": "Wan2.1-T2V-1.3B",
    "i2v-14b-480p": "Wan2.1-I2V-14B-480P",
    "i2v-14b-720p": "Wan2.1-I2V-14B-720P",
    "flf2v-14b-720p": "Wan2.1-FLF2V-14B-720P",
}

def validate_model_parameters(model, model_name="model"):
    has_invalid = False
    total_params = 0
    trainable_params = 0
    
    for name, param in model.named_parameters():
        total_params += 1
        if param.requires_grad:
            trainable_params += 1
            
        if torch.isnan(param).any():
            logging.error(f"ERROR: {model_name} parameter {name} contains NaN values!")
            has_invalid = True
        if torch.isinf(param).any():
            logging.error(f"ERROR: {model_name} parameter {name} contains Inf values!")
            has_invalid = True
    
    logging.info(f"{model_name}: {trainable_params}/{total_params} parameters are trainable")
    
    if has_invalid:
        logging.warning(f"WARNING: {model_name} has invalid parameters!")
        return False
    return True

def basic_init(config):
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    dist.init_process_group("nccl")
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dtype = PRECISION_TO_TYPE[config.train.precision]
    initialize_sequence_parallel_state(config.dataset.sp_size)
    set_logging(local_rank)

    set_manual_seed(config.train.seed + nccl_info.group_id)
    logging.info(f"lanuch with seed {config.train.seed + rank}")

    config.save.ckpt_dir = os.path.join(
        config.save.output_dir, f"{config.train_id}/checkpoints"
    )
    config.save.log_dir = os.path.join(
        config.save.output_dir, f"{config.train_id}/logs"
    )
    config.save.sanity_check_dir = f"outputs/sanity_check/wanx/{config.train_id}"
    config.save.tensorboard_dir = os.path.join(config.save.output_dir, f"{config.train_id}/tensorboard")
    config.save.mlp_dir = os.path.join(config.save.output_dir, f"{config.train_id}/mlp")
    log_path = os.path.join(config.save.log_dir, "log.txt")

    if rank == 0:
        os.makedirs(config.save.output_dir, exist_ok=True)
        os.makedirs(config.save.ckpt_dir, exist_ok=True)
        os.makedirs(config.save.log_dir, exist_ok=True)
        os.makedirs(config.save.tensorboard_dir, exist_ok=True)
        OmegaConf.save(config, os.path.join(config.save.log_dir, "train_config.yaml"))
        if not os.path.exists(log_path):
            with open(log_path, "w") as f:
                f.write(f"Start logging {config.train_id}:\n")
        if config.train.sanity_check_interval > 0:
            os.makedirs(config.save.sanity_check_dir, exist_ok=True)
        logging.info(f"save ckpt directory {config.save.ckpt_dir}")

    if config.train.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True
        logging.info(f"enable TF32")

    basic_kwargs = EasyDict(
        {
            "local_rank": local_rank,
            "rank": rank,
            "world_size": world_size,
            "device": device,
            "dtype": dtype,
            "log_path": log_path,
        }
    )

    return config, basic_kwargs

def model_init(config, basic_kwargs):
    assert config.task in NAME_MAPPING.keys()
    base_dir = config.model.base_path

    if config.model.init_transformer_path:
        logging.info(f"loading model tranformer from {config.model.init_transformer_path}") 
        transformer = WanModel.from_pretrained(config.model.init_transformer_path)
        resume_step = 0
    else:
        if config.task in [
            "t2v-1.3b",
            "t2v-14b",
            "i2v-14b-480p",
            "i2v-14b-720p",
            "flf2v-14b-720p",
        ]:
            logging.info(f"loading model tranformer from {base_dir}")
            transformer = WanModel.from_pretrained(base_dir)
        elif config.task in ["i2v-1.3b"]:
            transformer_config = json.load(
                open(os.path.join(base_dir, "config.json"), "r")
            )
            transformer_config["in_dim"] = 36
            transformer_config["model_type"] = "i2v"
            transformer = WanModel.from_config(transformer_config)
            transformer = transformer_zero_init(transformer)
            state_dict = load_state_dict(model_dir=base_dir)

            del state_dict["patch_embedding.bias"]
            del state_dict["patch_embedding.weight"]

            m, u = transformer.load_state_dict(state_dict, strict=False)
            logging.info(f"load lora from {base_dir}.")
            logging.info(f"miss {len(m)}; unexpect {len(u)}.")
        resume_step = 0

    frozen_modules = [
        'patch_embedding',
        'text_embedding',
        'time_embedding',
        'time_projection',
        'img_emb',
        # 'freqs',
    ]

    for module_name in frozen_modules:
        if hasattr(transformer, module_name):
            module = getattr(transformer, module_name)
            for param in module.parameters():
                param.requires_grad = False

    trainable_blocks = config.lrm.trainable_blocks

    logging.info(f"freezing all blocks except for {trainable_blocks}")

    new_blocks = []
    for i, block in enumerate(transformer.blocks):
        if i in trainable_blocks:
            logging.info(f"block {i} is set to be trainable.")
            for param in block.parameters():
                param.requires_grad = True
            new_blocks.append(block)
        else:
            logging.info(f"block {i} is frozen and removed.")
            for param in block.parameters():
                param.requires_grad = False

    transformer.blocks = nn.ModuleList(new_blocks)

    if hasattr(transformer, 'head'):
        del transformer.head
        transformer.head = None

    transformer.__class__.enable_teacache = False

    if config.model.lora.use_lora:
        lora_config = LoraConfig(
            r=config.model.lora.lora_rank,
            lora_alpha=config.model.lora.lora_rank,
            init_lora_weights=True,
            target_modules=config.model.lora.target_modules,
        )
        transformer = get_peft_model(transformer, lora_config)
        if config.model.lora.resume_lora_path:
            lora_state_dict = load_lora_state_dict(config.model.lora.resume_lora_path)
            m, u = transformer.load_state_dict(lora_state_dict, strict=False)
            logging.info(f"load lora from {config.model.lora.resume_lora_path}.")
            logging.info(f"miss {len(m)}; unexpect {len(u)}.")
            resume_step = int(config.model.lora.resume_lora_path.split("-")[-1])
    
    if config.model.resume_transformer_path:
        logging.info(f"loading model tranformer from {config.model.resume_transformer_path}")
        state_dict = load_state_dict(model_dir=config.model.resume_transformer_path)
        m, u = transformer.load_state_dict(state_dict, strict=False)
        logging.info(f"miss {len(m)}; unexpect {len(u)}.")
        resume_step = int(config.model.resume_transformer_path.split("-")[-1].split('.')[0])
    
    transformer = transformer.to(dtype=torch.float32)

    if config.model.ema.use_ema:
        logging.info("loading ema model")
        ema_transformer = deepcopy(transformer)

    else:
        ema_transformer = None

    fsdp_kwargs, no_split_modules = get_dit_fsdp_kwargs(
        transformer,
        config.model.fsdp.fsdp_sharding_startegy,
        config.model.lora.use_lora,
        config.model.fsdp.use_cpu_offload,
        master_weight_type="fp32",
    )

    if config.model.lora.use_lora:
        transformer.config.lora_rank = config.model.lora.lora_rank
        transformer.config.lora_alpha = config.model.lora.lora_rank
        transformer.config.lora_target_modules = config.model.lora.target_modules
        transformer._no_split_modules = [cls.__name__ for cls in no_split_modules]
        fsdp_kwargs["auto_wrap_policy"] = fsdp_kwargs["auto_wrap_policy"](transformer)

    transformer = FSDP(transformer, **fsdp_kwargs)

    if config.model.ema.use_ema:
        ema_transformer = FSDP(ema_transformer, **fsdp_kwargs)

    if config.model.gradient_checkpointing:
        apply_fsdp_checkpointing(
            transformer, no_split_modules, config.model.selective_checkpointing
        )
        if config.model.ema.use_ema:
            apply_fsdp_checkpointing(
                ema_transformer, no_split_modules, config.model.selective_checkpointing
            )
        logging.info("enable gradient checkpointing")

    transformer.train()
    print_parameters_information(transformer, "WAN", basic_kwargs.rank)

    if not validate_model_parameters(transformer, "Transformer"):
        logging.warning("transformer has invalid parameters!")

    if config.model.ema.use_ema:
        ema_transformer.requires_grad_(False)
        print_parameters_information(ema_transformer, "WAN EMA", basic_kwargs.rank)
        if not validate_model_parameters(ema_transformer, "EMA Transformer"):
            logging.warning("EMA Transformer has invalid parameters!")
    
    feature_layer = config.lrm.feature_layer
    mlp_input_dim = config.lrm.mlp_dim
    loss_type = config.lrm.loss
    
    mlp = MLP(mlp_input_dim)
    if config.model.resume_mlp_path:
        logging.info(f"loading model mlp from {config.model.resume_mlp_path}")
        checkpoint = torch.load(config.model.resume_mlp_path)
        mlp.load_state_dict(checkpoint)
        resume_step = int(config.model.resume_mlp_path.split("_")[-1].split('.')[0])
    mlp.train()
    mlp = mlp.to(device=basic_kwargs.device, dtype=torch.float32)
    
    if not validate_model_parameters(mlp, "MLP"):
        logging.error("MLP has invalid parameters!")
        raise ValueError("MLP initialization failed")

    query_attention_config = getattr(config.lrm, 'query_attention', {})
    num_queries = query_attention_config.get('num_queries', 1)
    num_heads = query_attention_config.get('num_heads', 8)
    dropout = query_attention_config.get('dropout', 0.)
    layer_norm = query_attention_config.get('layer_norm', False)
    return_type = query_attention_config.get('return_type', None)
    product_text = query_attention_config.get('product_text', False)
    text_dim = query_attention_config.get('text_dim', 4096)

    query_attention = QueryAttention(
        feature_dim=mlp_input_dim, 
        num_queries=num_queries, 
        num_heads=num_heads, 
        dropout=dropout,
        return_type=return_type,
        product_text=product_text,
        text_dim=text_dim
    )
    if hasattr(config.model, 'resume_query_attention_path') and config.model.resume_query_attention_path:
        logging.info(f"loading model query_attention from {config.model.resume_query_attention_path}")
        checkpoint = torch.load(config.model.resume_query_attention_path)
        query_attention.load_state_dict(checkpoint)
    query_attention.train()
    query_attention = query_attention.to(device=basic_kwargs.device, dtype=torch.float32)
    
    if not validate_model_parameters(query_attention, "QueryAttention"):
        logging.error("QueryAttention has invalid parameters!")
        raise ValueError("QueryAttention initialization failed")

    criterion = nn.BCELoss()
    criterion = criterion.to(device=basic_kwargs.device)

    model_kwargs = EasyDict(
        {
            "transformer": transformer,
            "ema_transformer": ema_transformer,
            "resume_step": resume_step,
            "feature_layer": feature_layer,
            "loss_type": loss_type,
            "mlp": mlp,
            "query_attention": query_attention,
            "criterion": criterion,
        }
    )
    return model_kwargs

def extra_model_init(config, basic_kwargs):
    base_dir = config.model.base_path

    noise_scheduler = FlowMatchDiscreteScheduler(
        shift=config.extra_model.scheduler.flow_shift
    )
    noise_scheduler.set_timesteps(
        config.extra_model.scheduler.num_train_timesteps, dtype=torch.int64
    )

    if config.train.sanity_check_interval > 0:
        vae = WanVAE(
            vae_pth=os.path.join(base_dir, config.extra_model.vae.name),
            device=basic_kwargs.device,
        )
        print_parameters_information(vae.model, "VAE", basic_kwargs.rank)
    else:
        vae = None

    tokenizer = None
    text_encoder = None
    image_encoder = None

    extra_model_kwargs = EasyDict(
        {
            "noise_scheduler": noise_scheduler,
            "vae": vae,
            "tokenizer": tokenizer,
            "text_encoder": text_encoder,
            "image_encoder": image_encoder,
        }
    )

    logging.info(f"extra model initialized")

    return extra_model_kwargs

def evaluate_model(config, model_kwargs, extra_model_kwargs, basic_kwargs, log_kwargs, writer, step, bucket_timesteps):
    val_dataset = Image2VideoTrainDataset(
        dataset_type="lrm_ce",
        task=config.task,
        meta_file_list=config.dataset.val_meta_file_list,
        uncond_prob=config.dataset.uncond_prob,
        sp_size=config.dataset.sp_size,
        patch_size=config.model.patch_size
    )
    logging.info(f"val dataset length {len(val_dataset)}")
    val_sampler = BlockDistributedSampler(
        val_dataset,
        num_replicas=basic_kwargs.world_size // nccl_info.sp_size,
        rank=nccl_info.group_id,
        shuffle=False,
        seed=config.train.seed,
        drop_last=False,
        batch_size=config.dataset.batch_size
    )
    val_dataloader = DataLoader(
        val_dataset,
        sampler=val_sampler,
        batch_size=config.dataset.batch_size,
        num_workers=0,
        drop_last=False,
        worker_init_fn=set_worker_seed_builder(basic_kwargs.rank), 
        persistent_workers=False
    )

    transformer = model_kwargs.transformer
    MLP = model_kwargs.mlp
    query_attention = model_kwargs.query_attention
    noise_scheduler = extra_model_kwargs.noise_scheduler
    criterion = model_kwargs.criterion
    
    transformer.eval()
    MLP.eval()
    query_attention.eval()
    
    total_loss = 0.0
    num_batches = 0
    
    all_predictions = []
    all_labels = []

    with torch.no_grad():
        dataloader_iter = tqdm(val_dataloader, desc="Evaluating", disable=(basic_kwargs.rank != 0))
        for batch in dataloader_iter:
            (
                latents,
                text_states,
                uncond_text_states,
                image_embeds,
                latents_condition,
                data_from_model,
                text_alignment,
                blur_quality,
                physics_quality,
                human_quality
            ) = batch

            latents = latents.to(basic_kwargs.device, dtype=basic_kwargs.dtype)
            text_states = text_states.to(basic_kwargs.device, dtype=basic_kwargs.dtype)
            
            latents_condition = (
                latents_condition.to(basic_kwargs.device, dtype=basic_kwargs.dtype)
                if "i2v" in config.task or "flf2v" in config.task
                else None
            )
            image_embeds = (
                image_embeds.to(basic_kwargs.device, dtype=basic_kwargs.dtype)
                if "i2v" in config.task or "flf2v" in config.task
                else None
            )
            
            if config.lrm.task == "text_alignment":
                label = text_alignment
            elif config.lrm.task == "blur_quality":
                label = blur_quality
            elif config.lrm.task == "physics_quality":
                label = physics_quality
            elif config.lrm.task == "human_quality":
                label = human_quality
            elif config.lrm.task == "motion_quality":
                label = physics_quality and human_quality
            else:
                label = None
                
            if label is not None:
                label = label.to(basic_kwargs.device, dtype=torch.float32)

            if latents_condition is not None and latents_condition.shape[1] == 16:
                b, _, f, h, w = latents_condition.shape
                mask_lat_size = torch.ones((b, 4, f, h, w), dtype=basic_kwargs.dtype, device=basic_kwargs.device) 
                mask_lat_size[:, :, 1:, ...] = 0.0
                latents_condition = torch.cat([mask_lat_size, latents_condition], dim=1)

            if image_embeds is not None:
                N = image_embeds.shape[1] // 257
                image_embeds = rearrange(image_embeds, "b (n s) d -> (b n) s d", n=N)
            
            if config.dataset.sp_size <= 1:
                latents, latents_condition = crop_tensor(
                    latents,
                    latents_condition,
                    config.dataset.crop_ratio[0],
                    config.dataset.crop_ratio[1],
                    config.dataset.crop_type,
                    crop_time_ratio=config.dataset.crop_ratio[2],
                )

            _, _, latents_t, latents_h, latents_w = latents.shape
            max_sequence_length = (
                latents_t
                * latents_h
                * latents_w
                // (config.model.patch_size[1] * config.model.patch_size[2])
            )

            seed_g = torch.Generator(device=latents.device)
            seed_g.manual_seed(config.eval.seed)
            bsz = latents.shape[0]
            noise = torch.randn(latents.shape, device=latents.device, dtype=latents.dtype, generator=seed_g)
            
            t = bucket_timesteps[0]
            t_sample = random.choice(bucket_timesteps)
            timestep = torch.full((1,), t_sample, device=latents.device, dtype=torch.int64)
            sigma = noise_scheduler.get_train_sigma(
                timestep, 
                n_dim=latents.ndim, 
                device=latents.device, 
                dtype=latents.dtype
            )

            if config.dataset.sp_size > 1:
                if "i2v" in config.task or "flf2v" in config.task: 
                    broadcast(latents_condition)
                    broadcast(image_embeds)
                broadcast(sigma)
                broadcast(noise)
                broadcast(timestep)
                broadcast(latents)
                broadcast(text_states)
            
            noisy_latents = noise_scheduler.add_noise(latents, noise, sigma)

            cond_kwargs = {
                "x": batch2list(noisy_latents),
                "t": timestep,
                "context": batch2list(text_states),
                "seq_len": max_sequence_length,
                "clip_fea": image_embeds,
                "y": (
                    batch2list(latents_condition)
                    if "i2v" in config.task or "flf2v" in config.task
                    else None
                ),
                "output_features": True,
                "selected_layers": model_kwargs.feature_layer,
            }
            
            with torch.autocast("cuda", dtype=basic_kwargs.dtype):
                model_pred = transformer(**cond_kwargs)
                model_pred = list2batch(model_pred)
                
                if config.dataset.sp_size > 1:
                    if len(model_pred.shape) == 4:
                        if config.lrm.pool == 'q_attn':
                            model_pred_final = query_attention(model_pred)
                        elif config.lrm.pool == 'max':
                            model_pred_pooled, _ = model_pred.max(dim=2)
                            model_pred_final, _ = model_pred_pooled.max(dim=0)
                        else:
                            model_pred_pooled = model_pred.mean(dim=2)
                            model_pred_final = model_pred_pooled.mean(dim=0)
                    else:
                        original_batch_size = bsz
                        model_pred_flat = model_pred.view(original_batch_size, -1)
                        model_pred_final = model_pred_flat.mean(dim=1, keepdim=True)
                else:
                    if len(model_pred.shape) == 3:
                        if config.lrm.pool == 'q_attn':
                            model_pred_final = query_attention(model_pred)
                        elif config.lrm.pool == 'max':
                            model_pred_final, _ = model_pred.max(dim=1)
                        else:
                            model_pred_final = model_pred.mean(dim=1)
                    else:
                        batch_size = model_pred.shape[0]
                        model_pred_final = model_pred.view(batch_size, -1).mean(dim=1, keepdim=True)

                outputs = forward_mlp(MLP, model_pred_final)

            if label is not None:
                outputs_squeezed = outputs.squeeze()
                label_squeezed = label.squeeze()
                
                outputs_for_loss = outputs_squeezed.float()
                label_for_loss = label_squeezed.float()
                
                if config.dataset.sp_size > 1:
                    broadcast(label_for_loss)
                
                loss = criterion(outputs_for_loss, label_for_loss)
                total_loss += loss.item()
                
                predictions = (outputs_squeezed > 0.5).long()
                
                if predictions.ndim == 0:
                    predictions = predictions.unsqueeze(0)
                if label_for_loss.ndim == 0:
                    label_for_loss = label_for_loss.unsqueeze(0)
                
                pred_np = predictions.cpu().numpy()
                label_np = label_for_loss.long().cpu().numpy()
                
                if pred_np.ndim > 1:
                    pred_np = pred_np.flatten()
                if label_np.ndim > 1:
                    label_np = label_np.flatten()
                
                all_predictions.append(pred_np)
                all_labels.append(label_np)
            
            num_batches += 1

    total_loss_tensor = torch.tensor(total_loss, device=basic_kwargs.device, dtype=torch.float32)
    num_batches_tensor = torch.tensor(num_batches, device=basic_kwargs.device, dtype=torch.float32)
    all_predictions = torch.tensor(np.stack(all_predictions), device=basic_kwargs.device, dtype=torch.float32)
    all_labels = torch.tensor(np.stack(all_labels), device=basic_kwargs.device, dtype=torch.float32)
    
    dist.all_reduce(total_loss_tensor, op=dist.ReduceOp.SUM)
    dist.all_reduce(num_batches_tensor, op=dist.ReduceOp.SUM)
    all_predictions_list = [torch.zeros_like(all_predictions) for _ in range(basic_kwargs.world_size)]
    dist.all_gather(all_predictions_list, all_predictions)
    all_labels_list = [torch.zeros_like(all_labels) for _ in range(basic_kwargs.world_size)]
    dist.all_gather(all_labels_list, all_labels)

    avg_loss = total_loss_tensor.item() / num_batches_tensor.item() if num_batches_tensor.item() > 0 else 0
    all_preds = torch.concat(all_predictions_list).cpu().numpy()
    all_labs = torch.concat(all_labels_list).cpu().numpy()

    if len(all_predictions) > 0 and len(all_labels) > 0:
        try:
            accuracy = accuracy_score(all_labs, all_preds)
            precision = precision_score(all_labs, all_preds, zero_division=0)
            recall = recall_score(all_labs, all_preds, zero_division=0)
            f1 = f1_score(all_labs, all_preds, zero_division=0)
        except Exception as e:
            logging.error(f"Error: {e}")
            accuracy = precision = recall = f1 = 0.0
    else:
        accuracy = precision = recall = f1 = 0.0
    
    if basic_kwargs.rank == 0:
        logging.info(f"✨ Evaluation - Accuracy: {accuracy:.4f}, Precision: {precision:.4f}, Recall: {recall:.4f}, F1: {f1:.4f}, Avg Loss: {avg_loss:.4f}")

        log_info = (
            f"│ Rank {basic_kwargs.rank:02d} │ Workers: {basic_kwargs.world_size} │"
            f"CKPT Step: {step} │"
            f"Timstep: {t} │"
            f"VAL Loss: {avg_loss:.4f} │"
            f"VAL Acc:{accuracy:.4f} │"
            f"VAL Prec:{precision:.4f} │"
            f"VAL Recall:{recall:.4f} │"
            f"VAL F1:{f1:.4f} │"
        )

        with open(basic_kwargs.log_path, "a", encoding="utf-8") as f:
            f.write(log_info + "\n")
    
        if writer is not None:
            writer.add_scalar(f'val/loss_{t}', avg_loss, step)
            writer.add_scalar(f'val/acc_{t}', accuracy, step)
            writer.add_scalar(f'val/precision_{t}', precision, step)
            writer.add_scalar(f'val/recall_{t}', recall, step)
            writer.add_scalar(f'val/f1_{t}', f1, step)

    transformer.eval()
    MLP.eval()
     
    return accuracy, avg_loss, precision, recall, f1
    
def main(config):
    config, basic_kwargs = basic_init(config)

    model_kwargs = model_init(config, basic_kwargs)
    extra_model_kwargs = extra_model_init(config, basic_kwargs)

    dist.barrier()
    free_memory()

    writer = SummaryWriter(config.save.tensorboard_dir) if basic_kwargs.rank == 0 else None
    total_batch_size = (
        config.dataset.batch_size
        * (basic_kwargs.world_size // nccl_info.sp_size)
        * config.train.gradient_accumulation_steps
    )
    logging.info("***** Running evaluation *****")

    bucket_intervals = [(0, 200), (201, 400), (401, 600), (601,800), (801, 1000)]
    for i, (start_bound, end_bound) in enumerate(bucket_intervals):
        bucket_timesteps = []
        for t in extra_model_kwargs.noise_scheduler.timesteps:
            if t >= start_bound and t <= end_bound:
                bucket_timesteps.append(t)
        try:
            random.seed(config.eval.seed)
            evaluate_model(config, model_kwargs, extra_model_kwargs, basic_kwargs, EasyDict(), writer, model_kwargs.resume_step, bucket_timesteps)
        except:
            continue

    if basic_kwargs.rank == 0 and writer is not None:
        writer.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config_path",
        type=str,
        required=True,
        default="scripts/train/train_wanx.yaml",
    )
    args = parser.parse_args()

    main(OmegaConf.load(args.config_path))