nodes.py

import os
import math
import random
import logging
import datetime

from pathlib import Path
from tqdm.auto import tqdm
from einops import rearrange
from omegaconf import OmegaConf

import torch
import torch.nn.functional as F

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler
from diffusers.optimization import get_scheduler

from transformers import CLIPTextModel, CLIPTokenizer

from .animatediff.models.unet import UNet3DConditionModel
from .animatediff.pipelines.pipeline_animation import AnimationPipeline
from .animatediff.utils.util import save_videos_grid, load_weights
from .animatediff.utils.lora_handler import LoraHandler
from .animatediff.utils.lora import extract_lora_child_module

from .motion_lora import MotionLoraInfo, MotionLoraList

from lion_pytorch import Lion
import comfy.model_management
import comfy.utils
import folder_paths

try:
    import xformers
    import xformers.ops

    XFORMERS_IS_AVAILABLE = True
except:
    XFORMERS_IS_AVAILABLE = False

script_directory = os.path.dirname(os.path.abspath(__file__))
folder_paths.add_model_folder_path("animatediff_models", str(Path(__file__).parent.parent / "models"))
folder_paths.add_model_folder_path("animatediff_models", str(Path(folder_paths.models_dir) / "animatediff_models"))

def create_save_paths(output_dir: str):
    directories = [
        output_dir,
        f"{output_dir}/samples",
        f"{output_dir}/sanity_check",
    ]
    for directory in directories:
        os.makedirs(directory, exist_ok=True)

def resize_and_pad_images(images, output_size):
    images = images.permute(0, 3, 1, 2)
    
    # Calculate aspect ratio of the output size
    aspect_ratio = output_size[1] / output_size[0]
    
    # Calculate resizing size that maintains aspect ratio
    batch_size, channels, height, width = images.size()
    if width / height > aspect_ratio:
        # If image is wider than desired aspect ratio, fit to width
        resize_width = output_size[1]
        resize_height = round(resize_width / width * height)
    else:
        # If image is taller than desired aspect ratio, fit to height
        resize_height = output_size[0]
        resize_width = round(resize_height / height * width)
    
    # Resize the images while maintaining the aspect ratio
    resized_images = F.interpolate(images, size=(resize_height, resize_width), mode='bilinear', align_corners=False)

    # Calculate the padding required to make the images the same size
    pad_width = max(0, output_size[1] - resized_images.size(3))
    pad_height = max(0, output_size[0] - resized_images.size(2))
    pad_left = pad_width // 2
    pad_right = pad_width - pad_left
    pad_top = pad_height // 2
    pad_bottom = pad_height - pad_top

    # Pad the images to the desired size with black bars
    padded_images = F.pad(resized_images, (pad_left, pad_right, pad_top, pad_bottom), value=0)

    return padded_images

def do_sanity_check(
    sanity_check,  
    output_dir,
    text_prompt
):
    if isinstance(sanity_check, list):
        resized_images = []
        for image in sanity_check:
            image = resize_and_pad_images(image, (512, 512))
            resized_images.append(image)
        sanity_check = torch.cat(resized_images, dim=0)
        sanity_check = sanity_check.unsqueeze(0)
        
    else:
        sanity_check = sanity_check.permute(0, 3, 1, 2).unsqueeze(0)#B,H,W,C to B,F,C,H,W
    sanity_check = sanity_check * 2.0 - 1.0 #normalize to the expected range (-1, 1)
    

    sanity_check, texts = sanity_check.cpu(), text_prompt  
    sanity_check = rearrange(sanity_check, "b f c h w -> b c f h w")
    for idx, (pixel_value, text) in enumerate(zip(sanity_check, texts)):
        pixel_value = pixel_value[None, ...]
        text = text
        save_name = f"{'-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'-{idx}'}.mp4"
        save_videos_grid(pixel_value, f"{output_dir}/sanity_check/{save_name}", rescale=False)
    return(sanity_check)

def sample_noise(latents, noise_strength, use_offset_noise=False):
    b, c, f, *_ = latents.shape
    noise_latents = torch.randn_like(latents, device=latents.device)

    if use_offset_noise:
        offset_noise = torch.randn(b, c, f, 1, 1, device=latents.device)
        noise_latents = noise_latents + noise_strength * offset_noise

    return noise_latents

def param_optim(model, condition, extra_params=None, is_lora=False, negation=None):
    extra_params = extra_params if len(extra_params.keys()) > 0 else None
    return {
        "model": model,
        "condition": condition,
        'extra_params': extra_params,
        'is_lora': is_lora,
        "negation": negation
    }

def create_optim_params(name='param', params=None, lr=5e-6, extra_params=None):
    params = {
        "name": name,
        "params": params,
        "lr": lr
    }
    if extra_params is not None:
        for k, v in extra_params.items():
            params[k] = v

    return params

def create_optimizer_params(model_list, lr):
    import itertools
    optimizer_params = []

    for optim in model_list:
        model, condition, extra_params, is_lora, negation = optim.values()
        # Check if we are doing LoRA training.
        if is_lora and condition and isinstance(model, list):
            params = create_optim_params(
                params=itertools.chain(*model),
                extra_params=extra_params
            )
            optimizer_params.append(params)
            continue

        if is_lora and condition and not isinstance(model, list):
            for n, p in model.named_parameters():
                if 'lora' in n:
                    params = create_optim_params(n, p, lr, extra_params)
                    optimizer_params.append(params)
            continue

        # If this is true, we can train it.
        if condition:
            for n, p in model.named_parameters():
                should_negate = 'lora' in n and not is_lora
                if should_negate: continue
            
                params = create_optim_params(n, p, lr, extra_params)
                optimizer_params.append(params)

    return optimizer_params

def scale_loras(lora_list: list, scale: float, step=None): 
    # Assumed enumerator
    if step is not None:
        process_list = range(0, len(lora_list), 1)
    else:
        process_list = lora_list

    for lora_i in process_list:
        if step is not None:
            lora_list[lora_i].scale = scale
        else:
            lora_i.scale = scale

