nodes/MinusZoneAI-ComfyUI-TrainTools-MZ
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

'update'

Browse Source
This commit is contained in:
wailovet
2024-06-27 23:01:02 +08:00
parent 740e059dd9
commit 464b3b5c67
5 changed files with 223 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2 -2
View File
@@ -58,7 +58,7 @@ class MZ_KohyaSSDatasetConfig:
"workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
"images": ("IMAGE",),
"enable_bucket": (["enable", "disable"], {"default": "enable"}),
"resolution": ("INT", {"default": 512}),
"resolution": ("INT", {"default": 1024}),
"num_repeats": ("INT", {"default": 1}),
"batch_size": ("INT", {"default": 1}),
"force_clear": (["enable", "disable"], {"default": "disable"}),
@@ -539,7 +539,7 @@ class MZ_HYDiTTrain:
"resolution": ("INT", {"default": 1024}),
"batch_size": ("INT", {"default": 1}),
"epochs": ("INT", {"default": 50}),
"save_every_n_epochs": ("INT", {"default": 10}),
"ckpt_every": ("INT", {"default": 10}),
"base_lora": (["latest", "empty"], {"default": "latest"}),
"sample_generate": (["enable", "disable"], {"default": "enable"}),
"sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}),
hook_HYDiT_main.py
+147 -14
View File
@@ -83,6 +83,14 @@ T5_ENCODER = {
}
global TRAIN_CONFIG
def set_train_config(train_config):
global TRAIN_CONFIG
TRAIN_CONFIG = train_config
def set_t5_encoder_path(path):
global T5_ENCODER
T5_ENCODER['MT5'] = path
@@ -135,12 +143,9 @@ def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoi
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)
else:
model.save_pretrained(cur_ckpt_save_dir)
checkpoint_path = "[Not rank 0. Disabled output.]"
@@ -183,12 +188,6 @@ def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoi
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
@@ -441,6 +440,7 @@ def Core(args, LOG):
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,
@@ -537,10 +537,15 @@ def Core(args, LOG):
optimizer = torch.optim.AdamW(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
target_height=768, target_width=1280, train_steps=0)
# 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}")
@@ -605,8 +610,9 @@ def Core(args, LOG):
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}")
pbar.set_description(
f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {mean_loss / step:.4f}")
# ===========================================================================
# Log loss values:
# ===========================================================================
@@ -640,9 +646,11 @@ def Core(args, LOG):
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:
) and train_steps > 0:
logger.info(
f" Saving checkpoint at step {train_steps}.")
easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5,
target_height=768, target_width=1280, train_steps=train_steps)
save_checkpoint(args, rank, logger, model, ema,
epoch, train_steps, checkpoint_dir)
@@ -708,3 +716,128 @@ def model_resume(args, model, ema, logger):
" “If `resume` is True, then either `resume_split` must be true.”")
return model, ema, start_epoch, start_epoch_step, train_steps
def easy_sample_images(
args,
vae=None,
text_encoder=None,
tokenizer=None,
model=None,
embedder_t5=None,
target_height=768,
target_width=1280,
prompt="A photo of a girl with a hat on a sunny day",
negative_prompt="",
batch_size=1,
guidance_scale=5.0,
infer_steps=20,
sampler='ddpm',
train_steps=0,
):
with torch.no_grad():
from hydit.diffusion.pipeline import StableDiffusionPipeline
from diffusers import schedulers
from hydit.constants import SAMPLER_FACTORY
from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed
from hydit.modules.models import HUNYUAN_DIT_CONFIG
workspace_dir = TRAIN_CONFIG.get("workspace_dir")
sample_config_file = TRAIN_CONFIG.get("sample_config_file", None)
if sample_config_file is None:
print("sample_config_file is not set.")
return
try:
sample_config = json.load(open(sample_config_file, "r"))
except Exception as e:
print(f"Failed to load sample_config_file: {sample_config_file}")
return
sample_images_dir = os.path.join(workspace_dir, "sample_images")
os.makedirs(sample_images_dir, exist_ok=True)
# Load sampler from factory
kwargs = SAMPLER_FACTORY[sampler]['kwargs']
scheduler = SAMPLER_FACTORY[sampler]['scheduler']
# Build scheduler according to the sampler.
scheduler_class = getattr(schedulers, scheduler)
scheduler = scheduler_class(**kwargs)
# Set timesteps for inference steps.
scheduler.set_timesteps(infer_steps, "cuda")
def calc_rope(height, width):
model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"]
patch_size = model_config['patch_size']
head_size = model_config['hidden_size'] // model_config['num_heads']
th = height // 8 // patch_size
tw = width // 8 // patch_size
base_size = 512 // 8 // patch_size
start, stop = get_fill_resize_and_crop((th, tw), base_size)
sub_args = [start, stop, (th, tw)]
rope = get_2d_rotary_pos_embed(head_size, *sub_args)
return rope
pipeline = StableDiffusionPipeline(vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=model,
scheduler=scheduler,
feature_extractor=None,
safety_checker=None,
requires_safety_checker=False,
embedder_t5=embedder_t5,
)
pipeline.to("cuda")
style = torch.as_tensor([0, 0] * batch_size, device="cuda")
src_size_cond = (target_width, target_height)
size_cond = list(src_size_cond) + [target_width, target_height, 0, 0]
image_meta_size = torch.as_tensor(
[size_cond] * 2 * batch_size, device="cuda",)
if type(sample_config) != list:
sample_config = [sample_config]
for i, sample in enumerate(sample_config):
prompt = sample.get("prompt", "")
negative_prompt = sample.get("negative_prompt", "")
guidance_scale = sample.get("cfg", guidance_scale)
infer_steps = sample.get("steps", infer_steps)
width = sample.get("width", target_width)
height = sample.get("height", target_height)
freqs_cis_img = calc_rope(height, width)
samples = pipeline(
height=height,
width=width,
prompt=prompt,
negative_prompt=negative_prompt,
num_images_per_prompt=batch_size,
guidance_scale=guidance_scale,
num_inference_steps=infer_steps,
style=style,
return_dict=False,
use_fp16=True,
learn_sigma=args.learn_sigma,
freqs_cis_img=freqs_cis_img,
image_meta_size=image_meta_size,
)[0]
# print("samples:",type(samples),)
# input("Press Enter to continue...")
# print("samples:",samples,)
if type(samples) == list:
pil_image = samples[0]
else:
pil_image = samples
sample_filename = f"{args.task_flag}_train_steps_{train_steps:07d}.png"
sample_filename_path = os.path.join(sample_images_dir, sample_filename)
pil_image.save(sample_filename_path)
return None
hook_HYDiT_run.py
+9 -6
View File
@@ -5,6 +5,8 @@ import os
import sys
from types import SimpleNamespace
import torch
# python -m torch.distributed.launch ./hook_HYDi_run.py --sys_path /data/ComfyUI/models/minus_zone_models/train_tools/HunyuanDiT/
@@ -87,11 +89,9 @@ if __name__ == "__main__":
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
@@ -109,8 +109,8 @@ if __name__ == "__main__":
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_img = train_config.get("rope_img", "base1024")
margs.image_size = train_config.get("image_size", 1024)
margs.rope_real = train_config.get("rope_real", True)
margs.index_file = train_config.get("index_file", None)
@@ -122,7 +122,8 @@ if __name__ == "__main__":
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")
margs.mse_loss_weight_type = train_config.get(
"mse_loss_weight_type", "constant")
for k, v in train_config.items():
if hasattr(margs, k):
@@ -143,4 +144,6 @@ if __name__ == "__main__":
hook_HYDiT_main.set_t5_encoder_path(
train_config.get("t5_encoder_path"))
hook_HYDiT_main.set_train_config(train_config)
hook_HYDiT_main.Core(margs, LOG)
mz_train_tools_core_HYDiT.py
+63 -30
View File
@@ -110,7 +110,7 @@ def MZ_HYDiTDatasetConfig_call(args={}):
pil_images = Utils.tensors2pil_list(images)
resolution = args.get("resolution", 512)
resolution = args.get("resolution", 1024)
if workspace_name is None or workspace_name == "":
raise Exception("训练名称不能为空(workspace_name is required)")
@@ -282,13 +282,20 @@ def MZ_HYDiTTrain_call(args={}):
# "tokenizer": (["auto"] + folders, {"default": "auto"}),
# "t5_encoder": (["auto"] + folders, {"default": "auto"}),
epochs = args.get("epochs", 50)
output_dir = os.path.join(workspace_dir, "output")
datetime = time.strftime("%Y%m%d_%H%M%S", time.localtime())
output_name = f"{workspace_name}_{datetime}"
train_config.update({
"workspace_name": workspace_name,
"workspace_dir": workspace_dir,
"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),
"epochs": epochs,
"ckpt_every": int(args.get("ckpt_every")),
"rope_img": f"base{resolution}",
"rope_real": args.get("rope_real", "enable") == "enable",
"image_size": resolution,
@@ -298,11 +305,31 @@ def MZ_HYDiTTrain_call(args={}):
"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),
"task_flag": output_name,
})
sample_generate = args.get("sample_generate", "enable")
sample_prompt = args.get("sample_prompt", "")
sample_config_file = os.path.join(workspace_dir, "sample_config.json")
if sample_generate == "enable":
resolution = args.get("resolution")
sample_config = {
"prompt": sample_prompt,
"negative_prompt": "",
"cfg": 5.0,
"steps": 20,
"width": resolution,
"height": resolution,
}
with open(sample_config_file, "w", encoding="utf-8") as f:
json.dump(sample_config, f, indent=4, ensure_ascii=False)
train_config["sample_config_file"] = sample_config_file
else:
train_config["sample_config_file"] = None
with open(train_config_path, "w", encoding="utf-8") as f:
json.dump(train_config, f, indent=4)
json.dump(train_config, f, indent=4, ensure_ascii=False)
run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir)
return (
@@ -311,7 +338,7 @@ def MZ_HYDiTTrain_call(args={}):
def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir):
pb = Utils.progress_bar(train_config.get("epochs"), "sdxl")
pb = Utils.progress_bar(0, "sdxl")
import traceback
import comfy.model_management
stop_server = None
@@ -365,30 +392,36 @@ def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir):
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))
if train_config.get("sample_config_file") is None:
return None
# image_files = sorted(image_files, key=lambda x: x)
from PIL import Image
output_name = train_config.get("task_flag")
# 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
workspace_dir = train_config.get("workspace_dir")
sample_images_dir = os.path.join(workspace_dir, "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
mz_train_tools_utils.py
+2
View File
@@ -787,6 +787,8 @@ class Utils:
# 去掉pre_render_texts_x中所有相同字符串前缀
def get_common_prefix(pre_render_texts_x):
if len(pre_render_texts_x) == 0:
return ""
common_prefix = ""
for i in range(len(pre_render_texts_x[0])):
c = pre_render_texts_x[0][i]