diff --git a/.gitignore b/.gitignore index 5130811..03b2b0f 100644 --- a/.gitignore +++ b/.gitignore @@ -171,3 +171,4 @@ exclude.txt *.msgpack *.onnx_data configs/models_config/**/*.png +HunyuanDiT \ No newline at end of file diff --git a/__init__.py b/__init__.py index 289db33..189e9df 100644 --- a/__init__.py +++ b/__init__.py @@ -50,7 +50,7 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSInitWorkspace"] = MZ_KohyaSSInitWorkspace NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSSInitWorkspace" -class MZ_ImagesCopyWorkspace: +class MZ_KohyaSSDatasetConfig: @classmethod def INPUT_TYPES(s): return { @@ -65,9 +65,6 @@ class MZ_ImagesCopyWorkspace: "force_clear_only_images": (["enable", "disable"], {"default": "disable"}), "same_caption_generate": (["enable", "disable"], {"default": "disable"}), "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), - "description": ("STRING", {"default": """ -如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images) -""", "multiline": True}), }, "optional": { "conditioning_images": ("IMAGE",), @@ -78,6 +75,9 @@ class MZ_ImagesCopyWorkspace: RETURN_NAMES = ("workspace_images_dir",) # OUTPUT_NODE = True + MZ_DESC = """ +如果训练类型是controlnet,必须传入预处理后的图片(conditioning_images) +""" FUNCTION = "start" CATEGORY = CATEGORY_NAME @@ -87,9 +87,13 @@ class MZ_ImagesCopyWorkspace: return mz_train_tools_core.MZ_ImageSelecter_call(kwargs) -NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_ImagesCopyWorkspace +NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_KohyaSSDatasetConfig NODE_DISPLAY_NAME_MAPPINGS["MZ_ImagesCopyWorkspace"] = f"{AUTHOR_NAME} - ImagesCopyWorkspace" +# 别名 +NODE_CLASS_MAPPINGS["MZ_KohyaSSDatasetConfig"] = MZ_KohyaSSDatasetConfig +NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSDatasetConfig"] = f"{AUTHOR_NAME} - KohyaSSDatasetConfig" + class MZ_KohyaSSUseConfig: train_config_template_dir = os.path.join( @@ -290,6 +294,7 @@ NODE_CLASS_MAPPINGS["MZ_KohyaSSLoraTrain"] = MZ_KohyaSSLoraTrain NODE_DISPLAY_NAME_MAPPINGS[ "MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)" + class MZ_KohyaSSControlnetTrain: @classmethod @@ -382,6 +387,229 @@ NODE_DISPLAY_NAME_MAPPINGS[ "MZ_LoadImagesFromDirectoryPath"] = f"{AUTHOR_NAME} - LoadImagesFromDirectoryPath" -# 别名 -NODE_CLASS_MAPPINGS["MZ_DatasetConfig"] = MZ_ImagesCopyWorkspace -NODE_DISPLAY_NAME_MAPPINGS["MZ_DatasetConfig"] = f"{AUTHOR_NAME} - DatasetConfig" +from . import mz_train_tools_core_HYDiT + + +class MZ_HYDiTInitWorkspace: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "train_name": ("STRING", {"default": ""}), + "branch": ("STRING", {"default": "5657364143e44ac90f72aeb47b81bd505a95665d"}), + "source": ([ + "github", + "githubfast", + "521github", + "kkgithub", + ], {"default": "github"}), + "seed": ("INT", {"default": 0}), + }, + } + + RETURN_TYPES = ("MZ_TT_HYDiT_WorkspaceConfig",) + RETURN_NAMES = ("workspace_config",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTInitWorkspace_call(kwargs) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTInitWorkspace"] = MZ_HYDiTInitWorkspace +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTInitWorkspace"] = f"{AUTHOR_NAME} - HYDiTInitWorkspace" + + +class MZ_HYDiTDatasetConfig: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), + "resolution": ("INT", {"default": 1024}), + }, + "optional": { + "images": ("IMAGE",), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("workspace_images_dir",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTDatasetConfig_call(kwargs) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTDatasetConfig"] = MZ_HYDiTDatasetConfig +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTDatasetConfig"] = f"{AUTHOR_NAME} - HYDiTDatasetConfig" + + +# class MZ_HYDiTUseConfig: +# @classmethod +# def INPUT_TYPES(s): + +# return { +# "required": { +# "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), +# "workspace_images_dir": ("STRING", {"forceInput": True}), +# }, +# "optional": { +# } +# } + +# RETURN_TYPES = ("MZ_TT_HYDiT_TrainConfig",) +# RETURN_NAMES = ("train_config",) + +# FUNCTION = "start" + +# CATEGORY = CATEGORY_NAME + +# def start(self, **kwargs): +# importlib.reload(mz_train_tools_core_HYDiT) +# # kwargs["train_config_template_dir"] = self.train_config_template_dir +# return mz_train_tools_core_HYDiT.MZ_HYDiTUseConfig_call(kwargs) + + +# NODE_CLASS_MAPPINGS["MZ_HYDiTUseConfig"] = MZ_HYDiTUseConfig +# NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTUseConfig"] = f"{AUTHOR_NAME} - HYDiTUseConfig" + + +class MZ_HYDiTAdvConfig: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "learn_sigma": (["enable", "disable"], {"default": "enable"}), + "text_states_dim": ("INT", {"default": 1024}), + "text_states_dim_t5": ("INT", {"default": 2048}), + "text_len": ("INT", {"default": 77}), + "text_len_t5": ("INT", {"default": 256}), + "norm": (["layer", "instance"], {"default": "layer"}), + "infer_mode": (["torch", "deepspeed"], {"default": "torch"}), + "use_fp16": (["enable", "disable"], {"default": "enable"}), + } + } + + RETURN_TYPES = ("MZ_TT_HYDiT_AdvConfig",) + RETURN_NAMES = ("advanced_config",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + return (kwargs,) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTAdvConfig"] = MZ_HYDiTAdvConfig +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTAdvConfig"] = f"{AUTHOR_NAME} - HYDiTAdvConfig" + + +class MZ_HYDiTTrain: + @classmethod + def INPUT_TYPES(s): + + hunyuan_models_path = os.path.join( + Utils.get_comfyui_models_path(), "hunyuan") + os.makedirs(hunyuan_models_path, exist_ok=True) + + models = Utils.get_models_by_folder(hunyuan_models_path) + folders = Utils.get_folders_by_folder(hunyuan_models_path) + + vae_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "vae")) + unet_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "unet")) + return { + "required": { + "workspace_config": ("MZ_TT_HYDiT_WorkspaceConfig",), + "unet_path": (["auto"] + models + unet_models, {"default": "auto"}), + "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}), + "text_encoder_path": (["auto"] + folders, {"default": "auto"}), + "tokenizer_path": (["auto"] + folders, {"default": "auto"}), + "t5_encoder_path": (["auto"] + folders, {"default": "auto"}), + "resolution": ("INT", {"default": 1024}), + "batch_size": ("INT", {"default": 1}), + "epochs": ("INT", {"default": 50}), + "save_every_n_epochs": ("INT", {"default": 10}), + "base_lora": (["latest", "empty"], {"default": "latest"}), + "sample_generate": (["enable", "disable"], {"default": "enable"}), + "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + }, + "optional": { + "workspace_images_dir": ("STRING", {"forceInput": True}), + "has_no_effect": (AlwaysEqualProxy("*"),), + } + } + + RETURN_TYPES = () + RETURN_NAMES = () + + FUNCTION = "start" + + MZ_DESC = """ +base_size The base resolution (n, n) from which to create multiple resolutions | Recommended values: 256/512/1024 +""" + + OUTPUT_NODE = True + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTTrain_call(kwargs) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTTrain"] = MZ_HYDiTTrain +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTTrain"] = f"{AUTHOR_NAME} - HYDiTTrain" + + +class MZ_TrainToolsDebug: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "object": (AlwaysEqualProxy("*"),), + "indent": ("INT", {"default": 2}), + "depth": ("INT", {"default": 5}), + "width": ("INT", {"default": 80}), + "compact": (["enable", "disable"], {"default": "enable"}), + "sort_keys": (["enable", "disable"], {"default": "enable"}), + "underscore_numbers": (["enable", "disable"], {"default": "enable"}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("debug",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + + from pprint import pprint, pp + object = kwargs["object"] + indent = kwargs["indent"] + depth = kwargs["depth"] + width = kwargs["width"] + compact = kwargs["compact"] == "enable" + sort_keys = kwargs["sort_keys"] == "enable" + underscore_numbers = kwargs["underscore_numbers"] == "enable" + + debug = pp(object, stream=None, indent=indent, depth=depth, width=width, + compact=compact, sort_dicts=sort_keys, underscore_numbers=underscore_numbers) + + return (debug,) + + +NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug +NODE_DISPLAY_NAME_MAPPINGS["MZ_TrainToolsDebug"] = f"{AUTHOR_NAME} - TrainToolsDebug" diff --git a/hook_HYDiT_idk_run.py b/hook_HYDiT_idk_run.py new file mode 100644 index 0000000..50b8f27 --- /dev/null +++ b/hook_HYDiT_idk_run.py @@ -0,0 +1,155 @@ +import argparse + +from index_kits.dataset.make_dataset_core import startup, make_multireso +from index_kits.common import show_index_info +from index_kits import __version__ + + +def common_args(parser): + parser.add_argument('-t', '--target', type=str, + required=True, help='Save path') + + +def get_args(): + parser = argparse.ArgumentParser(description=""" + IndexKits is a tool to build and manage index files for large-scale datasets. + It supports both base index and multi-resolution index. + + Introduction + ------------ + This command line tool provides the following functionalities: + 1. Show index v2 information + 2. Build base index v2 + 3. Build multi-resolution index v2 + + Examples + -------- + 1. Show index v2 information + index_kits show /path/to/index.json + + 2. Build base index v2 + Default usage: + index_kits base -c /path/to/config.yaml -t /path/to/index.json + + Use multiple processes: + index_kits base -c /path/to/config.yaml -t /path/to/index.json -w 40 + + 3. Build multi-resolution index v2 + + Build with a configuration file: + index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json + + Build by specifying arguments without a configuration file: + index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + + Build by specifying target-ratios: + index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json + + Build with multiple source index files. + index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + """, formatter_class=argparse.RawTextHelpFormatter) + sub_parsers = parser.add_subparsers(dest='task', required=True) + + # Show index message + show_parser = sub_parsers.add_parser('show', description=""" + Show base/multireso index v2 information. + + Example + ------- + index_kits show /path/to/index.json + """, formatter_class=argparse.RawTextHelpFormatter) + show_parser.add_argument( + 'src', type=str, help='Path to a base/multireso index file.') + show_parser.add_argument( + '--arrow-files', action='store_true', help='Show arrow files only.') + show_parser.add_argument('--depth', type=int, default=1, + help='Arrow file depth. Default is 1, the level of last folder in the arrow file path. ' + 'Set it to 0 to show the full path including `xxx/last_folder/*.arrow`.') + + # Single resolution bucket + base_parser = sub_parsers.add_parser('base', description=""" + Build base index v2. + + Example + ------- + index_kits base -c /path/to/config.yaml -t /path/to/index.json + """, formatter_class=argparse.RawTextHelpFormatter) + base_parser.add_argument('-c', '--config', type=str, + required=True, help='Configuration file path') + common_args(base_parser) + base_parser.add_argument('-w', '--world-size', type=int, default=1) + base_parser.add_argument('--work-dir', type=str, + default='.', help='Work directory') + base_parser.add_argument('--use-cache', action='store_true', help='Use cache to avoid reprocessing. ' + 'Perform merge pkl results directly.') + + # Multi-resolution bucket + mo_parser = sub_parsers.add_parser('multireso', description=""" + Build multi-resolution index v2 + + Example + ------- + Build with a configuration file: + index_kits multireso -c /path/to/config.yaml -t /path/to/index_mb_gt512.json + + Build by specifying arguments without a configuration file: + index_kits multireso --src /path/to/index.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + + Build by specifying target-ratios: + index_kits multireso --src /path/to/index.json --base-size 512 --target-ratios 1:1 4:3 3:4 16:9 9:16 --min-size 512 -t /path/to/index_mb_gt512.json + + Build with multiple source index files. + index_kits multireso --src /path/to/index1.json /path/to/index2.json --base-size 512 --reso-step 32 --min-size 512 -t /path/to/index_mb_gt512.json + """, formatter_class=argparse.RawTextHelpFormatter) + mo_parser.add_argument('-c', '--config', type=str, default=None, + help='Configuration file path in a yaml format. Either --config or --src must be provided.') + mo_parser.add_argument('-s', '--src', type=str, nargs='+', default=None, + help='Source index files. Either --config or --src must be provided.') + common_args(mo_parser) + mo_parser.add_argument('--base-size', type=int, default=None, + help="Base size. Typically set as 256/512/1024 according to image size you train model.") + mo_parser.add_argument('--reso-step', type=int, default=None, + help="Resolution step. Either reso_step or target_ratios must be provided.") + mo_parser.add_argument('--target-ratios', type=str, nargs='+', default=None, + help="Target ratios. Either reso_step or target_ratios must be provided.") + mo_parser.add_argument('--md5-file', type=str, default=None, + help='You can provide an md5 to height and width file to accelerate the process. ' + 'It is a pickle file that contains a dict, which maps md5 to (height, width) tuple.') + mo_parser.add_argument('--align', type=int, default=16, + help="Used when --target-ratios is provided. Align size of source image height and width.") + mo_parser.add_argument('--min-size', type=int, default=0, + help="Minimum size. Images smaller than this size will be ignored.") + + # Common + parser.add_argument('-v', '--version', action='version', + version=f'%(prog)s {__version__}') + + args = parser.parse_args() + return args + + +if __name__ == '__main__': + args = get_args() + if args.task == 'show': + show_index_info(args.src, + args.arrow_files, + args.depth, + ) + elif args.task == 'base': + startup(args.config, + args.target, + args.world_size, + args.work_dir, + use_cache=args.use_cache, + ) + elif args.task == 'multireso': + make_multireso(args.target, + args.config, + args.src, + args.base_size, + args.reso_step, + args.target_ratios, + args.align, + args.min_size, + args.md5_file, + ) diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py new file mode 100644 index 0000000..d33d547 --- /dev/null +++ b/hook_HYDiT_main.py @@ -0,0 +1,710 @@ + +import argparse +import json +import logging +import os +import sys +from types import SimpleNamespace +import gc +import json +import os +import random +import sys +import time +from functools import partial +from glob import glob +from pathlib import Path +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader +from torchvision.transforms import functional as TF +from diffusers.models import AutoencoderKL +from transformers import BertModel, BertTokenizer, logging as tf_logging + +from hydit.config import get_args +# from hydit.constants import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER +from hydit.lr_scheduler import WarmupLR +from hydit.data_loader.arrow_load_stream import TextImageArrowStream +from hydit.diffusion import create_diffusion +from hydit.ds_config import deepspeed_config_from_args +from hydit.modules.ema import EMA +from hydit.modules.fp16_layers import Float16Module +from hydit.modules.models import HUNYUAN_DIT_MODELS +from hydit.modules.posemb_layers import init_image_posemb +from hydit.utils.tools import set_seeds, get_trainable_params +from IndexKits.index_kits import ResolutionGroup +from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler +from peft import LoraConfig, get_peft_model +from tqdm import tqdm + + +UNET_PATH = "ckpts/t2i/model/pytorch_model_ema.pt" + + +def set_unet_path(path): + global UNET_PATH + UNET_PATH = path + + +VAE_EMA_PATH = "ckpts/t2i/sdxl-vae-fp16-fix" + + +def set_vae_ema_path(path): + global VAE_EMA_PATH + VAE_EMA_PATH = path + + +TOKENIZER = "ckpts/t2i/tokenizer" + + +def set_tokenizer_path(path): + global TOKENIZER + TOKENIZER = path + + +TEXT_ENCODER = 'ckpts/t2i/clip_text_encoder' + + +def set_text_encoder_path(path): + global TEXT_ENCODER + TEXT_ENCODER = path + + +T5_ENCODER = { + 'MT5': 'ckpts/t2i/mt5', + 'attention_mask': True, + 'layer_index': -1, + 'attention_pool': True, + 'torch_dtype': torch.float16, + 'learnable_replace': True +} + + +def set_t5_encoder_path(path): + global T5_ENCODER + T5_ENCODER['MT5'] = path + + +def create_logger(log_dir=None): + logger = logging.getLogger(__name__) + logger.addHandler(logging.NullHandler()) + return logger + + +def create_exp_folder(args, rank): + if rank == 0: + os.makedirs(args.results_dir, exist_ok=True) + + # print('args.results_dir', args.results_dir) + _existed_experiments = os.listdir(args.results_dir) + # print('_existed_experiments', _existed_experiments) + existed_experiments = [] + for x in _existed_experiments: + if os.path.isdir(os.path.join(args.results_dir, x)): + existed_experiments.append(x) + + if len(existed_experiments) == 0: + experiment_index = 1 + else: + existed_experiments.sort() + # print('existed_experiments', existed_experiments) + experiment_index = max([int(x.split('-')[0]) + for x in existed_experiments]) + 1 + model_string_name = args.task_flag if args.task_flag else args.model.replace( + "/", "-") + # Create an experiment folder + experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" + # Stores saved model checkpoints + checkpoint_dir = f"{experiment_dir}/checkpoints" + if rank == 0: + os.makedirs(checkpoint_dir, exist_ok=True) + logger = create_logger(experiment_dir) + logger.info(f"Experiment directory created at {experiment_dir}") + else: + logger = create_logger() + experiment_dir = "" + logger.info = print + return experiment_dir, checkpoint_dir, logger + + +def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir): + def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"): + cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}" + if rank == 0: + if args.use_fp16: + try: + model.module.save_pretrained(cur_ckpt_save_dir) + except Exception as e: + model.save_pretrained(cur_ckpt_save_dir) + else: + model.module.save_pretrained(cur_ckpt_save_dir) + + checkpoint_path = "[Not rank 0. Disabled output.]" + + client_state = { + "steps": train_steps, + "epoch": epoch, + "args": args + } + if ema is not None: + client_state['ema'] = ema.state_dict() + + dst_paths = [] + + if train_steps % args.ckpt_every == 0: + checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt" + try: + if args.training_parts == "lora": + save_lora_weight(checkpoint_dir, client_state, + tag=f"{train_steps:07d}.pt") + else: + model.save_checkpoint( + checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt") + dst_paths.append(checkpoint_path) + logger.info(f"Saved checkpoint to {checkpoint_path}") + except: + logger.error(f"Saved failed to {checkpoint_path}") + + if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps: + save_name = "latest.pt" + checkpoint_path = f"{checkpoint_dir}/{save_name}" + try: + if args.training_parts == "lora": + save_lora_weight(checkpoint_dir, client_state, + tag=f"{save_name}") + else: + model.save_checkpoint( + checkpoint_dir, client_state=client_state, tag=f"{save_name}") + dst_paths.append(checkpoint_path) + logger.info(f"Saved checkpoint to {checkpoint_path}") + except: + logger.error(f"Saved failed to {checkpoint_path}") + + if rank == 0 and len(dst_paths) > 0: + # Delete optimizer states to avoid occupying too much disk space. + for dst_path in dst_paths: + for opt_state_path in glob(f"{dst_path}/zero_dp_rank_*_tp_rank_00_pp_rank_00_optim_states.pt"): + os.remove(opt_state_path) + + return checkpoint_path + + +@torch.no_grad() +def prepare_model_inputs(args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img): + image, text_embedding, text_embedding_mask, text_embedding_t5, text_embedding_mask_t5, kwargs = batch + + # clip & mT5 text embedding + text_embedding = text_embedding.to(device) + text_embedding_mask = text_embedding_mask.to(device) + encoder_hidden_states = text_encoder( + text_embedding.to(device), + attention_mask=text_embedding_mask.to(device), + )[0] + text_embedding_t5 = text_embedding_t5.to(device).squeeze(1) + text_embedding_mask_t5 = text_embedding_mask_t5.to(device).squeeze(1) + with torch.no_grad(): + output_t5 = text_encoder_t5( + input_ids=text_embedding_t5, + attention_mask=text_embedding_mask_t5 if T5_ENCODER['attention_mask'] else None, + output_hidden_states=True + ) + encoder_hidden_states_t5 = output_t5['hidden_states'][T5_ENCODER['layer_index']].detach( + ) + + # additional condition + image_meta_size = kwargs['image_meta_size'].to(device) + style = kwargs['style'].to(device) + + if args.extra_fp16: + image = image.half() + image_meta_size = image_meta_size.half() if image_meta_size is not None else None + + # Map input images to latent space + normalize latents: + image = image.to(device) + vae_scaling_factor = vae.config.scaling_factor + latents = vae.encode(image).latent_dist.sample().mul_( + vae_scaling_factor).to(device) + + # positional embedding + _, _, height, width = image.shape + reso = f"{height}x{width}" + + # print("freqs_cis_img:", freqs_cis_img.keys()) + # print("freqs_cis_img[reso]:", freqs_cis_img[reso]) + + cos_cis_img, sin_cis_img = freqs_cis_img[reso] + + # Model conditions + model_kwargs = dict( + encoder_hidden_states=encoder_hidden_states, + text_embedding_mask=text_embedding_mask, + encoder_hidden_states_t5=encoder_hidden_states_t5, + text_embedding_mask_t5=text_embedding_mask_t5, + image_meta_size=image_meta_size, + style=style, + cos_cis_img=cos_cis_img, + sin_cis_img=sin_cis_img, + ) + + return latents, model_kwargs + + +def Core(args, LOG): + if args.training_parts == "lora": + args.use_ema = False + + assert torch.cuda.is_available(), "Training currently requires at least one GPU." + + batch_size = args.batch_size + grad_accu_steps = args.grad_accu_steps + global_batch_size = batch_size * grad_accu_steps + + seed = args.global_seed + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + print(f"Starting ,seed={seed}.") + rank = 0 + # Setup an experiment folder + experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank) + + # Log all the arguments + logger.info(sys.argv) + logger.info(str(args)) + # Save to a json file + args_dict = vars(args) + with open(f"{experiment_dir}/args.json", 'w') as f: + json.dump(args_dict, f, indent=4) + + # Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel." + # If needed, just comment the following line. + tf_logging.set_verbosity_error() + + # =========================================================================== + # Building HYDIT + # =========================================================================== + + logger.info("Building HYDIT Model.") + + # --------------------------------------------------------------------------- + # Training sample base size, such as 256/512/1024. Notice that this size is + # just a base size, not necessary the actual size of training samples. Actual + # size of the training samples are correlated with `resolutions` when enabling + # multi-resolution training. + # --------------------------------------------------------------------------- + image_size = args.image_size + if type(image_size) == int: + image_size = [image_size, image_size] + if len(image_size) == 1: + image_size = [image_size[0], image_size[0]] + if len(image_size) != 2: + raise ValueError(f"Invalid image size: {args.image_size}") + assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \ + f"got {image_size}" + latent_size = [image_size[0] // 8, image_size[1] // 8] + + # initialize model by deepspeed + model = HUNYUAN_DIT_MODELS[args.model](args, + input_size=latent_size, + log_fn=logger.info, + ) + # Multi-resolution / Single-resolution training. + if args.multireso: + resolutions = ResolutionGroup(image_size[0], + align=16, + step=args.reso_step, + target_ratios=args.target_ratios).data + else: + resolutions = ResolutionGroup(image_size[0], + align=16, + target_ratios=['1:1']).data + + freqs_cis_img = init_image_posemb(args.rope_img, + resolutions=resolutions, + patch_size=model.patch_size, + hidden_size=model.hidden_size, + num_heads=model.num_heads, + log_fn=logger.info, + rope_real=args.rope_real, + ) + + # Create EMA model and convert to fp16 if needed. + ema = None + device = "cuda" + if args.use_ema: + ema = EMA(args, model, device, logger) + + # Setup FP16 main model: + if args.use_fp16: + model = Float16Module(model, args) + logger.info( + f" Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}") + + diffusion = create_diffusion( + noise_schedule=args.noise_schedule, + predict_type=args.predict_type, + learn_sigma=args.learn_sigma, + mse_loss_weight_type=args.mse_loss_weight_type, + beta_start=args.beta_start, + beta_end=args.beta_end, + noise_offset=args.noise_offset, + ) + + # Setup VAE + logger.info(f" Loading vae from {VAE_EMA_PATH}") + + if os.path.isdir(VAE_EMA_PATH): + vae = AutoencoderKL.from_pretrained(VAE_EMA_PATH) + else: + vae = AutoencoderKL.from_single_file(VAE_EMA_PATH) + # Setup BERT text encoder + logger.info(f" Loading Bert text encoder from {TEXT_ENCODER}") + text_encoder = BertModel.from_pretrained( + TEXT_ENCODER, False, revision=None) + # Setup BERT tokenizer: + logger.info(f" Loading Bert tokenizer from {TOKENIZER}") + tokenizer = BertTokenizer.from_pretrained(TOKENIZER) + # Setup T5 text encoder + from hydit.modules.text_encoder import MT5Embedder + mt5_path = T5_ENCODER['MT5'] + embedder_t5 = MT5Embedder( + mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5) + tokenizer_t5 = embedder_t5.tokenizer + text_encoder_t5 = embedder_t5.model + + if args.extra_fp16: + logger.info(f" Using fp16 for extra modules: vae, text_encoder") + vae = vae.half().to(device) + text_encoder = text_encoder.half().to(device) + text_encoder_t5 = text_encoder_t5.half().to(device) + else: + vae = vae.to(device) + text_encoder = text_encoder.to(device) + text_encoder_t5 = text_encoder_t5.to(device) + + logger.info( + f" Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}") + + # =========================================================================== + # Building Dataset + # =========================================================================== + + logger.info(f"Building Streaming Dataset.") + logger.info(f" Loading index file {args.index_file} (v2)") + world_size = 1 + + dataset = TextImageArrowStream(args=args, + resolution=image_size[0], + random_flip=args.random_flip, + log_fn=logger.info, + index_file=args.index_file, + multireso=args.multireso, + batch_size=batch_size, + world_size=world_size, + random_shrink_size_cond=args.random_shrink_size_cond, + merge_src_cond=args.merge_src_cond, + uncond_p=args.uncond_p, + text_ctx_len=args.text_len, + tokenizer=tokenizer, + uncond_p_t5=args.uncond_p_t5, + text_ctx_len_t5=args.text_len_t5, + tokenizer_t5=tokenizer_t5, + ) + if args.multireso: + sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, + shuffle=False, drop_last=True, batch_size=batch_size) + else: + sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed, + shuffle=False, drop_last=True) + loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, sampler=sampler, + num_workers=args.num_workers, pin_memory=True, drop_last=True) + logger.info(f" Dataset contains {len(dataset):,} images.") + logger.info(f" Index file: {args.index_file}.") + if args.multireso: + logger.info(f' Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} ' + f'and base size {dataset.index_manager.base_size}') + logger.info(f'\n {dataset.index_manager.resolutions}') + + # =========================================================================== + # Loading parameter + # =========================================================================== + + logger.info(f"Loading parameter") + start_epoch = 0 + start_epoch_step = 0 + train_steps = 0 + # Resume checkpoint if needed + model, ema, start_epoch, start_epoch_step, train_steps = model_resume( + args, model, ema, logger) + + if args.training_parts == "lora": + loraconfig = LoraConfig( + r=args.rank, + lora_alpha=args.rank, + target_modules=args.target_modules + ) + if args.use_fp16: + model.module = get_peft_model(model.module, loraconfig) + else: + model = get_peft_model(model, loraconfig) + + logger.info(f" Training parts: {args.training_parts}") + + # =========================================================================== + # Training + # =========================================================================== + + model.to(device) + model.train() + if args.use_ema: + ema.eval() + + print(f" Worker {rank} ready.") + + iters_per_epoch = len(loader) + logger.info( + " ****************************** Running training ******************************") + logger.info(f" Number GPUs: {world_size}") + logger.info(f" Number training samples: {len(dataset):,}") + logger.info( + f" Number parameters: {sum(p.numel() for p in model.parameters()):,}") + logger.info( + f" Number trainable params: {sum(p.numel() for p in get_trainable_params(model)):,}") + logger.info( + " ------------------------------------------------------------------------------") + logger.info(f" Iters per epoch: {iters_per_epoch:,}") + logger.info(f" Batch size per device: {batch_size}") + logger.info( + f" Batch size all device: {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)") + logger.info(f" Gradient Accu steps: {args.grad_accu_steps}") + logger.info( + f" Total optimization steps: {args.epochs * iters_per_epoch // grad_accu_steps:,}") + + logger.info( + f" Training epochs: {start_epoch}/{args.epochs}") + logger.info( + f" Training epoch steps: {start_epoch_step:,}/{iters_per_epoch:,}") + logger.info( + f" Training total steps: {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}") + logger.info( + " ------------------------------------------------------------------------------") + logger.info(f" Noise schedule: {args.noise_schedule}") + logger.info( + f" Beta limits: ({args.beta_start}, {args.beta_end})") + logger.info(f" Learn sigma: {args.learn_sigma}") + logger.info(f" Prediction type: {args.predict_type}") + logger.info(f" Noise offset: {args.noise_offset}") + + logger.info( + " ------------------------------------------------------------------------------") + logger.info( + f" Using EMA model: {args.use_ema} ({args.ema_dtype})") + if args.use_ema: + logger.info( + f" Using EMA decay: {ema.max_value if args.use_ema else None}") + logger.info( + f" Using EMA warmup power: {ema.power if args.use_ema else None}") + logger.info(f" Using main model fp16: {args.use_fp16}") + logger.info(f" Using extra modules fp16: {args.extra_fp16}") + logger.info( + " ------------------------------------------------------------------------------") + logger.info(f" Experiment directory: {experiment_dir}") + logger.info( + " *******************************************************************************") + + if args.gc_interval > 0: + gc.disable() + gc.collect() + + # Variables for monitoring/logging purposes: + log_steps = 0 + running_loss = 0 + start_time = time.time() + + if args.async_ema: + ema_stream = torch.cuda.Stream() + + total_steps = args.epochs * len(loader) + + if args.use_fp16: + optimizer = torch.optim.AdamW( + model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=4e-6) + else: + optimizer = torch.optim.AdamW( + model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + + # torch.autograd.set_detect_anomaly(True) + # with torch.autograd.detect_anomaly(): + with tqdm(total=total_steps, initial=train_steps) as pbar: + + # Training loop + for epoch in range(start_epoch, args.epochs): + logger.info(f" Start random shuffle with seed={seed}") + # Makesure all processors use the same seed to shuffle dataset. + dataset.shuffle(seed=args.global_seed + epoch, fast=True) + logger.info(f" End of random shuffle") + + # Move sampler to start_index + if not args.multireso: + start_index = start_epoch_step * world_size * batch_size + if start_index != sampler.start_index: + sampler.start_index = start_index + # Reset start_epoch_step to zero, to ensure next epoch will start from the beginning. + start_epoch_step = 0 + logger.info( + f" Iters left this epoch: {len(loader):,}") + + logger.info(f" Beginning epoch {epoch}...") + step = 0 + mean_loss = 0 + for batch in loader: + step += 1 + + # Zero gradients, backward pass, and optimizer step: + optimizer.zero_grad() + latents, model_kwargs = prepare_model_inputs( + args, batch, device, vae, text_encoder, text_encoder_t5, freqs_cis_img) + # with torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True): + + if True: + # training model by deepspeed while use fp16 + if args.use_fp16: + if args.use_ema and args.async_ema: + with torch.cuda.stream(ema_stream): + ema.update(model.module.module, step=step) + torch.cuda.current_stream().wait_stream(ema_stream) + + loss_dict = diffusion.training_losses( + model=model, x_start=latents, model_kwargs=model_kwargs) + + loss = loss_dict["loss"].mean() + + # print(f"step={step}, loss={loss.item()}") + + # model.backward(loss) + + loss.backward() + + # last_batch_iteration = ( + # train_steps + 1) // (global_batch_size // (batch_size * world_size)) + # model.step( + # lr_kwargs={'last_batch_iteration': last_batch_iteration}) + + optimizer.step() + + if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader) - 1): + if args.use_fp16: + ema.update(model.module.module, step=step) + else: + ema.update(model.module, step=step) + + train_steps += 1 + mean_loss += loss.item() + pbar.update(1) + pbar.set_description(f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}") + + # =========================================================================== + # Log loss values: + # =========================================================================== + # running_loss += loss.item() + # log_steps += 1 + # if train_steps % args.log_every == 0: + # # Measure training speed: + # torch.cuda.synchronize() + # end_time = time.time() + # steps_per_sec = log_steps / (end_time - start_time) + # # Reduce loss history over all processes: + # avg_loss = torch.tensor( + # running_loss / log_steps, device=device) + # # dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM) + + # # avg_loss = avg_loss.item() / world_size + # # get lr from deepspeed fused optimizer + # logger.info(f"(step={train_steps:07d}) " + + # (f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") + + # f"Train Loss: {avg_loss:.4f}, " + # # f"Lr: {opt.param_groups[0]['lr']:.6g}, " + # f"Steps/Sec: {steps_per_sec:.2f}, " + # f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}") + # # Reset monitoring variables: + # running_loss = 0 + # log_steps = 0 + # start_time = time.time() + + # collect gc: + if args.gc_interval > 0 and (step % args.gc_interval == 0): + gc.collect() + + if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0 # or train_steps == args.max_training_steps + ) and train_steps > 0: + logger.info( + f" Saving checkpoint at step {train_steps}.") + save_checkpoint(args, rank, logger, model, ema, + epoch, train_steps, checkpoint_dir) + + if train_steps >= args.max_training_steps: + logger.info(f"Breaking step loop at {train_steps}.") + break + + LOG({ + "type": "sample_images", + "global_step": train_steps, + "total_steps": total_steps, + # "latent": noise_pred_latent_path, + }) + + if train_steps >= args.max_training_steps: + logger.info(f"Breaking epoch loop at {epoch}.") + break + + +def model_resume(args, model, ema, logger): + """ + Load pretrained weights. + """ + start_epoch = 0 + start_epoch_step = 0 + train_steps = 0 + resume_path = UNET_PATH + if not Path(resume_path).exists(): + raise FileNotFoundError( + f" Cannot find checkpoint from {resume_path}") + + logger.info(f"Resume from checkpoint {resume_path}") + + if args.resume_split: + # Resume main model + + resume_ckpt_module = torch.load( + resume_path, map_location=lambda storage, loc: storage) + model.load_state_dict(resume_ckpt_module, strict=False) + + # Resume ema model + if args.use_ema: + if args.module_to_ema: + if "resume_ckpt_module" in locals(): + logger.info(f" Resume ema model from main states.") + ema.load_state_dict(resume_ckpt_module, strict=args.strict) + else: + logger.info(f" Resume ema model from module states.") + resume_ckpt_module = torch.load( + resume_path, map_location=lambda storage, loc: storage) + ema.load_state_dict(resume_ckpt_module, strict=args.strict) + else: + if "resume_ckpt_ema" in locals(): + logger.info(f" Resume ema model from EMA states.") + ema.load_state_dict(resume_ckpt_ema, strict=args.strict) + else: + logger.info(f" Resume ema model from EMA states.") + resume_ckpt_ema = torch.load(resume_path, + map_location=lambda storage, loc: storage) + ema.load_state_dict(resume_ckpt_ema, strict=args.strict) + else: + raise ValueError( + " “If `resume` is True, then either `resume_split` must be true.”") + + return model, ema, start_epoch, start_epoch_step, train_steps diff --git a/hook_HYDiT_run.py b/hook_HYDiT_run.py new file mode 100644 index 0000000..3ced37c --- /dev/null +++ b/hook_HYDiT_run.py @@ -0,0 +1,146 @@ + +import argparse +import json +import os +import sys +from types import SimpleNamespace + + +# python -m torch.distributed.launch ./hook_HYDi_run.py --sys_path /data/ComfyUI/models/minus_zone_models/train_tools/HunyuanDiT/ + + +class SimpleNamespaceCNWarrper(SimpleNamespace): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__dict__.update(kwargs) + self.__iter__ = lambda: iter(kwargs.keys()) + # is not iterable + + def __iter__(self): + return iter(self.__dict__.keys()) + # object has no attribute 'num_attention_heads' + + def __getattr__(self, name): + return self.__dict__.get(name, None) + + +import requests + +master_port = 0 + + +def LOG(log): + if master_port == 0: + raise Exception("master_port is 0") + # 发送http + resp = requests.request("post", f"http://127.0.0.1:{master_port}/log", data=json.dumps(log), headers={ + "Content-Type": "application/json"}) + if resp.status_code != 200: + raise Exception(f"LOG failed: {resp.text}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + conflict_handler='resolve', + ) + parser.add_argument("--sys_path", type=str, default="") + parser.add_argument("--train_config_file", type=str, default="") + parser.add_argument("--master_port", type=int, default=0) + args = parser.parse_args() + + master_port = args.master_port + + print(f"master_port = {master_port}") + + sys_path = args.sys_path + if sys_path != "": + sys.path.append(sys_path) + + print("HYDi run hook") + + try: + from . import hook_HYDiT_main + except Exception as e: + import hook_HYDiT_main + + import hydit.config + + def _handle_conflict_error(self, *args, **kwargs): + pass + + def parse_args(self, args=None, namespace=None): + args, argv = self.parse_known_args(args, namespace) + return args + + train_config_file = args.train_config_file + + if train_config_file == "": + raise ValueError("train_config_file is empty") + + train_config = {} + with open(train_config_file, "r") as f: + train_config = json.load(f) + + argparse.ArgumentParser._handle_conflict_error = _handle_conflict_error + argparse.ArgumentParser._handle_conflict_resolve = _handle_conflict_error + argparse.ArgumentParser.parse_args = parse_args + margs = hydit.config.get_args() + margs.model = train_config.get("model", "DiT-g/2") + + + margs.task_flag = train_config.get( + "task_flag", "lora_porcelain_ema_rank64") + + + margs.resume_split = train_config.get("resume_split", True) + margs.ema_to_module = train_config.get("ema_to_module", True) + margs.deepspeed = False + margs.predict_type = train_config.get("predict_type", "v_prediction") + margs.training_parts = train_config.get("training_parts", "lora") + margs.batch_size = train_config.get("batch_size", 1) + margs.grad_accu_steps = train_config.get("grad_accu_steps", 1) + margs.global_seed = train_config.get("global_seed", 0) + margs.use_flash_attn = train_config.get("use_flash_attn", False) + margs.use_fp16 = train_config.get("use_fp16", True) + margs.qk_norm = train_config.get("qk_norm", True) + margs.ema_dtype = train_config.get("ema_dtype", "fp32") + margs.async_ema = False + margs.multireso = train_config.get("multireso", True) + margs.epochs = train_config.get("epochs", 50) + margs.target_ratios = train_config.get( + "target_ratios", ['1:1', '3:4', '4:3', '16:9', '9:16']) + margs.rope_img = train_config.get("rope_img", "base512") + margs.image_size = train_config.get("image_size", 512) + margs.rope_real = train_config.get("rope_real", True) + + margs.index_file = train_config.get("index_file", None) + + margs.lr = train_config.get("lr", 1e-5) + + margs.noise_offset = train_config.get("noise_offset", 0.1) + + margs.log_every = train_config.get("log_every", 10) + + margs.results_dir = train_config.get("results_dir") + margs.mse_loss_weight_type = train_config.get("mse_loss_weight_type", "constant") + + for k, v in train_config.items(): + if hasattr(margs, k): + setattr(margs, k, v) + + hook_HYDiT_main.set_unet_path( + train_config.get("unet_path")) + + hook_HYDiT_main.set_vae_ema_path( + train_config.get("vae_ema_path")) + + hook_HYDiT_main.set_text_encoder_path( + train_config.get("text_encoder_path")) + + hook_HYDiT_main.set_tokenizer_path( + train_config.get("tokenizer_path")) + + hook_HYDiT_main.set_t5_encoder_path( + train_config.get("t5_encoder_path")) + + hook_HYDiT_main.Core(margs, LOG) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py new file mode 100644 index 0000000..7c04555 --- /dev/null +++ b/mz_train_tools_core_HYDiT.py @@ -0,0 +1,394 @@ + + +import argparse +import csv +import hashlib +import os +import shutil +import sys +import json +import subprocess +import time + +import torch + +from .mz_train_tools_utils import Utils +import folder_paths +import nodes + + +git_accelerate_urls = { + "githubfast": "githubfast.com", + "521github": "521github.com", + "kkgithub": "kkgithub.com", +} + +# 初始化工具仓库和工作区 + + +def MZ_HYDiTInitWorkspace_call(args={}): + mz_dir = Utils.get_minus_zone_models_path() + git_url = "https://github.com/Tencent/HunyuanDiT" + source = args.get("source", "github") + kohya_ss_lora_dir = os.path.join(mz_dir, "train_tools", "HunyuanDiT") + if git_accelerate_urls.get(source, None) is not None: + git_url = f"https://{git_accelerate_urls[source]}/Tencent/HunyuanDiT" + try: + if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")): + subprocess.run( + ["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True) + + # 切换远程分支 git remote set-branches origin 'main' + branch = args.get("branch", "main") + + # 查看本地分支是否一致 + short_result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True) + result = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=kohya_ss_lora_dir, stdout=subprocess.PIPE, check=True) + + short_current_branch = short_result.stdout.decode().strip() + long_current_branch = result.stdout.decode().strip() + print( + f"当前分支(current branch): {long_current_branch}({short_current_branch})") + print(f"目标分支(target branch): {branch}") + + if branch != result.stdout.decode() and branch != short_result.stdout.decode(): + subprocess.run( + ["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True) + subprocess.run( + ["git", "fetch", "--depth", "1", "origin", branch], cwd=kohya_ss_lora_dir, check=True) + + # 恢复所有文件 + subprocess.run( + ["git", "checkout", "."], cwd=kohya_ss_lora_dir, check=True) + + subprocess.run( + ["git", "checkout", branch], cwd=kohya_ss_lora_dir, check=True) + + except Exception as e: + raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...") + + workspace_name = args.get("train_name", None) + workspace_name = workspace_name.strip() + + if workspace_name is None or workspace_name == "": + raise Exception("训练名称不能为空(train_name is required)") + + args["workspace_name"] = workspace_name + workspaces_dir = os.path.join( + folder_paths.output_directory, "mz_train_workspaces") + + os.makedirs(workspaces_dir, exist_ok=True) + + workspace_dir = os.path.join(workspaces_dir, workspace_name) + if not os.path.exists(workspace_dir): + os.makedirs(workspace_dir) + + args["workspace_dir"] = workspace_dir + + return ( + args, + ) + + +def MZ_HYDiTDatasetConfig_call(args={}): + workspace_config = args.get("workspace_config", {}) + workspace_name = workspace_config.get("workspace_name", None) + workspace_dir = os.path.join( + folder_paths.output_directory, "mz_train_workspaces", workspace_name) + train_images_dir = os.path.join(workspace_dir, "train_images") + os.makedirs(train_images_dir, exist_ok=True) + + images = args.get("images", None) + + if images is None: + print(f"训练数据位于:{train_images_dir},未检测到传入图片默认已有训练数据,将直接跳过数据准备步骤...") + return ( + train_images_dir, + ) + + pil_images = Utils.tensors2pil_list(images) + + resolution = args.get("resolution", 512) + + if workspace_name is None or workspace_name == "": + raise Exception("训练名称不能为空(workspace_name is required)") + + saved_images_path = [] + + for i, pil_image in enumerate(pil_images): + pil_image = Utils.resize_max(pil_image, resolution, resolution) + width, height = pil_image.size + filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png" + pil_image.save(os.path.join(train_images_dir, filename)) + saved_images_path.append(filename) + + return ( + train_images_dir, + ) + + +def MZ_HYDiTTrain_call(args={}): + # raise Exception(args) + resolution = args.get("resolution") + + workspace_config = args.get("workspace_config", {}) + workspace_name = workspace_config.get("workspace_name", None) + workspace_dir = workspace_config.get("workspace_dir", None) + + workspace_images_dir = args.get("workspace_images_dir", None) + if workspace_images_dir is None: + workspace_images_dir = os.path.join(workspace_dir, "train_images") + + full_filenames = os.listdir(workspace_images_dir) + + # 创建image_text.csv文件 + csv_filename = os.path.join(workspace_dir, "image_text.csv") + + if os.path.exists(csv_filename): + os.remove(csv_filename) + + max_width, max_height, min_width, min_height = 0, 0, 99999, 99999 + + from PIL import Image + + with open(csv_filename, "w", encoding="utf-8", newline="") as file: + writer = csv.writer(file) + + writer.writerow(["image_path", "text_zh"]) + for filename in full_filenames: + if filename.endswith(".png"): + + image_path = os.path.join(workspace_images_dir, filename) + pil_image = Image.open(image_path) + width, height = pil_image.size + if width > max_width: + max_width = width + if height > max_height: + max_height = height + if width < min_width: + min_width = width + if height < min_height: + min_height = height + + text_zh = "" + + # 去除后缀 + _caption_filename = filename.split(".")[0] + caption_filename = "" + if os.path.exists(os.path.join(workspace_images_dir, _caption_filename + ".txt")): + caption_filename = _caption_filename + ".txt" + elif os.path.exists(os.path.join(workspace_images_dir, _caption_filename + ".caption")): + caption_filename = _caption_filename + ".caption" + else: + caption_filename = "" + if caption_filename != "": + with open(os.path.join(workspace_images_dir, caption_filename), "r") as f: + caption = f.read() + text_zh = caption + + writer.writerow([image_path, text_zh]) + + if resolution > max_width and resolution > max_height: + print("resolution 大于所有图片的宽高,强制使用最大宽高作为resolution") + resolution = max(max_width, max_height) // 8 * 8 + + print(f"csv文件已生成: {csv_filename}") + + HYDiT_tool_dir = os.path.join( + Utils.get_minus_zone_models_path(), "train_tools", "HunyuanDiT") + + # python ./hydit/data_loader/csv2arrow.py ./dataset/porcelain/csvfile/image_text.csv ./dataset/porcelain/arrows + arrows_dir = os.path.join(workspace_dir, "arrows") + os.makedirs(arrows_dir, exist_ok=True) + csv2arrow_exec = os.path.join( + HYDiT_tool_dir, "hydit", "data_loader", "csv2arrow.py") + try: + subprocess.run([ + sys.executable, csv2arrow_exec, csv_filename, arrows_dir + ], check=True) + except Exception as e: + raise Exception(f"生成arrow文件时出现异常,详细信息请查看控制台...") + + arrows_files = os.listdir(arrows_dir) + + dataset_yaml_path = os.path.join(workspace_dir, "dataset.yaml") + + import yaml + with open(dataset_yaml_path, "w", encoding="utf-8") as f: + yaml.dump({ + "source": [os.path.join(arrows_dir, arrow_file) for arrow_file in arrows_files], + }, f) + + print(f"dataset.yaml文件已生成: {dataset_yaml_path}") + + idk_exec = os.path.join( + os.path.dirname(__file__), "hook_HYDiT_idk_run.py") + + try: + from index_kits import __version__ + except Exception as e: + subprocess.run([ + sys.executable, "-m", "pip", "install", "-e", os.path.join( + HYDiT_tool_dir, "IndexKits") + ], check=True) + # idk base -c dataset/yamls/porcelain.yaml -t dataset/porcelain/jsons/porcelain.json + + dataset_json = os.path.join(workspace_dir, "dataset.json") + try: + subprocess.run([ + sys.executable, idk_exec, "base", "-c", dataset_yaml_path, "-t", dataset_json, + ], check=True) + except Exception as e: + raise Exception(f"生成index文件时出现异常,详细信息请查看控制台...") + + print(f"index文件已生成: {dataset_json}") + + dataset_mt_yaml_path = os.path.join(workspace_dir, "dataset_mt.yaml") + + with open(dataset_mt_yaml_path, "w", encoding="utf-8") as f: + yaml.dump({ + "src": [dataset_json], + "base_size": resolution, + "target_ratios": ["1:1", "3:4", "4:3", "16:9", "9:16"], + }, f) + + print(f"dataset_mt.yaml文件已生成: {dataset_mt_yaml_path}") + + # idk multireso -c dataset/yamls/porcelain_mt.yaml -t dataset/porcelain/jsons/porcelain_mt.json + + dataset_mt_json = os.path.join(workspace_dir, "dataset_mt.json") + try: + subprocess.run([ + sys.executable, idk_exec, "multireso", "-c", dataset_mt_yaml_path, "-t", dataset_mt_json, + ], check=True) + except Exception as e: + raise Exception(f"生成多分辨率index文件时出现异常,详细信息请查看控制台...") + + print(f"多分辨率index文件已生成: {dataset_mt_json}") + + train_config = {} + + train_config_path = os.path.join(workspace_dir, "mz_train_config.json") + + if os.path.exists(train_config_path): + with open(train_config_path, "r", encoding="utf-8") as f: + train_config = json.load(f) + +# "unet": (["auto"] + models + unet_models, {"default": "auto"}), +# "vae": (["auto"] + folders + vae_models, {"default": "auto"}), +# "text_encoder": (["auto"] + folders, {"default": "auto"}), +# "tokenizer": (["auto"] + folders, {"default": "auto"}), +# "t5_encoder": (["auto"] + folders, {"default": "auto"}), + + output_dir = os.path.join(workspace_dir, "output") + train_config.update({ + "batch_size": args.get("batch_size", 1), + "target_ratios": train_config.get("target_ratios", ["1:1", "3:4", "4:3", "16:9", "9:16"]), + "index_file": dataset_mt_json, + "epochs": args.get("epochs", 50), + "ckpt_every": args.get("save_every_n_epochs", 1), + "rope_img": f"base{resolution}", + "rope_real": args.get("rope_real", "enable") == "enable", + "image_size": resolution, + "unet_path": args.get("unet_path"), + "vae_ema_path": args.get("vae_ema_path"), + "text_encoder_path": args.get("text_encoder_path"), + "tokenizer_path": args.get("tokenizer_path"), + "t5_encoder_path": args.get("t5_encoder_path"), + "results_dir": output_dir, + "task_flag": args.get("task_flag", workspace_name), + }) + + with open(train_config_path, "w", encoding="utf-8") as f: + json.dump(train_config, f, indent=4) + + run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir) + return ( + args, + ) + + +def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir): + pb = Utils.progress_bar(train_config.get("epochs"), "sdxl") + import traceback + import comfy.model_management + stop_server = None + + is_running = True + + def log_callback(log): + try: + comfy.model_management.throw_exception_if_processing_interrupted() + except Exception as e: + stop_server() + return is_running + + try: + resp = log + if resp.get("type") == "sample_images": + total_steps = resp.get("total_steps") + global_step = resp.get("global_step") + xy_img = get_sample_images(train_config) + if xy_img is None: + pb.update( + int(global_step), int(total_steps), None) + return is_running + max_side = max(xy_img.width, xy_img.height) + + pb.update( + int(global_step), int(total_steps), ("JPEG", xy_img, max_side)) + else: + print(f"LOG: {log}") + except Exception as e: + print(f"LOG: {log} e: {e} ") + print(f"stack: {traceback.format_exc()}") + return is_running + stop_server, port = Utils.Simple_Server(log_callback) + + exec_pyfile = os.path.join(os.path.dirname( + __file__), "hook_HYDiT_run.py",) + + try: + subprocess.run( + [sys.executable, exec_pyfile, "--sys_path", HYDiT_tool_dir, + "--master_port", str(port), "--train_config_file", train_config_path], + check=True, + ) + stop_server() + is_running = False + except Exception as e: + stop_server() + is_running = False + raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") + + +def get_sample_images(train_config): + # output_name = train_config.get("output_name") + # sample_images_dir = os.path.join( + # os.path.dirname(train_config.get("dataset_config")), "sample_images" + # ) + # pil_images = [] + # pre_render_texts_x = [] + # if os.path.exists(sample_images_dir): + # image_files = os.listdir(sample_images_dir) + # image_files = list( + # filter(lambda x: x.endswith(".png"), image_files)) + # # 筛选 output_name 前缀 + # image_files = list( + # filter(lambda x: x.startswith(output_name), image_files)) + + # image_files = sorted(image_files, key=lambda x: x) + + # for image_file in image_files: + # pil_image = Image.open(os.path.join(sample_images_dir, image_file)) + # pil_images.append([pil_image]) + # pre_render_texts_x.append(image_file) + # result = Utils.xy_image( + # pre_render_images=pil_images, + # pre_render_texts_x=pre_render_texts_x, + # pre_render_texts_y=[""], + # ) + # return result + return None diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index 49ac5eb..c5345eb 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -913,3 +913,18 @@ class Utils: label_text, font=font, fill="black") return full_canvas + + def get_models_by_folder(dir_path): + models = [] + for root, dirs, files in os.walk(dir_path): + for file in files: + if file.endswith(".pth") or file.endswith(".pt") or file.endswith(".pkl") or file.endswith(".onnx") or file.endswith(".safetensors"): + models.append(os.path.join(root, file)) + return models + + def get_folders_by_folder(dir_path): + folders = [] + for root, dirs, files in os.walk(dir_path): + for dir in dirs: + folders.append(os.path.join(root, dir)) + return folders \ No newline at end of file