def tensor_to_vae_latent(t, vae):
    video_length = t.shape[1]
    t = rearrange(t, "b f c h w -> (b f) c h w") 
    latents = vae.encode(t).latent_dist.sample()
    latents = rearrange(latents, "(b f) c h w -> b c f h w", f=video_length)
    latents = latents * 0.18215

    return latents

def get_spatial_latents(
        pixel_values: torch.Tensor, 
        noisy_latents:torch.Tensor, 
        target: torch.Tensor,
    ):
    ran_idx = torch.randint(0, pixel_values.shape[2], (1,)).item()

    noisy_latents_input = None
    target_spatial = None
    
    noisy_latents_input = noisy_latents[:, :, ran_idx, :, :]
    target_spatial = target[:, :, ran_idx, :, :]

    return noisy_latents_input, target_spatial

def create_ad_temporal_loss(
        model_pred: torch.Tensor, 
        loss_temporal: torch.Tensor, 
        target: torch.Tensor
    ):

    beta = 1
    alpha = (beta ** 2 + 1) ** 0.5

    ran_idx = torch.randint(0, model_pred.shape[2], (1,)).item()

    model_pred_decent = alpha * model_pred - beta * model_pred[:, :, ran_idx, :, :].unsqueeze(2)
    target_decent = alpha * target - beta * target[:, :, ran_idx, :, :].unsqueeze(2)

    loss_ad_temporal = F.mse_loss(model_pred_decent.float(), target_decent.float(), reduction="mean")
    loss_temporal = loss_temporal + loss_ad_temporal

    return loss_temporal

class ADMD_InitializeTraining:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "pipeline": ("PIPELINE", ),
            "lora_name": ("STRING", {"multiline": False, "default": "motiondirectorlora",}),
            "images": ("IMAGE", ),
            "prompt": ("STRING", {"multiline": True, "default": "",}),
            "max_train_steps": ("INT", {"default": 300, "min": 0, "max": 100000, "step": 1}),
            "learning_rate": ("FLOAT", {"default": 5e-4, "min": 0, "max": 10000, "step": 0.00001}),
            "learning_rate_spatial": ("FLOAT", {"default": 1e-4, "min": 0, "max": 10000, "step": 0.00001}),    
            "lora_rank": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 8}),
            "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
            
            "optimization_method": (
            [   
                'Lion',
                'AdamW',
            ], {
               "default": 'Lion'
            }),
            "include_resnet": ("BOOLEAN", {"default": False}),
            },
            
            }
    
    RETURN_TYPES = ("IMAGE", "ADMDPIPELINE", "LORAINFO")
    RETURN_NAMES =("sanitycheck", "admd_pipeline", "lora_info",)
    FUNCTION = "process"

    CATEGORY = "AD_MotionDirector"

    def process(self, pipeline, images, prompt,  
                lora_name, learning_rate, learning_rate_spatial, 
                lora_rank, seed, optimization_method, max_train_steps, include_resnet):
        with torch.inference_mode(False):
                      
            validation_pipeline = pipeline["validation_pipeline"]
            train_noise_scheduler = pipeline["train_noise_scheduler"]
            train_noise_scheduler_spatial = pipeline["train_noise_scheduler_spatial"]

            unet = pipeline["unet"]
            text_encoder = pipeline["text_encoder"]
            vae = pipeline["vae"]
            tokenizer = pipeline["tokenizer"]

            torch.manual_seed(seed)

            text_prompt = []
            text_prompt.append(prompt)
  
            scale_lr = False
            lr_warmup_steps = 0
            lr_scheduler = "constant"

            train_batch_size = 1
            adam_beta1 = 0.9
            adam_beta2 = 0.999
            adam_weight_decay = 1e-2
            gradient_accumulation_steps = 1
            
            is_debug = False

            lora_unet_dropout = 0.1
            if include_resnet:
                target_spatial_modules = ["Transformer3DModel", "ResnetBlock2D"]
            else: 
                target_spatial_modules = ["Transformer3DModel"]

            target_temporal_modules = ["TemporalTransformerBlock"]          
  
            name = lora_name
            date_calendar = datetime.datetime.now().strftime("%Y-%m-%d")
            date_time = datetime.datetime.now().strftime("%H-%M-%S")
            folder_name = "debug" if is_debug else name + date_time
            
            output_dir = os.path.join(script_directory, "outputs", date_calendar, folder_name)

            if is_debug and os.path.exists(output_dir):
                os.system(f"rm -rf {output_dir}")

            # 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,
            )

            # set paths
            spatial_lora_path = os.path.join(folder_paths.models_dir,"loras", "trained_spatial", date_calendar, date_time, lora_name) 
            temporal_lora_path = os.path.join(folder_paths.models_dir,"animatediff_motion_lora", date_calendar, date_time, lora_name)
            temporal_lora_base_path = os.path.join(date_calendar, date_time, lora_name)
            
            lora_info = {
                "lora_name": lora_name,
                "lora_rank": lora_rank,
                "spatial_lora_path": spatial_lora_path,
                "temporal_lora_path": temporal_lora_path,
                "temporal_lora_base_path": temporal_lora_base_path
            }

            if optimization_method == "AdamW":   
                print("Using AdamW optimizer for training") 
                optimizer = torch.optim.AdamW
            else:
                print("Using Lion optimizer for training")
                optimizer = Lion
                learning_rate, learning_rate_spatial = map(lambda lr: lr / 10, (learning_rate, learning_rate_spatial))
                adam_weight_decay *= 10

            if scale_lr:
                learning_rate = (learning_rate * gradient_accumulation_steps * train_batch_size)

            # Temporal LoRA
            
            # one temporal lora
            lora_manager_temporal = LoraHandler(use_unet_lora=True, unet_replace_modules=target_temporal_modules)
            
            unet_lora_params_temporal, unet_negation_temporal = lora_manager_temporal.add_lora_to_model(
                True, unet, lora_manager_temporal.unet_replace_modules, 0,
                temporal_lora_path, r=lora_rank)

            optimizer_temporal = optimizer(
                create_optimizer_params([param_optim(unet_lora_params_temporal, True, is_lora=True,
                                                    extra_params={**{"lr": learning_rate}}
                                                    )], learning_rate),
                lr=learning_rate,
                betas=(adam_beta1, adam_beta2),
                weight_decay=adam_weight_decay
            )
        
            lr_scheduler_temporal = get_scheduler(
                lr_scheduler,
                optimizer=optimizer_temporal,
                num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
                num_training_steps=max_train_steps * gradient_accumulation_steps,
            )

            # Spatial LoRAs

            unet_lora_params_spatial_list = []
            optimizer_spatial_list = []
            lr_scheduler_spatial_list = []

            lora_manager_spatial = LoraHandler(use_unet_lora=True, unet_replace_modules=target_spatial_modules)
            unet_lora_params_spatial, unet_negation_spatial = lora_manager_spatial.add_lora_to_model(
                True, unet, lora_manager_spatial.unet_replace_modules, lora_unet_dropout,
                spatial_lora_path, r=lora_rank)

            unet_lora_params_spatial_list.append(unet_lora_params_spatial)

            optimizer_spatial = optimizer(
                create_optimizer_params([param_optim(unet_lora_params_spatial, True, is_lora=True,
                                                    extra_params={**{"lr": learning_rate_spatial}}
                                                    )], learning_rate_spatial),
                lr=learning_rate_spatial,
                betas=(adam_beta1, adam_beta2),
                weight_decay=adam_weight_decay
            )

            optimizer_spatial_list.append(optimizer_spatial)

            # Scheduler
            lr_scheduler_spatial = get_scheduler(
                lr_scheduler,
                optimizer=optimizer_spatial,
                num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
                num_training_steps=max_train_steps * gradient_accumulation_steps,
            )
            lr_scheduler_spatial_list.append(lr_scheduler_spatial)
           
            # Support mixed-precision training
            if 'scaler' not in globals():
                scaler = torch.cuda.amp.GradScaler()
                print("initialize scaler")
            else:
                scaler.reset()
                print("reset scaler")

        admd_pipeline = {
            "optimizer_temporal": optimizer_temporal,
            "optimizer_spatial_list": optimizer_spatial_list,
            "lr_scheduler_spatial_list": lr_scheduler_spatial_list,
            "lr_scheduler_temporal": lr_scheduler_temporal,
            "text_prompt": text_prompt,
            "unet": unet,
            "text_encoder": text_encoder,
            "vae": vae,
            "tokenizer": tokenizer,
            "pixel_values": images,
            "train_noise_scheduler": train_noise_scheduler,
            "train_noise_scheduler_spatial": train_noise_scheduler_spatial,
            "validation_pipeline": validation_pipeline,
            "global_step": 0,
            "max_train_steps": max_train_steps,
            "scaler": scaler,
            "include_resnet": include_resnet,
            "seed": seed
        }
        #Data batch sanity check
        
        sanitycheck = do_sanity_check(
                images,  
                output_dir=output_dir, 
                text_prompt=text_prompt
        )
 
        sanitycheck = sanitycheck.view(*sanitycheck.shape[1:])
        sanitycheck = sanitycheck.permute(1, 2, 3, 0).cpu()
        sanitycheck = (sanitycheck + 1.0) / 2.0
        return (sanitycheck, admd_pipeline, lora_info,)

