'update'
This commit is contained in:
wailovet
@@ -171,3 +171,4 @@ exclude.txt
|
||||
*.msgpack
|
||||
*.onnx_data
|
||||
configs/models_config/**/*.png
|
||||
HunyuanDiT
|
||||
+236
-8
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user