class ADMD_DiffusersLoader:
    @classmethod
    def IS_CHANGED(s):
        return ""
    @classmethod
    def INPUT_TYPES(cls):
        paths = []
        for search_path in folder_paths.get_folder_paths("diffusers"):
            if os.path.exists(search_path):
                for root, subdir, files in os.walk(search_path, followlinks=True):
                    if "model_index.json" in files:
                        paths.append(os.path.relpath(root, start=search_path))

        return {"required":
                {
                "additional_models": ("ADDITIONAL_MODELS", ),  
                "download_default": ("BOOLEAN", {"default": False},),
                "scheduler": (
            [   
                'DDIMScheduler',
                'DDPMScheduler',
            ], {
               "default": 'DDIMScheduler'
            }),
            "use_xformers": ("BOOLEAN", {"default": False}),
                },
                "optional": {
                 "model": (paths,),
                }
                
            }
    RETURN_TYPES = ("PIPELINE",)

    FUNCTION = "load_checkpoint"

    CATEGORY = "AD_MotionDirector"

    def load_checkpoint(self, download_default, scheduler, use_xformers, additional_models, model=""):
        with torch.inference_mode(False):
            device = comfy.model_management.get_torch_device()
            target_path = os.path.join(folder_paths.models_dir,'diffusers', "stable-diffusion-v1-5")      
            if download_default and model != os.path.exists(target_path):
                from huggingface_hub import snapshot_download
                download_to = os.path.join(folder_paths.models_dir,'diffusers')
                snapshot_download(repo_id="runwayml/stable-diffusion-v1-5", ignore_patterns=["*.safetensors","*.ckpt", "*.pt", "*.png", "*non_ema*", "*safety_checker*", "*fp16*"], 
                                    local_dir=f"{download_to}/stable-diffusion-v1-5", local_dir_use_symlinks=False)   
                model_path = "stable-diffusion-v1-5"
            else:
                model_path = model
            
            for search_path in folder_paths.get_folder_paths("diffusers"):
                if os.path.exists(search_path):
                    path = os.path.join(search_path, model_path)
                    if os.path.exists(path):
                        model_path = path
                        break      
            
            config = OmegaConf.load(os.path.join(script_directory, f"configs/training/motion_director/training.yaml"))
       
            vae          = AutoencoderKL.from_pretrained(model_path, subfolder="vae")
            tokenizer    = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
            text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="text_encoder")
         
            unet_additional_kwargs = config.unet_additional_kwargs
            unet = UNet3DConditionModel.from_pretrained_2d(
                model_path, subfolder="unet", 
                unet_additional_kwargs=unet_additional_kwargs
            )
            
            # Load scheduler, tokenizer and models.
            noise_scheduler_kwargs = {
                'num_train_timesteps': 1000,
                'beta_start':    0.00085,
                'beta_end':      0.012,
                'beta_schedule': "linear",
                'clip_sample':   False,
                'steps_offset': 1
            }
            # Determine the scheduler class based on the scheduler variable
            SchedulerClass = DDPMScheduler if scheduler == "DDPMScheduler" else DDIMScheduler
            print(f"using {SchedulerClass.__name__} for training")

            # Set the beta_schedule and create the default noise scheduler
            noise_scheduler_kwargs['beta_schedule'] = 'linear'
            noise_scheduler = SchedulerClass(**noise_scheduler_kwargs)

            # Set the beta_schedule for the spatial noise scheduler and create it
            noise_scheduler_kwargs['beta_schedule'] = 'scaled_linear'
            train_noise_scheduler_spatial = SchedulerClass(**noise_scheduler_kwargs)

            # Reset the beta_schedule for the linear noise scheduler and create it
            noise_scheduler_kwargs['beta_schedule'] = 'linear'
            train_noise_scheduler = SchedulerClass(**noise_scheduler_kwargs)
            
            # Freeze all models for LoRA training
            unet.requires_grad_(False)
            vae.requires_grad_(False)
            text_encoder.requires_grad_(False)

            #xformers
            if XFORMERS_IS_AVAILABLE:
                if use_xformers:
                    unet.enable_xformers_memory_efficient_attention()
                else:
                    unet.disable_xformers_memory_efficient_attention()

            # Enable gradient checkpointing
            unet.enable_gradient_checkpointing()

            # Validation pipeline
            validation_pipeline = AnimationPipeline(
                unet=unet, vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=noise_scheduler,
            ).to(device)

            motion_model, domain_adapter_path = additional_models 

            validation_pipeline = load_weights(
                validation_pipeline, 
                motion_model=motion_model,
                adapter_lora_path=domain_adapter_path, 
                dreambooth_model_path=""
            )

            validation_pipeline.enable_vae_slicing()         

            pipeline = {
                'validation_pipeline': validation_pipeline,
                'train_noise_scheduler': train_noise_scheduler,
                'train_noise_scheduler_spatial': train_noise_scheduler_spatial,
                'unet': unet,
                'vae': vae,
                'text_encoder': text_encoder,
                'tokenizer': tokenizer
            }

            return (pipeline,)

class ADMD_CheckpointLoader:
    @classmethod
    def IS_CHANGED(s):
        return ""
    @classmethod
    def INPUT_TYPES(cls):

        return {"required":
                {
                "additional_models": ("ADDITIONAL_MODELS", ),
                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),  
                "scheduler": (
            [   
                'DDIMScheduler',
                'DDPMScheduler',
            ], {
               "default": 'DDIMScheduler'
            }),
            "use_xformers": ("BOOLEAN", {"default": False}),
                },                
            }
    RETURN_TYPES = ("PIPELINE",)

    FUNCTION = "load_checkpoint"

    CATEGORY = "AD_MotionDirector"

    def load_checkpoint(self, scheduler, use_xformers, additional_models, ckpt_name):
        with torch.inference_mode(False):            
            model_path = folder_paths.get_full_path("checkpoints", ckpt_name)
            original_config = OmegaConf.load(os.path.join(script_directory, f"configs/v1-inference.yaml"))
            ad_unet_config = OmegaConf.load(os.path.join(script_directory, f"configs/ad_unet_config.yaml"))

            from diffusers.loaders.single_file_utils import (convert_ldm_vae_checkpoint, convert_ldm_unet_checkpoint, create_text_encoder_from_ldm_clip_checkpoint, create_vae_diffusers_config, create_unet_diffusers_config)
            from safetensors import safe_open

            if model_path.endswith(".safetensors"):
                dreambooth_state_dict = {}
                with safe_open(model_path, framework="pt", device="cpu") as f:
                    for key in f.keys():
                        dreambooth_state_dict[key] = f.get_tensor(key)
            elif model_path.endswith(".ckpt"):
                dreambooth_state_dict = torch.load(model_path, map_location="cpu")
                while "state_dict" in dreambooth_state_dict:
                    dreambooth_state_dict = dreambooth_state_dict["state_dict"]

            tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
            text_encoder = create_text_encoder_from_ldm_clip_checkpoint("openai/clip-vit-large-patch14",dreambooth_state_dict)

            noise_scheduler_kwargs = {
                'num_train_timesteps': 1000,
                'beta_start':    0.00085,
                'beta_end':      0.012,
                'beta_schedule': "linear",
                'clip_sample':   False,
                'steps_offset': 1
            }
            #Determine the scheduler class based on the scheduler variable
            SchedulerClass = DDPMScheduler if scheduler == "DDPMScheduler" else DDIMScheduler
            print(f"using {SchedulerClass.__name__} for training")

            # Set the beta_schedule and create the default noise scheduler
            noise_scheduler_kwargs['beta_schedule'] = 'linear'
            noise_scheduler = SchedulerClass(**noise_scheduler_kwargs)

            # Set the beta_schedule for the spatial noise scheduler and create it
            noise_scheduler_kwargs['beta_schedule'] = 'scaled_linear'
            train_noise_scheduler_spatial = SchedulerClass(**noise_scheduler_kwargs)

            # Reset the beta_schedule for the linear noise scheduler and create it
            noise_scheduler_kwargs['beta_schedule'] = 'linear'
            train_noise_scheduler = SchedulerClass(**noise_scheduler_kwargs)

            # 1. vae
            converted_vae_config = create_vae_diffusers_config(original_config, image_size=512)
            converted_vae = convert_ldm_vae_checkpoint(dreambooth_state_dict, converted_vae_config)
            vae = AutoencoderKL(**converted_vae_config)
            vae.load_state_dict(converted_vae, strict=False)

            # 2. unet
            converted_unet_config = create_unet_diffusers_config(original_config, image_size=512)
            converted_unet = convert_ldm_unet_checkpoint(dreambooth_state_dict, converted_unet_config)
            unet = UNet3DConditionModel(**ad_unet_config)
            unet.load_state_dict(converted_unet, strict=False)

            del dreambooth_state_dict, converted_unet, converted_vae

            # Validation pipeline
            validation_pipeline = AnimationPipeline(
                unet=unet, vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=noise_scheduler,
            )
            # Freeze all models for LoRA training
            unet.requires_grad_(False)
            vae.requires_grad_(False)
            text_encoder.requires_grad_(False)

            #xformers
            if XFORMERS_IS_AVAILABLE:
                if use_xformers:
                    unet.enable_xformers_memory_efficient_attention()
                else:
                    unet.disable_xformers_memory_efficient_attention()

            # Enable gradient checkpointing
            unet.enable_gradient_checkpointing()

            motion_model, domain_adapter_path = additional_models 

            validation_pipeline = load_weights(
                validation_pipeline, 
                motion_model=motion_model,
                adapter_lora_path=domain_adapter_path, 
                dreambooth_model_path=""
            )

            validation_pipeline.enable_vae_slicing()         

            pipeline = {
                'validation_pipeline': validation_pipeline,
                'train_noise_scheduler': train_noise_scheduler,
                'train_noise_scheduler_spatial': train_noise_scheduler_spatial,
                'unet': unet,
                'vae': vae,
                'text_encoder': text_encoder,
                'tokenizer': tokenizer
            }

            return (pipeline,)
        
class ADMD_ComfyModelLoader:
    @classmethod
    def IS_CHANGED(s):
        return ""
    @classmethod
    def INPUT_TYPES(cls):

        return {"required":
                {
                "model" :("MODEL",),
                "clip": ("CLIP",),
                "vae": ("VAE",),
                "motion_model":("MOTION_MODEL_ADE",),
                "scheduler": (
            [   
                'DDIMScheduler',
                'DDPMScheduler',
            ], {
               "default": 'DDIMScheduler'
            }),
            "use_xformers": ("BOOLEAN", {"default": False}),
                },                
            }
    RETURN_TYPES = ("PIPELINE",)

    FUNCTION = "load_checkpoint"

    CATEGORY = "AD_MotionDirector"

    def load_checkpoint(self, model, clip, vae, scheduler, use_xformers, motion_model):
        with torch.inference_mode(False):
        
            pbar = comfy.utils.ProgressBar(4)
            original_config = OmegaConf.load(os.path.join(script_directory, f"configs/v1-inference.yaml"))
            ad_unet_config = OmegaConf.load(os.path.join(script_directory, f"configs/ad_unet_config.yaml"))

            from diffusers.loaders.single_file_utils import (convert_ldm_vae_checkpoint, convert_ldm_unet_checkpoint, create_text_encoder_from_ldm_clip_checkpoint, create_vae_diffusers_config, create_unet_diffusers_config)

            tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")

            clip_sd = None
            load_models = [model]
            load_models.append(clip.load_model())
            clip_sd = clip.get_sd()
            
            comfy.model_management.load_models_gpu(load_models)
            sd = model.model.state_dict_for_saving(clip_sd, vae.get_sd(), None)
     
            text_encoder = create_text_encoder_from_ldm_clip_checkpoint("openai/clip-vit-large-patch14",sd)
            pbar.update(1)
            noise_scheduler_kwargs = {
                'num_train_timesteps': 1000,
                'beta_start':    0.00085,
                'beta_end':      0.012,
                'beta_schedule': "linear",
                'clip_sample':   False,
                'steps_offset': 1
            }
            #Determine the scheduler class based on the scheduler variable
            SchedulerClass = DDPMScheduler if scheduler == "DDPMScheduler" else DDIMScheduler
            print(f"using {SchedulerClass.__name__} for training")

            # Set the beta_schedule and create the default noise scheduler
            noise_scheduler_kwargs['beta_schedule'] = 'linear'
            noise_scheduler = SchedulerClass(**noise_scheduler_kwargs)

            # Set the beta_schedule for the spatial noise scheduler and create it
            noise_scheduler_kwargs['beta_schedule'] = 'scaled_linear'
            train_noise_scheduler_spatial = SchedulerClass(**noise_scheduler_kwargs)

            # Reset the beta_schedule for the linear noise scheduler and create it
            noise_scheduler_kwargs['beta_schedule'] = 'linear'
            train_noise_scheduler = SchedulerClass(**noise_scheduler_kwargs)

            # 1. vae
            converted_vae_config = create_vae_diffusers_config(original_config, image_size=512)
            converted_vae = convert_ldm_vae_checkpoint(sd, converted_vae_config)
            vae = AutoencoderKL(**converted_vae_config)
            vae.load_state_dict(converted_vae, strict=False)
            pbar.update(1)
            # 2. unet
            converted_unet_config = create_unet_diffusers_config(original_config, image_size=512)
            converted_unet = convert_ldm_unet_checkpoint(sd, converted_unet_config)
            unet = UNet3DConditionModel(**ad_unet_config)
            unet.load_state_dict(converted_unet, strict=False)
            
            del sd, converted_unet, converted_vae

            pbar.update(1)

            # Validation pipeline
            validation_pipeline = AnimationPipeline(
                unet=unet, vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, scheduler=noise_scheduler,
            )
            # Freeze all models for LoRA training
            unet.requires_grad_(False)
            vae.requires_grad_(False)
            text_encoder.requires_grad_(False)

            #xformers
            if XFORMERS_IS_AVAILABLE:
                if use_xformers:
                    unet.enable_xformers_memory_efficient_attention()
                else:
                    unet.disable_xformers_memory_efficient_attention()

            # Enable gradient checkpointing
            unet.enable_gradient_checkpointing()

            domain_adapter_path = ""

            validation_pipeline = load_weights(
                validation_pipeline, 
                motion_model=motion_model,
                adapter_lora_path=domain_adapter_path, 
                dreambooth_model_path=""
            )
            pbar.update(1)
            validation_pipeline.enable_vae_slicing()         

            pipeline = {
                'validation_pipeline': validation_pipeline,
                'train_noise_scheduler': train_noise_scheduler,
                'train_noise_scheduler_spatial': train_noise_scheduler_spatial,
                'unet': unet,
                'vae': vae,
                'text_encoder': text_encoder,
                'tokenizer': tokenizer
            }

            return (pipeline,)
        
class ADMD_AdditionalModelSelect:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": { 
                "motion_module": (folder_paths.get_filename_list("animatediff_models"),),
                "use_adapter_lora": ("BOOLEAN", {"default": True}),                                       
            },
            "optional": {
                "optional_adapter_lora": (folder_paths.get_filename_list("loras"),),       
            }
        }
    RETURN_TYPES = ("ADDITIONAL_MODELS",)
    RETURN_NAMES = ("additional_models",)
    FUNCTION = "select_models"

    CATEGORY = "AD_MotionDirector"

    def select_models(self, motion_module, use_adapter_lora, optional_adapter_lora=""):
        additional_models = []
        motion_model = folder_paths.get_full_path("animatediff_models", motion_module)
        if not Path(motion_model).is_file():
            raise ValueError(f"Motion model {motion_model} does not exist")
        if use_adapter_lora:
            adapter_lora_path = folder_paths.get_full_path("loras", optional_adapter_lora)
            if not Path(adapter_lora_path).is_file():
                raise ValueError(f"Adapter LoRA path {adapter_lora_path} does not exist")
        else:
            adapter_lora_path = ""        

        additional_models.append(motion_model)
        additional_models.append(adapter_lora_path)  
        return (additional_models,)

class ADMD_ValidationSettings:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
        "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
        "inference_steps": ("INT", {"default": 25, "min": 0, "max": 256, "step": 1}),
        "guidance_scale": ("FLOAT", {"default": 8, "min": 0, "max": 32, "step": 0.1}),
        "spatial_scale": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),                                        
        "validation_prompt": ("STRING", {"multiline": True, "default": "",}),
        },
          
        }
    RETURN_TYPES = ("VALIDATION_SETTINGS",)
    RETURN_NAMES = ("validation_settings",)
    FUNCTION = "create_validation_settings"

    CATEGORY = "AD_MotionDirector"

    def create_validation_settings(self, inference_steps, guidance_scale, spatial_scale, seed, validation_prompt):
        # Create a dictionary with the local variables
        local_vars = locals()
        
        # Filter the dictionary to include only the variables you want
        validation_settings = {
            "inference_steps": local_vars["inference_steps"],
            "guidance_scale": local_vars["guidance_scale"],
            "spatial_scale": local_vars["spatial_scale"],
            "seed": local_vars["seed"],
            "validation_prompt": local_vars["validation_prompt"]
        }
       
        return validation_settings,

class ADMD_LoadLora:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "lora_path": ("STRING", {"forceInput": True, "multiline": False, "default": "",}),
                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}),
            },
            "optional": {
                "prev_motion_lora": ("MOTION_LORA",),
            }
        }
    
    RETURN_TYPES = ("MOTION_LORA",)
    CATEGORY = "AD_MotionDirector"
    FUNCTION = "load_motion_lora"

    def load_motion_lora(self, lora_path: str, strength: float, prev_motion_lora: MotionLoraList=None):
        
        if prev_motion_lora is None:
            prev_motion_lora = MotionLoraList()
        else:
            prev_motion_lora = prev_motion_lora.clone()
        full_lora_path = os.path.join(folder_paths.models_dir,"animatediff_motion_lora",lora_path)
        # check if motion lora with name exists
        if not Path(full_lora_path).is_file():
            raise FileNotFoundError(f"Motion lora not found at {full_lora_path}")
        # create motion lora info to be loaded in AnimateDiff Loader
        lora_info = MotionLoraInfo(name=lora_path, strength=strength)
        prev_motion_lora.add_lora(lora_info)

        return (prev_motion_lora,)

class ADMD_SaveLora:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "admd_pipeline": ("ADMDPIPELINE", ),
                "lora_info": ("LORAINFO", ),
            },
        }
    
    RETURN_TYPES = ("STRING", "ADMDPIPELINE",)
    RETURN_NAMES = ("lora_path", "admd_pipeline",)
    CATEGORY = "AD_MotionDirector"
    FUNCTION = "save_motion_lora"

    def save_motion_lora(self, admd_pipeline, lora_info):
        with torch.inference_mode(False):
            validation_pipeline = admd_pipeline['validation_pipeline']
            global_step = admd_pipeline['global_step']

            import copy
            validation_pipeline.to('cpu')  # We do this to prevent VRAM spiking / increase from the new copy
            
            spatial_lora_path = lora_info['spatial_lora_path']
            temporal_lora_path = lora_info['temporal_lora_path']
            lora_name = lora_info['lora_name']
            lora_rank = lora_info['lora_rank']
            temporal_lora_base_path = lora_info['temporal_lora_base_path']

            lora_manager_spatial = LoraHandler(use_unet_lora=True, unet_replace_modules=["Transformer3DModel"])
            lora_manager_spatial.save_lora_weights(
                model=copy.deepcopy(validation_pipeline), 
                save_path=spatial_lora_path, 
                step=global_step,
                use_safetensors=True,
                lora_rank=lora_rank,
                lora_name=lora_name + "_r"+ str(lora_rank) + "_spatial",
            )
            lora_manager_temporal = LoraHandler(use_unet_lora=True, unet_replace_modules=["TemporalTransformerBlock"])
            if lora_manager_temporal is not None:
                lora_manager_temporal.save_lora_weights(
                    model=copy.deepcopy(validation_pipeline), 
                    save_path=temporal_lora_path, 
                    step=global_step,
                    use_safetensors=True,
                    lora_rank=lora_rank,
                    lora_name=lora_name + "_r"+ str(lora_rank) + "_temporal",
                    use_motion_lora_format=True
                )

            final_temporal_lora_name = os.path.join(temporal_lora_base_path, (str(global_step) + "_" + lora_name + "_r"+ str(lora_rank) + "_temporal_unet.safetensors"))
       
            return (final_temporal_lora_name, admd_pipeline)

class ADMD_TrainLora:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": { 
                "admd_pipeline": ("ADMDPIPELINE", ),
                "steps": ("INT", {"default": 100, "min": 0, "max": 10000, "step": 1}),
            },
            "optional": {
                "trigger_input": ("VHS_FILENAMES", ), #attempt to force comfy execution order
                "opt_images_override": ("IMAGE", ),
            }
        }
    
    RETURN_TYPES = ("ADMDPIPELINE",)
    RETURN_NAMES = ("admd_pipeline",)
    CATEGORY = "AD_MotionDirector"
    FUNCTION = "train"

    def train(self, admd_pipeline, steps, opt_images_override=None, trigger_input=None):
        with torch.inference_mode(False):
            train_noise_scheduler = admd_pipeline["train_noise_scheduler"]
            train_noise_scheduler_spatial = admd_pipeline["train_noise_scheduler_spatial"]
            
            text_encoder = admd_pipeline["text_encoder"]
            tokenizer = admd_pipeline["tokenizer"]
            
            optimizer_temporal = admd_pipeline["optimizer_temporal"]
            optimizer_spatial_list = admd_pipeline["optimizer_spatial_list"]
            lr_scheduler_spatial_list = admd_pipeline["lr_scheduler_spatial_list"]
            lr_scheduler_temporal = admd_pipeline["lr_scheduler_temporal"]
            text_prompt = admd_pipeline["text_prompt"]
            pixel_values = admd_pipeline["pixel_values"]
            scaler = admd_pipeline["scaler"]
            seed = admd_pipeline["seed"]
            include_resnet = admd_pipeline["include_resnet"]
            use_offset_noise = False

            torch.manual_seed(seed)

            device = comfy.model_management.get_torch_device()
            comfy.model_management.unload_all_models()

            unet = admd_pipeline["unet"]
            vae = admd_pipeline["vae"]

            unet.to(device)
            vae.to(device)
            text_encoder.to(device)
            unet.enable_gradient_checkpointing()
            unet.train()

            if include_resnet:
                target_spatial_modules = ["Transformer3DModel", "ResnetBlock2D"]
            else: 
                target_spatial_modules = ["Transformer3DModel"]

            target_temporal_modules = ["TemporalTransformerBlock"]
      
            first_epoch = 0
            gradient_accumulation_steps = 1
            global_step = admd_pipeline["global_step"]
            print(f"global_step: {global_step}")
            
            
            # Get the text embedding for conditioning
            with torch.no_grad():
                prompt_ids = tokenizer(
                    text_prompt, 
                    max_length=tokenizer.model_max_length, 
                    padding="max_length", 
                    truncation=True, 
                    return_tensors="pt"
                ).input_ids.to(device)

                #text encoding
                text_encoder.to(device)
                encoder_hidden_states = text_encoder(prompt_ids)[0]
                text_encoder.to('cpu')

                pixel_list = []
                latent_list = []

                if opt_images_override is not None:
                    pixel_values = opt_images_override

                if isinstance(pixel_values, list):
                    print(f"Received {len(pixel_values)} batches:")
                    for p in pixel_values:
                        print("input batch shape:", p.shape)
                        p = p * 2.0 - 1.0 #normalize to the expected range (-1, 1)
                        p = p.permute(0, 3, 1, 2).unsqueeze(0)#B,H,W,C to B,F,C,H,W
                        p = p.to(device)
                        pixel_list.append(p)
                        latent_list.append(tensor_to_vae_latent(p, vae))
                        batch_size = len(pixel_list)
                else:
                    print("Received a single batch")
                    print("input batch shape:", pixel_values.shape)
                    pixel_values = pixel_values * 2.0 - 1.0 #normalize to the expected range (-1, 1)
                    pixel_values = pixel_values.permute(0, 3, 1, 2).unsqueeze(0)#B,H,W,C to B,F,C,H,W
                    pixel_values = pixel_values.to(device)
                    latents = tensor_to_vae_latent(pixel_values, vae)
                    pixel_list.append(pixel_values)
                    latent_list.append(latents)
                    batch_size = 1

                
                print("batch_size:", batch_size)
                vae.to('cpu')

                #num_update_steps_per_epoch = math.ceil(batch_size) / gradient_accumulation_steps
                #num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch)

                progress_bar = tqdm(range(0, steps))
                progress_bar.set_description("Steps")
                pbar = comfy.utils.ProgressBar(steps)

            import itertools
            pixel_cycle = itertools.cycle(pixel_list)
            latent_cycle = itertools.cycle(latent_list)
            ### <<<< Training <<<< ###
            #for epoch in range(first_epoch, steps):
            for step in range(steps):
                pixel_values = next(pixel_cycle)
                latents = next(latent_cycle)
                spatial_scheduler_lr = 0.0
                temporal_scheduler_lr = 0.0

                # Handle Lora Optimizers & Conditions
                for optimizer_spatial in optimizer_spatial_list:
                    optimizer_spatial.zero_grad(set_to_none=True)

                if optimizer_temporal is not None:
                    optimizer_temporal.zero_grad(set_to_none=True)
    
                mask_spatial_lora = random.uniform(0, 1) < 0.2
                #mask_spatial_lora = 0

                # Sample a random timestep for each video
                timesteps = torch.randint(0, 1000, (1,), device=pixel_values.device)
                timesteps = timesteps.long()

                # Add noise to the latents according to the noise magnitude at each timestep
                # (this is the forward diffusion process)
                                    
                noise = sample_noise(latents, 0, use_offset_noise=use_offset_noise)
                comfy.model_management.soft_empty_cache()
                target = noise

                with torch.cuda.amp.autocast():
                    if mask_spatial_lora:
                        loras = extract_lora_child_module(unet, target_replace_module=target_spatial_modules)
                        scale_loras(loras, 0.)
                        loss_spatial = None
                    else:
                        loras = extract_lora_child_module(unet, target_replace_module=target_spatial_modules)
                        scale_loras(loras, 1.0)

                        loras = extract_lora_child_module(unet, target_replace_module=target_temporal_modules)
                        if len(loras) > 0:
                            scale_loras(loras, 0.)
                        
                        ### >>>> Spatial LoRA Prediction >>>> ###
                        noisy_latents = train_noise_scheduler_spatial.add_noise(latents, noise, timesteps)
                        noisy_latents_input, target_spatial = get_spatial_latents(
                            pixel_values,  
                            noisy_latents,
                            target,
                        )
                        model_pred_spatial = unet(noisy_latents_input.unsqueeze(2), timesteps,
                                                encoder_hidden_states=encoder_hidden_states).sample
                        loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
                                                target_spatial.float(), reduction="mean")
                        
                    loras = extract_lora_child_module(unet, target_replace_module=target_temporal_modules)
                    scale_loras(loras, 1.0)
                    
                    ### >>>> Temporal LoRA Prediction >>>> ###
                    noisy_latents = train_noise_scheduler.add_noise(latents, noise, timesteps)
                    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
                    
                    loss_temporal = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
                    loss_temporal = create_ad_temporal_loss(model_pred, loss_temporal, target)
                        
                    # Backpropagate
                    if not mask_spatial_lora:
                        scaler.scale(loss_spatial).backward(retain_graph=True)
                        scaler.step(optimizer_spatial_list[0])
                                        
                    scaler.scale(loss_temporal).backward()
                    scaler.step(optimizer_temporal)
        
                    lr_scheduler_spatial_list[0].step()
                    spatial_scheduler_lr = lr_scheduler_spatial_list[0].get_lr()[0]
                        
                    if lr_scheduler_temporal is not None:
                        lr_scheduler_temporal.step()
                        temporal_scheduler_lr = lr_scheduler_temporal.get_lr()[0]
            
                scaler.update()
                progress_bar.update(1)
                pbar.update(1)
                global_step += 1
                logs = {
                    "Temporal Loss": loss_temporal.detach().item(),
                    "Temporal LR": temporal_scheduler_lr, 
                    "Spatial Loss": loss_spatial.detach().item() if loss_spatial is not None else 0,
                    "Spatial LR": spatial_scheduler_lr
                }
                progress_bar.set_postfix(**logs)

                # if global_step >= steps:
                #     print("BREAK")
                #     break

            admd_pipeline.update({
                "global_step": global_step,
                "unet": unet,
                "scaler": scaler,
            })
            comfy.model_management.soft_empty_cache()
            return (admd_pipeline,)


class ADMD_ValidationSampler:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "validation_settings": ("VALIDATION_SETTINGS", ),
                "admd_pipeline": ("ADMDPIPELINE", ),
            },
        }
    
    RETURN_TYPES = ("ADMDPIPELINE", "IMAGE",)
    RETURN_NAMES = ("admd_pipeline", "images",)
    CATEGORY = "AD_MotionDirector"
    FUNCTION = "train"

    def train(self, admd_pipeline, validation_settings):
        with torch.inference_mode(False):       
            unet = admd_pipeline["unet"]
            text_encoder = admd_pipeline["text_encoder"]
            vae = admd_pipeline["vae"]           
            text_prompt = admd_pipeline["text_prompt"]
            pixel_values = admd_pipeline["pixel_values"]
            
            validation_pipeline = admd_pipeline['validation_pipeline']
          
            device = comfy.model_management.get_torch_device()
            if isinstance(pixel_values, list):
                B, H, W, C = pixel_values[0].shape
            else:
                B, H, W, C = pixel_values.shape

            unet.to(device)
            vae.to(device)
            text_encoder.to(device)
            
            validation_inference_steps = validation_settings["inference_steps"]
            validation_guidance_scale = validation_settings["guidance_scale"]
            validation_spatial_scale = validation_settings["spatial_scale"]
            validation_seed = validation_settings["seed"]
            validation_prompt = validation_settings["validation_prompt"]
                      
            with torch.inference_mode(True):
                samples = []
                generator = torch.Generator(device=device)
                generator.manual_seed(validation_seed)
                
                with torch.cuda.amp.autocast(enabled=True):
                    unet.disable_gradient_checkpointing()
                    unet.eval()
                    loras = extract_lora_child_module(
                        unet, 
                        target_replace_module=["Transformer3DModel"]
                    )
                    scale_loras(loras, validation_spatial_scale)
                    
                    with torch.inference_mode(True):
                        if len(validation_prompt) == 0:
                            prompt = text_prompt
                        else: 
                            prompt = validation_prompt
                        
                        sample = validation_pipeline(
                            prompt,
                            generator    = generator,
                            video_length = B,
                            height       = H,
                            width        = W,
                            num_inference_steps = validation_inference_steps,
                            guidance_scale = validation_guidance_scale,
                        ).videos
                        samples.append(sample)
                # Reshape the sample tensor for returning
                samples = torch.concat(samples)
                samples = samples.view(*samples.shape[1:])
                samples = samples.permute(1, 2, 3, 0).cpu() 
                return (admd_pipeline, samples,)
            
class ADMD_MakeBatchList:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE", ),
                "images2": ("IMAGE", ),
            },
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image_batch_list",)
    CATEGORY = "AD_MotionDirector"
    FUNCTION = "batchlist"

    def batchlist(self, images, images2):
        batch_list = []
        
        if isinstance(images, list):
            batch_list.extend(images)
        else:
            batch_list.append(images)

        if isinstance(images2, list):
            batch_list.extend(images2)
        else:
            batch_list.append(images2)
        
        return (batch_list,)

NODE_CLASS_MAPPINGS = {
    "ADMD_InitializeTraining": ADMD_InitializeTraining,
    "ADMD_DiffusersLoader": ADMD_DiffusersLoader,
    "ADMD_AdditionalModelSelect": ADMD_AdditionalModelSelect,
    "ADMD_ValidationSettings": ADMD_ValidationSettings,
    "ADMD_LoadLora": ADMD_LoadLora,
    "ADMD_SaveLora": ADMD_SaveLora,
    "ADMD_TrainLora": ADMD_TrainLora,
    "ADMD_CheckpointLoader": ADMD_CheckpointLoader,
    "ADMD_ValidationSampler": ADMD_ValidationSampler,
    "ADMD_MakeBatchList": ADMD_MakeBatchList,
    "ADMD_ComfyModelLoader": ADMD_ComfyModelLoader
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "ADMD_InitializeTraining": "ADMD_InitializeTraining",
    "ADMD_DiffusersLoader": "ADMD_DiffusersLoader",
    "ADMD_AdditionalModelSelect": "ADMD_AdditionalModelSelect",
    "ADMD_ValidationSettings": "ADMD_ValidationSettings",
    "ADMD_LoadLora": "ADMD_LoadLora",
    "ADMD_SaveLora": "ADMD_SaveLora",
    "ADMD_TrainLora": "ADMD_TrainLora",
    "ADMD_CheckpointLoader": "ADMD_CheckpointLoader",
    "ADMD_ValidationSampler": "ADMD_ValidationSampler",
    "ADMD_MakeBatchList": "ADMD_MakeBatchList",
    "ADMD_ComfyModelLoader": "ADMD_ComfyModelLoader"
}