From 1b6ff2d4575fb3923db74374549fb5e2edf80852 Mon Sep 17 00:00:00 2001 From: wailovet Date: Sat, 6 Jul 2024 18:42:28 +0800 Subject: [PATCH 01/10] 'update' --- mz_train_tools_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index da336c8..b52a7b6 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -758,6 +758,7 @@ class Utils: step, total_steps, pil_img_info) except Exception as e: print("progress_bar:", e) + raise e return pb(steps) From 5e297aefe637a7e4f06d11fe5b9c2cd0252addd8 Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 7 Jul 2024 16:58:13 +0800 Subject: [PATCH 02/10] 'update' --- mz_train_tools_core.py | 2 +- tools/face_resize_copy_dir2.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 80fa5ea..06d3194 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -129,7 +129,7 @@ def MZ_ImageSelecter_call(args={}): if force_clear_only_images: images_files = Utils.listdir(train_images_dir) for file in images_files: - if file.lower().endswith(".png") or file.lower().endswith(".jpg"): + if file.lower().endswith(".png") or file.lower().endswith(".jpg") or file.lower().endswith(".webp"): os.remove(os.path.join(train_images_dir, file)) else: shutil.rmtree(train_images_dir) diff --git a/tools/face_resize_copy_dir2.py b/tools/face_resize_copy_dir2.py index 7a78823..79c1253 100644 --- a/tools/face_resize_copy_dir2.py +++ b/tools/face_resize_copy_dir2.py @@ -100,7 +100,8 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("-i", "--input_dir", type=str, default="") parser.add_argument("-o", "--output_dir", type=str, default="") - parser.add_argument("-m", "--face_model_path", type=str, default="") + parser.add_argument("-m", "--face_model_path", type=str, + default="D:\\数据集\\model.onnx") parser.add_argument("-r", "--resize", type=int, default=1024) parser.add_argument("-p", "--face_proportion", type=float, default=0.5) parser.add_argument("-f", "--format", type=str, default="jpg") @@ -111,7 +112,7 @@ def main(): input_dir = args.input_dir output_dir = args.output_dir resize = args.resize - img_format = args.format + img_format = args.format total = 0 for root, dirs, files in os.walk(input_dir): # 排除隐藏文件夹 @@ -129,12 +130,15 @@ def main(): # 排除隐藏文件夹 dirs[:] = [d for d in dirs if not d.startswith('.')] + for root, dirs, files in os.walk(input_dir): for file in files: if not file.endswith(".png") and not file.endswith(".jpg"): continue - file_path = os.path.join(root, file) - pbar.update(1) + + pbar.update(1) + + file_path = os.path.join(root, file) with open(file_path, "rb") as f: md5_str = Utils.Md5(f.read()) From a56d99d1dd2f44f924d0e261ace1accea25f14db Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 7 Jul 2024 21:28:33 +0800 Subject: [PATCH 03/10] 'update' --- __init__.py | 3 +- mz_train_tools_core.py | 5 +- mz_train_tools_core_HYDiT.py | 6 +- tools/aesthetic.py | 210 -------------------------- tools/aesthetic2_5.py | 260 --------------------------------- tools/face_resize_copy_dir.py | 200 ------------------------- tools/face_resize_copy_dir2.py | 249 ------------------------------- 7 files changed, 9 insertions(+), 924 deletions(-) delete mode 100644 tools/aesthetic.py delete mode 100644 tools/aesthetic2_5.py delete mode 100644 tools/face_resize_copy_dir.py delete mode 100644 tools/face_resize_copy_dir2.py diff --git a/__init__.py b/__init__.py index 852d390..1b73304 100644 --- a/__init__.py +++ b/__init__.py @@ -64,6 +64,7 @@ class MZ_KohyaSSDatasetConfig: "force_clear_only_images": (["enable", "disable"], {"default": "disable"}), "same_caption_generate": (["enable", "disable"], {"default": "disable"}), "same_caption": ("STRING", {"default": "", "dynamicPrompts": True, "multiline": True}), + "image_format": (["png", "jpg", "webp"], {"default": "webp"}), }, "optional": { "conditioning_images": ("IMAGE",), @@ -632,7 +633,7 @@ class MZ_LoadImagesFromDirectoryPath: images = Utils.listdir(image_dir) images = [x for x in images if x.lower().endswith( - ".png") or x.lower().endswith(".jpg")] + ".png") or x.lower().endswith(".jpg") or file.lower().endswith(".webp")] images = [os.path.join(image_dir, x) for x in images] pil_images = [] diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 06d3194..a14ff3e 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -135,11 +135,14 @@ def MZ_ImageSelecter_call(args={}): shutil.rmtree(train_images_dir) os.makedirs(train_images_dir, exist_ok=True) + image_format = args.get("image_format") + file_extension = "." + image_format 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" + filename = hashlib.md5( + pil_image.tobytes()).hexdigest() + file_extension pil_image.save(os.path.join(train_images_dir, filename)) saved_images_path.append(filename) diff --git a/mz_train_tools_core_HYDiT.py b/mz_train_tools_core_HYDiT.py index f3427fe..4147e61 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -126,7 +126,7 @@ def MZ_HYDiTDatasetConfig_call(args={}): if same_caption != "": # 循环已经保存的图片 for i, filename in enumerate(saved_images_path): - if filename.lower().endswith(".png") or filename.lower().endswith(".jpg"): + if filename.lower().endswith(".png") or filename.lower().endswith(".jpg") or file.lower().endswith(".webp"): base_filename = os.path.splitext(filename)[0] caption_filename = base_filename + ".caption" with open(os.path.join(train_images_dir, caption_filename), "w", encoding="utf-8") as f: @@ -150,7 +150,7 @@ def MZ_HYDiTDatasetConfig_call(args={}): if force_clear_only_images: images_files = Utils.listdir(train_images_dir) for file in images_files: - if file.lower().endswith(".png") or file.lower().endswith(".jpg"): + if file.lower().endswith(".png") or file.lower().endswith(".jpg") or file.lower().endswith(".webp"): os.remove(os.path.join(train_images_dir, file)) else: shutil.rmtree(train_images_dir) @@ -427,7 +427,7 @@ def MZ_HYDiTTrain_call(args={}): writer.writerow(["image_path", "text_zh"]) for filename in full_filenames: print(f"处理文件: {filename}") - if filename.lower().endswith(".png") or filename.lower().endswith(".jpg"): + if filename.lower().endswith(".png") or filename.lower().endswith(".jpg") or file.lower().endswith(".webp"): image_path = os.path.join(workspace_images_dir, filename) pil_image = Image.open(image_path) diff --git a/tools/aesthetic.py b/tools/aesthetic.py deleted file mode 100644 index ee3891d..0000000 --- a/tools/aesthetic.py +++ /dev/null @@ -1,210 +0,0 @@ -import json -import os -import shutil -from PIL import Image -import matplotlib.pyplot as plt - -from warnings import filterwarnings - -import numpy as np -import torch -import pytorch_lightning as pl -import torch.nn as nn -from PIL import Image - - -from transformers import CLIPProcessor, CLIPModel -import torch.nn.functional as F - - -class MLP(pl.LightningModule): - def __init__(self, input_size, xcol='emb', ycol='avg_rating'): - super().__init__() - self.input_size = input_size - self.xcol = xcol - self.ycol = ycol - self.layers = nn.Sequential( - nn.Linear(self.input_size, 1024), - # nn.ReLU(), - nn.Dropout(0.2), - nn.Linear(1024, 128), - # nn.ReLU(), - nn.Dropout(0.2), - nn.Linear(128, 64), - # nn.ReLU(), - nn.Dropout(0.1), - - nn.Linear(64, 16), - # nn.ReLU(), - - nn.Linear(16, 1) - ) - - def forward(self, x): - return self.layers(x) - - def training_step(self, batch, batch_idx): - x = batch[self.xcol] - y = batch[self.ycol].reshape(-1, 1) - x_hat = self.layers(x) - loss = F.mse_loss(x_hat, y) - return loss - - def validation_step(self, batch, batch_idx): - x = batch[self.xcol] - y = batch[self.ycol].reshape(-1, 1) - x_hat = self.layers(x) - loss = F.mse_loss(x_hat, y) - return loss - - def configure_optimizers(self): - optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) - return optimizer - - -def normalized(a, axis=-1, order=2): - import numpy as np - - l2 = np.atleast_1d(np.linalg.norm(a, order, axis)) - l2[l2 == 0] = 1 - return a / np.expand_dims(l2, axis) - - -class Aestheic: - def __init__(self, clip_path, aes_model): - - model = MLP(768) # CLIP embedding dim is 768 for CLIP ViT L 14 - - # load the model you trained previously or the model available in this repo - s = torch.load(aes_model) - - model.load_state_dict(s) - - model.to("cuda") - model.eval() - - self.model = model - - self.device = "cuda" - - model2 = CLIPModel.from_pretrained( - clip_path, local_files_only=True).to("cuda") - preprocess = CLIPProcessor.from_pretrained( - clip_path, local_files_only=True) - - self.model2 = model2 - self.preprocess = preprocess - - def prediction(self, img_paths): - if not isinstance(img_paths, list): - img_paths = [img_paths] - - pil_images = [] - for img_path in img_paths: - pil_image = Image.open(img_path) - pil_images.append(pil_image) - - inputs = self.preprocess( - images=pil_images, return_tensors="pt", - ).to("cuda") - with torch.no_grad(): - image_features = self.model2.get_image_features(**inputs) - print(image_features.shape) - - im_emb_arr = normalized(image_features.cpu().detach().numpy()) - - prediction = self.model(torch.from_numpy( - im_emb_arr).to(self.device).type(torch.cuda.FloatTensor)) - - prediction_score = prediction.cpu().detach().numpy() - prediction_score = prediction_score.flatten() - - images_score = { - img_path: float(score) for img_path, score in zip(img_paths, prediction_score) - } - - # 按value排序 - images_score = dict( - sorted(images_score.items(), key=lambda x: x[1], reverse=True)) - - return images_score - - -import argparse -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-i", "--input_dir", type=str, default="") - parser.add_argument("-o", "--output_dir", type=str, default="") - parser.add_argument("-c", "--clip_path", type=str, - default="D:\\数据集\\clip-vit-large-patch14") - parser.add_argument("-m", "--aes_model", type=str, - default="D:\\数据集\\sac+logos+ava1-l14-linearMSE.pth") - parser.add_argument("-s", "--score", type=float, default=5.9) - - parser.add_argument("-b", "--batch_size", type=int, default=1) - - args = parser.parse_args() - limitscore = args.score - - aes = Aestheic(args.clip_path, args.aes_model) - - input_dir = args.input_dir - output_dir = args.output_dir - batch_size = args.batch_size - total = 0 - for root, dirs, files in os.walk(input_dir): - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - total += 1 - - import tqdm - print(f"total: {total}") - - images = [] - - with tqdm.tqdm(total=total) as pbar: - for root, dirs, files in os.walk(input_dir): - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - file_path = os.path.join(root, file) - pbar.update(1) - - pbar.set_description( - f"Processing {os.path.basename(file_path)}") - - if len(images) < batch_size: - images.append(file_path) - else: - image_scroe_map = aes.prediction(images) - images = [] - # print(scroe) - # print(json.dumps(scroe, indent=4, ensure_ascii=False)) - # exit(0) - - for img_path, score in image_scroe_map.items(): - print(f"{img_path}: {score}") - - if score > limitscore: - if not os.path.exists(output_dir): - os.makedirs(output_dir) - shutil.copy(img_path, output_dir) - - if len(images) > 0: - aes.prediction(images) - images = [] - - for img_path, score in image_scroe_map.items(): - print(f"{img_path}: {score}") - - if score > limitscore: - if not os.path.exists(output_dir): - os.makedirs(output_dir) - shutil.copy(img_path, output_dir) - - diff --git a/tools/aesthetic2_5.py b/tools/aesthetic2_5.py deleted file mode 100644 index d51f377..0000000 --- a/tools/aesthetic2_5.py +++ /dev/null @@ -1,260 +0,0 @@ -import json -import os -from collections import OrderedDict -from os import PathLike -import shutil -from typing import Final - -import torch -import torch.nn as nn -from transformers import ( - SiglipImageProcessor, - SiglipVisionConfig, - SiglipVisionModel, - logging, -) -from transformers.image_processing_utils import BatchFeature -from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention - -logging.set_verbosity_error() - -URL: Final[str] = ( - "https://github.com/discus0434/aesthetic-predictor-v2-5/raw/main/models/aesthetic_predictor_v2_5.pth" -) - - -class AestheticPredictorV2_5Head(nn.Module): - def __init__(self, config: SiglipVisionConfig) -> None: - super().__init__() - self.scoring_head = nn.Sequential( - nn.Linear(config.hidden_size, 1024), - nn.Dropout(0.5), - nn.Linear(1024, 128), - nn.Dropout(0.5), - nn.Linear(128, 64), - nn.Dropout(0.5), - nn.Linear(64, 16), - nn.Dropout(0.2), - nn.Linear(16, 1), - ) - - def forward(self, image_embeds: torch.Tensor) -> torch.Tensor: - return self.scoring_head(image_embeds) - - -class AestheticPredictorV2_5Model(SiglipVisionModel): - PATCH_SIZE = 14 - - def __init__(self, config: SiglipVisionConfig, *args, **kwargs) -> None: - super().__init__(config, *args, **kwargs) - self.layers = AestheticPredictorV2_5Head(config) - self.post_init() - - def forward( - self, - pixel_values: torch.FloatTensor | None = None, - labels: torch.Tensor | None = None, - return_dict: bool | None = None, - ) -> tuple | ImageClassifierOutputWithNoAttention: - return_dict = ( - return_dict if return_dict is not None else self.config.use_return_dict - ) - - outputs = super().forward( - pixel_values=pixel_values, - return_dict=return_dict, - ) - image_embeds = outputs.pooler_output - image_embeds_norm = image_embeds / \ - image_embeds.norm(dim=-1, keepdim=True) - prediction = self.layers(image_embeds_norm) - - loss = None - if labels is not None: - loss_fct = nn.MSELoss() - loss = loss_fct() - - if not return_dict: - return (loss, prediction, image_embeds) - - return ImageClassifierOutputWithNoAttention( - loss=loss, - logits=prediction, - hidden_states=image_embeds, - ) - - -class AestheticPredictorV2_5Processor(SiglipImageProcessor): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - - def __call__(self, *args, **kwargs) -> BatchFeature: - return super().__call__(*args, **kwargs) - - @classmethod - def from_pretrained( - self, - pretrained_model_name_or_path: str - | PathLike = "google/siglip-so400m-patch14-384", - *args, - **kwargs, - ) -> "AestheticPredictorV2_5Processor": - return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) - - -def convert_v2_5_from_siglip( - predictor_name_or_path: str | PathLike | None = None, - encoder_model_name: str = "google/siglip-so400m-patch14-384", - *args, - **kwargs, -) -> tuple[AestheticPredictorV2_5Model, AestheticPredictorV2_5Processor]: - model = AestheticPredictorV2_5Model.from_pretrained( - encoder_model_name, *args, **kwargs - ) - - processor = AestheticPredictorV2_5Processor.from_pretrained( - encoder_model_name, *args, **kwargs - ) - - if predictor_name_or_path is None or not os.path.exists(predictor_name_or_path): - state_dict = torch.hub.load_state_dict_from_url( - URL, map_location="cpu") - else: - state_dict = torch.load(predictor_name_or_path, map_location="cpu") - - assert isinstance(state_dict, OrderedDict) - - model.layers.load_state_dict(state_dict) - model.eval() - - return model, processor - - -from PIL import Image - - -class AestheticPredictor: - def __init__(self, clip_path, aesthetic_predictor_path): - # load model and preprocessor - self.model, self.preprocessor = convert_v2_5_from_siglip( - low_cpu_mem_usage=True, - trust_remote_code=True, - encoder_model_name=clip_path, - predictor_name_or_path=aesthetic_predictor_path, - ) - if torch.cuda.is_available(): - self.model = self.model.to(torch.bfloat16).cuda() - - def inference(self, img_paths) -> float: - if not isinstance(img_paths, list): - img_paths = [img_paths] - - pil_images = [] - for img_path in img_paths: - pil_image = Image.open(img_path) - pil_images.append(pil_image) - - # preprocess image - pixel_values = self.preprocessor( - images=pil_images, return_tensors="pt" - ).pixel_values - - if torch.cuda.is_available(): - pixel_values = pixel_values.to(torch.bfloat16).cuda() - - # predict aesthetic score - with torch.inference_mode(): - score = self.model( - pixel_values).logits.squeeze().float().cpu().numpy() - prediction_score = score.flatten() - - images_score = { - img_path: float(score) for img_path, score in zip(img_paths, prediction_score) - } - - # 按value排序 - images_score = dict( - sorted(images_score.items(), key=lambda x: x[1], reverse=True)) - - return images_score - - - -import argparse -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-i", "--input_dir", type=str, default="") - parser.add_argument("-o", "--output_dir", type=str, default="") - parser.add_argument("-c", "--clip_path", type=str, - default="D:\\数据集\\siglip-so400m-patch14-384") - parser.add_argument("-m", "--aes_model", type=str, - default="D:\\数据集\\aesthetic_predictor_v2_5.pth") - parser.add_argument("-s", "--score", type=float, default=5.9) - - parser.add_argument("-b", "--batch_size", type=int, default=1) - - args = parser.parse_args() - limitscore = args.score - - aes = AestheticPredictor(args.clip_path, args.aes_model) - - input_dir = args.input_dir - output_dir = args.output_dir - batch_size = args.batch_size - total = 0 - for root, dirs, files in os.walk(input_dir): - - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - total += 1 - - import tqdm - print(f"total: {total}") - - images = [] - - with tqdm.tqdm(total=total) as pbar: - for root, dirs, files in os.walk(input_dir): - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - file_path = os.path.join(root, file) - pbar.update(1) - - pbar.set_description( - f"Processing {os.path.basename(file_path)}") - - if os.path.exists(os.path.join(output_dir, os.path.basename(file_path))): - continue - - if len(images) < batch_size: - images.append(file_path) - else: - image_scroe_map = aes.inference(images) - images = [] - for img_path, score in image_scroe_map.items(): - print(f"{img_path}: {score}") - - if score > limitscore: - if not os.path.exists(output_dir): - os.makedirs(output_dir) - shutil.copy(img_path, output_dir) - - if len(images) > 0: - image_scroe_map = aes.inference(images) - images = [] - - for img_path, score in image_scroe_map.items(): - print(f"{img_path}: {score}") - - if score > limitscore: - if not os.path.exists(output_dir): - os.makedirs(output_dir) - shutil.copy(img_path, output_dir) - - diff --git a/tools/face_resize_copy_dir.py b/tools/face_resize_copy_dir.py deleted file mode 100644 index 03eacf5..0000000 --- a/tools/face_resize_copy_dir.py +++ /dev/null @@ -1,200 +0,0 @@ -import argparse -import hashlib -import io -import os -from PIL import Image -import numpy as np -import onnxruntime as ort -import cv2 - - -class Utils: - def Md5(bytes): - return hashlib.md5(bytes).hexdigest() - - def resize_by(image, percent): - # 判断类型是否为PIL - if not isinstance(image, Image.Image): - image = Image.fromarray(image) - - width, height = image.size - new_width = int(width * percent) - new_height = int(height * percent) - return image.resize((new_width, new_height), Image.LANCZOS) - - def resize_max(im, dst_w, dst_h): - src_w, src_h = im.size - - if src_h > src_w: - newWidth = dst_w - newHeight = dst_w * src_h // src_w - else: - newWidth = dst_h * src_w // src_h - newHeight = dst_h - - newHeight = newHeight // 8 * 8 - newWidth = newWidth // 8 * 8 - - return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) - - def resize_min(im, dst_w, dst_h): - src_w, src_h = im.size - - if src_h < src_w: - newWidth = dst_w - newHeight = dst_w * src_h // src_w - else: - newWidth = dst_h * src_w // src_h - newHeight = dst_h - - newHeight = newHeight // 8 * 8 - newWidth = newWidth // 8 * 8 - - return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) - - # 根据人脸进行剪裁 - def resize_crop(im, dst_w, dst_h, focus_point): - - src_w, src_h = im.size - - if src_h > src_w: - newWidth = dst_w - newHeight = dst_w * src_h // src_w - else: - newWidth = dst_h * src_w // src_h - newHeight = dst_h - - newHeight = int(newHeight // 8 * 8) - newWidth = int(newWidth // 8 * 8) - - im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) - - left = (newWidth - dst_w) // 2 - top = (newHeight - dst_h) // 2 - - right = left + dst_w - bottom = top + dst_h - - focus_point_x, focus_point_y = focus_point - offset_w = dst_w // 8 - offset_h = dst_h // 8 - if focus_point_x < left + dst_w // 8: - left = 0 - right = dst_w - elif focus_point_x > right - dst_w // 8: - right = newWidth - left = newWidth - dst_w - if focus_point_y < top + dst_h // 8: - top = 0 - bottom = dst_h - elif focus_point_y > bottom - dst_h // 8: - bottom = newHeight - top = newHeight - dst_h - - return im.crop((left, top, right, bottom)) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("-i", "--input_dir", type=str, default="") - parser.add_argument("-o", "--output_dir", type=str, default="") - parser.add_argument("-m", "--face_model_path", type=str, default="") - parser.add_argument("-r", "--resize", type=int, default=1024) - parser.add_argument("-s", "--score", type=float, default=0.8) - parser.add_argument("-f", "--format", type=str, default="jpg") - - args = parser.parse_args() - limitscore = args.score - sess = ort.InferenceSession(args.face_model_path) - sess.set_providers(['CUDAExecutionProvider']) - input_dir = args.input_dir - output_dir = args.output_dir - resize = args.resize - img_format = args.format - total = 0 - for root, dirs, files in os.walk(input_dir): - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - total += 1 - - import tqdm - print(f"total: {total}") - - os.makedirs(output_dir, exist_ok=True) - with tqdm.tqdm(total=total) as pbar: - - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for root, dirs, files in os.walk(input_dir): - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - file_path = os.path.join(root, file) - pbar.update(1) - - with open(file_path, "rb") as f: - md5_str = Utils.Md5(f.read()) - if os.path.exists(os.path.join(output_dir, md5_str + ".png")): - continue - - image = Image.open(file_path) - image = Utils.resize_min(image, resize, resize) - - cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) - - print(f"start check {file_path}") - scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( - None, {"input": cv2_img}) - if len(bboxes) != 1: - continue - - bboxes = bboxes[0] - scores = scores[0] - if limitscore > scores: - continue - - affine_matrices = affine_matrices[0] - angle = np.degrees(np.arctan2( - affine_matrices[1][0], affine_matrices[0][0])) - - real_angle = 0 - if angle < -45 and angle > -135: - # 向左旋转90度 - real_angle = 90 - elif angle > 45 and angle < 135: - # 向右旋转90度 - real_angle = -90 - elif angle < -135 or angle > 135: - # 向左旋转180度 - real_angle = 180 - - if real_angle != 0: - continue - - # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}") - - center_point = (bboxes[0] + bboxes[2]) / \ - 2, (bboxes[1] + bboxes[3]) / 2 - image = Utils.resize_crop( - image, resize * 3 / 4, resize, center_point) - cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) - scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( - None, {"input": cv2_img}) - if len(bboxes) != 1: - continue - # cv2.imshow("image", cv2.cvtColor( - # np.array(image), cv2.COLOR_RGB2BGR)) - # cv2.waitKey(0) - if img_format == "png": - output_path = os.path.join(output_dir, md5_str + ".png") - image.save(output_path) - else: - output_path = os.path.join(output_dir, md5_str + ".jpg") - image.save(output_path, "jpeg", quality=95) - - -if __name__ == "__main__": - main() diff --git a/tools/face_resize_copy_dir2.py b/tools/face_resize_copy_dir2.py deleted file mode 100644 index 79c1253..0000000 --- a/tools/face_resize_copy_dir2.py +++ /dev/null @@ -1,249 +0,0 @@ -import argparse -import hashlib -import io -import os -from PIL import Image -import numpy as np -import onnxruntime as ort -import cv2 - -import piexif - - -class Utils: - def Md5(bytes): - return hashlib.md5(bytes).hexdigest() - - def resize_by(image, percent): - # 判断类型是否为PIL - if not isinstance(image, Image.Image): - image = Image.fromarray(image) - - width, height = image.size - new_width = int(width * percent) - new_height = int(height * percent) - return image.resize((new_width, new_height), Image.LANCZOS) - - def resize_max(im, dst_w, dst_h): - src_w, src_h = im.size - - if src_h > src_w: - newWidth = dst_w - newHeight = dst_w * src_h // src_w - else: - newWidth = dst_h * src_w // src_h - newHeight = dst_h - - newHeight = newHeight // 8 * 8 - newWidth = newWidth // 8 * 8 - - return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) - - def resize_min(im, dst_w, dst_h): - src_w, src_h = im.size - - if src_h < src_w: - newWidth = dst_w - newHeight = dst_w * src_h // src_w - else: - newWidth = dst_h * src_w // src_h - newHeight = dst_h - - newHeight = newHeight // 8 * 8 - newWidth = newWidth // 8 * 8 - - return im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) - - # 根据人脸进行剪裁 - def resize_crop(im, dst_w, dst_h, focus_point): - - src_w, src_h = im.size - - if src_h > src_w: - newWidth = dst_w - newHeight = dst_w * src_h // src_w - else: - newWidth = dst_h * src_w // src_h - newHeight = dst_h - - newHeight = int(newHeight // 8 * 8) - newWidth = int(newWidth // 8 * 8) - - im = im.resize((newWidth, newHeight), Image.Resampling.LANCZOS) - - left = (newWidth - dst_w) // 2 - top = (newHeight - dst_h) // 2 - - right = left + dst_w - bottom = top + dst_h - - focus_point_x, focus_point_y = focus_point - offset_w = dst_w // 8 - offset_h = dst_h // 8 - if focus_point_x < left + dst_w // 8: - left = 0 - right = dst_w - elif focus_point_x > right - dst_w // 8: - right = newWidth - left = newWidth - dst_w - if focus_point_y < top + dst_h // 8: - top = 0 - bottom = dst_h - elif focus_point_y > bottom - dst_h // 8: - bottom = newHeight - top = newHeight - dst_h - - return im.crop((left, top, right, bottom)) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("-i", "--input_dir", type=str, default="") - parser.add_argument("-o", "--output_dir", type=str, default="") - parser.add_argument("-m", "--face_model_path", type=str, - default="D:\\数据集\\model.onnx") - parser.add_argument("-r", "--resize", type=int, default=1024) - parser.add_argument("-p", "--face_proportion", type=float, default=0.5) - parser.add_argument("-f", "--format", type=str, default="jpg") - - args = parser.parse_args() - sess = ort.InferenceSession(args.face_model_path) - sess.set_providers(['CUDAExecutionProvider']) - input_dir = args.input_dir - output_dir = args.output_dir - resize = args.resize - img_format = args.format - total = 0 - for root, dirs, files in os.walk(input_dir): - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - total += 1 - - import tqdm - print(f"total: {total}") - - os.makedirs(output_dir, exist_ok=True) - with tqdm.tqdm(total=total) as pbar: - - # 排除隐藏文件夹 - dirs[:] = [d for d in dirs if not d.startswith('.')] - - for root, dirs, files in os.walk(input_dir): - for file in files: - if not file.endswith(".png") and not file.endswith(".jpg"): - continue - - pbar.update(1) - - file_path = os.path.join(root, file) - - with open(file_path, "rb") as f: - md5_str = Utils.Md5(f.read()) - if os.path.exists(os.path.join(output_dir, md5_str + ".png")): - continue - - image = Image.open(file_path) - max_image = image.copy() - image = Utils.resize_min(image, 1920, 1920) - zoom = image.size[0] / max_image.size[0] - - cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) - - print(f"start check {file_path}") - scores, bboxes, keypoints, aligned_imgs, landmarks, affine_matrices = sess.run( - None, {"input": cv2_img}) - if len(bboxes) != 1: - continue - - bboxes = bboxes[0] - - affine_matrices = affine_matrices[0] - angle = np.degrees(np.arctan2( - affine_matrices[1][0], affine_matrices[0][0])) - - real_angle = 0 - if angle < -45 and angle > -135: - # 向左旋转90度 - real_angle = 90 - elif angle > 45 and angle < 135: - # 向右旋转90度 - real_angle = -90 - elif angle < -135 or angle > 135: - # 向左旋转180度 - real_angle = 180 - - if real_angle != 0: - continue - - # print(f"{file_path} bboxes: \n{bboxes} {scores} {real_angle}") - - bboxes = bboxes / zoom - - image = max_image - - # 预期人脸高度占图片高度的一定比例 - pr = args.face_proportion - face_width = bboxes[2] - bboxes[0] - pre_image_width = face_width / pr - - face_center_x = (bboxes[2] + bboxes[0]) / 2 - crop_x1 = face_center_x - pre_image_width / 2 - crop_x2 = face_center_x + pre_image_width / 2 - - # 判断是否超出图片边界 - if crop_x1 < 0 or crop_x2 > image.width: - print( - f"crop_x1: {crop_x1}, crop_x2: {crop_x2}, image.width: {image.width}") - continue - - face_center_y = (bboxes[3] + bboxes[1]) / 2 - # 总高度和宽度的比例是4:3 - pre_image_height = pre_image_width * 4 / 3 - - # 中点位于高度的0.382比例 - crop_y1 = face_center_y - pre_image_height * 0.382 - crop_y2 = crop_y1 + pre_image_height - - # 判断是否超出图片边界 - if crop_y1 < 0 or crop_y2 > image.height: - print( - f"crop_y1: {crop_y1}, crop_y2: {crop_y2}, image.height: {image.height}") - continue - - # 判断是否比resize小 - if pre_image_width < resize or pre_image_height < resize: - print( - f"pre_image_width: {pre_image_width}, pre_image_height: {pre_image_height}") - continue - - # 进行裁剪 - image = image.crop((crop_x1, crop_y1, crop_x2, crop_y2)) - - # 进行缩放并居中剪裁 - image = image.resize(( - resize * 3 // 4, resize - ), Image.LANCZOS) - - # cv2_img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) - # cv2.imshow("image", cv2_img) - # cv2.waitKey(0) - - if img_format == "png": - output_path = os.path.join(output_dir, md5_str + ".png") - image.save(output_path) - else: - output_path = os.path.join(output_dir, md5_str + ".jpg") - image.save(output_path, "jpeg", quality=95) - # 插入IPTC - # piexif.insert(output_path, piexif.ImageIFD, { - # "CameraOwnerName": "CameraOwnerName", - # }) - - print(f"save to {output_path}") - - -if __name__ == "__main__": - main() From d63d8dcf7bc8dccc7214790718981ce96fd26763 Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 7 Jul 2024 23:07:40 +0800 Subject: [PATCH 04/10] 'update' --- mz_train_tools_core.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index a14ff3e..b2abceb 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -37,7 +37,9 @@ def MZ_KohyaSSCloneRepo_call(args={}): 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", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT,) # 切换远程分支 git remote set-branches origin 'main' branch = args.get("branch", "main") @@ -443,6 +445,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], check=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, ) stop_server() is_running = False From c3cae3115d09108e618153c74016e8a1e4040ca6 Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 7 Jul 2024 23:33:39 +0800 Subject: [PATCH 05/10] 'update' --- mz_train_tools_core.py | 41 +++++++++++++++++++++++++---------------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index b2abceb..b460c03 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -2,6 +2,7 @@ import argparse import hashlib +import io import os import shutil import sys @@ -408,8 +409,7 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type) - import traceback - + import traceback import comfy.model_management stop_server = None @@ -440,21 +440,30 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot print(f"stack: {traceback.format_exc()}") return is_running stop_server, port = Utils.Simple_Server(log_callback) - try: - subprocess.run( - [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, - "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], - check=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - stop_server() - is_running = False - except Exception as e: - stop_server() - is_running = False - raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") + + process = subprocess.Popen( + [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, + "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + ) + stdout, stderr = process.communicate() + stdout_str = "" + stderr_str = "" + if stdout is not None: + stdout_str = stdout.decode("utf-8") + if stderr is not None: + stderr_str = stderr.decode("utf-8") + retcode = process.poll() + if retcode: + raise Exception(f"""请查看以上具体信息, +执行失败: +========================stdout======================== +{stdout_str} +========================stderr======================== +{stderr_str} +====================================================== +""") def MZ_KohyaSSTrain_call(args={}): args = args.copy() From 2d7d3e419f6f70b7517d4ef86c69cf54abe3b55a Mon Sep 17 00:00:00 2001 From: wailovet Date: Sun, 7 Jul 2024 23:37:54 +0800 Subject: [PATCH 06/10] 'update' --- mz_train_tools_core.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index b460c03..814d057 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -452,6 +452,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot stderr_str = "" if stdout is not None: stdout_str = stdout.decode("utf-8") + # 清除制表符 + stdout_str = stdout_str.replace("\t", " ") if stderr is not None: stderr_str = stderr.decode("utf-8") retcode = process.poll() From 84f9a2c9abdeb1be81e9182138e5673efe511c87 Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 8 Jul 2024 00:23:33 +0800 Subject: [PATCH 07/10] 'update' --- mz_train_tools_core.py | 59 ++++++++++++++++++++++-------------------- 1 file changed, 31 insertions(+), 28 deletions(-) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 814d057..0612e56 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -2,7 +2,6 @@ import argparse import hashlib -import io import os import shutil import sys @@ -38,9 +37,7 @@ def MZ_KohyaSSCloneRepo_call(args={}): 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, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT,) + ["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True) # 切换远程分支 git remote set-branches origin 'main' branch = args.get("branch", "main") @@ -409,7 +406,8 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type) - import traceback + import traceback + import comfy.model_management stop_server = None @@ -440,32 +438,37 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot print(f"stack: {traceback.format_exc()}") return is_running stop_server, port = Utils.Simple_Server(log_callback) + try: + subprocess.run( + [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, + "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], + check=True, + ) + stop_server() + is_running = False + except subprocess.CalledProcessError as e: + stop_server() + is_running = False + stdout_str = "" + if e.stdout is not None: + stdout_str = e.stdout.decode("utf-8") + + stderr_str = "" + if e.stderr is not None: + stderr_str = e.stderr.decode("utf-8") - - process = subprocess.Popen( - [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, - "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, - ) - stdout, stderr = process.communicate() - stdout_str = "" - stderr_str = "" - if stdout is not None: - stdout_str = stdout.decode("utf-8") - # 清除制表符 - stdout_str = stdout_str.replace("\t", " ") - if stderr is not None: - stderr_str = stderr.decode("utf-8") - retcode = process.poll() - if retcode: - raise Exception(f"""请查看以上具体信息, -执行失败: -========================stdout======================== + raise Exception(f"""训练失败!!!具体报错信息请查看控制台... +=======================stdout======================= {stdout_str} -========================stderr======================== +=======================stderr======================= {stderr_str} -====================================================== -""") +==================================================== + """) + except Exception as e: + stop_server() + is_running = False + raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") + def MZ_KohyaSSTrain_call(args={}): args = args.copy() From 82277e828554585032c55e02f7413a0fdc3e288c Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 8 Jul 2024 01:21:36 +0800 Subject: [PATCH 08/10] 'update' --- mz_train_tools_core.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index 0612e56..ace0150 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -438,11 +438,13 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot print(f"stack: {traceback.format_exc()}") return is_running stop_server, port = Utils.Simple_Server(log_callback) + class CusprocessError(Exception): pass try: subprocess.run( [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], - check=True, + check=True, + capture_output=True, ) stop_server() is_running = False @@ -450,18 +452,18 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot stop_server() is_running = False stdout_str = "" - if e.stdout is not None: - stdout_str = e.stdout.decode("utf-8") - + if e.output is not None: + stdout_str = e.output.decode("utf-8") + stderr_str = "" if e.stderr is not None: stderr_str = e.stderr.decode("utf-8") - raise Exception(f"""训练失败!!!具体报错信息请查看控制台... + raise CusprocessError(f"""训练失败!!!具体报错信息请查看控制台... =======================stdout======================= {stdout_str} =======================stderr======================= -{stderr_str} +{stderr_str} ==================================================== """) except Exception as e: From c1a7f2d65f64869214f6a3c7819881fb673b1833 Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 8 Jul 2024 13:21:51 +0800 Subject: [PATCH 09/10] 'update' --- __init__.py | 49 +++++----- hook_kohya_ss_run.py | 44 ++++++--- mz_train_tools_core.py | 201 ++++++++++++++++++++++++++++------------- 3 files changed, 194 insertions(+), 100 deletions(-) diff --git a/__init__.py b/__init__.py index 1b73304..ce8790b 100644 --- a/__init__.py +++ b/__init__.py @@ -214,7 +214,8 @@ class MZ_KohyaSSAdvConfig: CATEGORY = CATEGORY_NAME + "/kohya_ss" def start(self, **kwargs): - return (kwargs,) + importlib.reload(mz_train_tools_core) + return mz_train_tools_core.MZ_KohyaSSAdvConfig_call(kwargs) NODE_CLASS_MAPPINGS["MZ_KohyaSSAdvConfig"] = MZ_KohyaSSAdvConfig @@ -376,25 +377,25 @@ class MZ_KohyaSSLoraTrain: def start(self, **kwargs): importlib.reload(mz_train_tools_core) - train_config = { - "workspace_config": kwargs["workspace_config"], - "train_config_template": kwargs["train_config_template"], - "ckpt_name": kwargs["ckpt_name"], - "max_train_steps": kwargs["max_train_steps"], - "max_train_epochs": kwargs["max_train_epochs"], - "save_every_n_epochs": kwargs["save_every_n_epochs"], - "learning_rate": kwargs["learning_rate"], - } + # train_config = { + # "workspace_config": kwargs["workspace_config"], + # "train_config_template": kwargs["train_config_template"], + # "ckpt_name": kwargs["ckpt_name"], + # "max_train_steps": kwargs["max_train_steps"], + # "max_train_epochs": kwargs["max_train_epochs"], + # "save_every_n_epochs": kwargs["save_every_n_epochs"], + # "learning_rate": kwargs["learning_rate"], + # } - train_config["train_config_template_dir"] = self.train_config_template_dir + # train_config["train_config_template_dir"] = self.train_config_template_dir - advanced_config = kwargs.get("advanced_config", None) + # advanced_config = kwargs.get("advanced_config", None) - if advanced_config is not None: - for k, v in advanced_config.items(): - train_config[k] = v + # if advanced_config is not None: + # for k, v in advanced_config.items(): + # train_config[k] = v - kwargs["train_config"] = train_config + # kwargs["train_config"] = train_config return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs) @@ -574,7 +575,7 @@ class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I: "text_encoder_path": (["auto"] + folders, {"default": "auto"}), "tokenizer_path": (["auto"] + folders, {"default": "auto"}), "t5_encoder_path": (["none", "auto"] + folders, {"default": "none"}), - "lora_path": (["none"] + comfyui_full_loras, {"default": "none"}), + "lora_path": (["none"] + comfyui_full_loras, {"default": "none"}), "seed": ("INT", {"default": 0}), "steps": ("INT", {"default": 20}), "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), @@ -605,7 +606,8 @@ class MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I: NODE_CLASS_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I -NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I" +NODE_DISPLAY_NAME_MAPPINGS[ + "MZ_KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I"] = f"{AUTHOR_NAME} - KohyaSS_KohakuBlueleaf_HYHiDSimpleT2I" class MZ_LoadImagesFromDirectoryPath: @@ -945,7 +947,7 @@ class MZ_TrainToolsDebug: def start(self, **kwargs): - from pprint import pprint, pp + from pprint import pprint, pformat object = kwargs["object"] indent = kwargs["indent"] depth = kwargs["depth"] @@ -954,11 +956,10 @@ class MZ_TrainToolsDebug: 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) - + debug = pformat(object, 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" +NODE_CLASS_MAPPINGS["MZ_TrainToolsDebug"] = MZ_TrainToolsDebug +NODE_DISPLAY_NAME_MAPPINGS["MZ_TrainToolsDebug"] = f"{AUTHOR_NAME} - TrainToolsDebug" diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index ecdfd35..83ab62a 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -1,5 +1,6 @@ import os import random +import time import torch import logging @@ -38,7 +39,7 @@ other_config = {} original_save_model = None -train_config_json = "{}" +train_config = {} sample_images_pipe_class = None @@ -48,7 +49,6 @@ def utils_sample_images(*args, **kwargs): def get_datasets(): - train_config = json.loads(train_config_json) import library.config_util user_config = library.config_util.load_user_config( train_config.get("dataset_config", None)) @@ -78,7 +78,7 @@ def sample_images(self, *args, **kwargs): if epoch is not None and cmd_args.save_every_n_epochs is not None and epoch % cmd_args.save_every_n_epochs == 0: datasets = get_datasets() - resolution = datasets.get("resolution", (512, 512)) + resolution = datasets.get("resolution", (512, 512)) if isinstance(resolution, int): resolution = (resolution, resolution) height, width = resolution @@ -101,7 +101,7 @@ def sample_images(self, *args, **kwargs): "height": height, "width": width, } - # + # prompt_dict_list.append(prompt_dict) else: for i, prompt_dict in enumerate(prompt_dict_list): @@ -149,7 +149,6 @@ def run_lora_sd1_5(): sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline trainer = train_network.NetworkTrainer() - train_config = json.loads(train_config_json) train_args = config2args(train_network.setup_parser(), train_config) LOG({ @@ -171,7 +170,6 @@ def run_lora_sdxl(): sample_images_pipe_class = library.sdxl_train_util.SdxlStableDiffusionLongPromptWeightingPipeline trainer = sdxl_train_network.SdxlNetworkTrainer() - train_config = json.loads(train_config_json) train_args = config2args(sdxl_train_network.setup_parser(), train_config) LOG({ @@ -215,7 +213,6 @@ def run_controlnet_sd1_5(): global sample_images_pipe_class sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline - train_config = json.loads(train_config_json) train_args = config2args(train_controlnet.setup_parser(), train_config) LOG({ @@ -362,7 +359,6 @@ def run_lora_hunyuan1_2(): library.hunyuan_utils.load_model = hunyuan_load_model trainer = hunyuan_train_network.HunYuanNetworkTrainer() - train_config = json.loads(train_config_json) train_args = config2args( hunyuan_train_network.setup_parser(), train_config) print(f"train_args = {train_args}") @@ -395,15 +391,11 @@ def LOG(log): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--sys_path", type=str, default="") - parser.add_argument("--train_config_json", type=str, default="") + parser.add_argument("--config", type=str, default="") parser.add_argument("--train_func", type=str, default="") parser.add_argument("--master_port", type=int, default=0) - parser.add_argument("--other_config_json", type=str, default="{}") args = parser.parse_args() - other_config_json = args.other_config_json - other_config = json.loads(other_config_json) - master_port = args.master_port print(f"master_port = {master_port}") @@ -412,10 +404,34 @@ if __name__ == "__main__": if sys_path != "": sys.path.append(sys_path) - train_config_json = args.train_config_json + config_file = args.config + if config_file == "": + raise Exception("train_config is empty") + + global_config = {} + with open(config_file, "r") as f: + _global_config = f.read() + global_config = json.loads(_global_config) + + train_config = global_config.get("train_config") + print(f"""=======================train_config======================= +{json.dumps(train_config, indent=4, ensure_ascii=False)} + """) + + other_config = global_config.get("other_config", {}) + print(f"""=======================other_config======================= +{json.dumps(other_config, indent=4, ensure_ascii=False)} + """) train_func = args.train_func if train_func == "": raise Exception("train_func is empty") + print(f"train_func = {train_func}") + + time.sleep(2) + LOG({ + "type": "Read configuration completed!", + }) + func_map[train_func]() diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index ace0150..f5e04b3 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -183,7 +183,7 @@ def MZ_ImageSelecter_call(args={}): def MZ_KohyaSSUseConfig_call(args={}): - # raise Exception(f"MZ_KohyaSSUseConfig_call: {args}") + args = args.copy() workspace_config = args.get("workspace_config", {}) workspace_name = workspace_config.get("workspace_name", None) @@ -249,6 +249,8 @@ def MZ_KohyaSSUseConfig_call(args={}): advanced_config = args.get("advanced_config", {}).copy() for k in advanced_config: + print(f"{k}= {advanced_config[k]}") + if type(advanced_config[k]) == str and advanced_config[k] == "": if k in config["train_config"]: del config["train_config"][k] @@ -275,6 +277,12 @@ def MZ_KohyaSSUseConfig_call(args={}): ) +def MZ_KohyaSSAdvConfig_call(args={}): + return ( + args, + ) + + def config2args(train_parser: argparse.ArgumentParser, config): try: config_args_list = [] @@ -354,10 +362,9 @@ def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_ from PIL import Image -def get_sample_images(train_config): - output_name = train_config.get("output_name") +def get_sample_images(workspace_dir, output_name): sample_images_dir = os.path.join( - os.path.dirname(train_config.get("dataset_config")), "sample_images" + workspace_dir, "sample_images" ) pil_images = [] pre_render_texts_x = [] @@ -385,15 +392,13 @@ def get_sample_images(train_config): return result -def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, other_config={}): +def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, trainer_func): - other_config_str = json.dumps(other_config) + train_config_file = os.path.join(workspace_dir, "config.json") exec_pyfile = os.path.join(os.path.dirname( __file__), "hook_kohya_ss_run.py",) - train_config_str = json.dumps(train_config) - max_train_steps = train_config.get("max_train_steps") - max_train_epochs = train_config.get("max_train_epochs") + is_running = True taesd_type = "sd1_5" @@ -404,7 +409,7 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot if trainer_func.find("hunyuan1_1") != -1: taesd_type = "sdxl" - pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type) + pb = Utils.progress_bar(0, taesd_type) import traceback @@ -423,7 +428,7 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot resp = log if resp.get("type") == "sample_images": global_step = resp.get("global_step") - xy_img = get_sample_images(train_config) + xy_img = get_sample_images(workspace_dir, output_name) max_side = max(xy_img.width, xy_img.height) # print(f"global_step: {global_step}, max_train_steps: {max_train_steps}") @@ -438,48 +443,38 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot print(f"stack: {traceback.format_exc()}") return is_running stop_server, port = Utils.Simple_Server(log_callback) - class CusprocessError(Exception): pass try: + cmd_list = [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, + "--config", train_config_file, "--train_func", trainer_func, "--master_port", str(port)] + startup_script_path_sh = os.path.join(workspace_dir, "start_train.sh") + startup_script_path_bat = os.path.join( + workspace_dir, "start_train.bat") + with open(startup_script_path_sh, "w", encoding="utf-8") as f: + f.write(" ".join(cmd_list)) + with open(startup_script_path_bat, "w", encoding="utf-8") as f: + f.write(" ".join(cmd_list)) + subprocess.run( - [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, - "--train_config_json", train_config_str, "--train_func", trainer_func, "--master_port", str(port), "--other_config_json", other_config_str], + cmd_list, check=True, - capture_output=True, ) stop_server() is_running = False - except subprocess.CalledProcessError as e: - stop_server() - is_running = False - stdout_str = "" - if e.output is not None: - stdout_str = e.output.decode("utf-8") - - stderr_str = "" - if e.stderr is not None: - stderr_str = e.stderr.decode("utf-8") - - raise CusprocessError(f"""训练失败!!!具体报错信息请查看控制台... -=======================stdout======================= -{stdout_str} -=======================stderr======================= -{stderr_str} -==================================================== - """) except Exception as e: stop_server() is_running = False raise Exception(f"训练失败!!! 具体报错信息请查看控制台...") -def MZ_KohyaSSTrain_call(args={}): +def generate_kohya_ss_config(args): args = args.copy() - workspace_config = args.get("workspace_config", {}) - - train_config = args.get("train_config", {}) - MZ_KohyaSSUseConfig_call(train_config) + workspace_config = args.get("workspace_config", {}).copy() + advanced_config = args.get("advanced_config", {}).copy() + train_config = args.get("train_config", {}).copy() workspace_name = workspace_config.get("workspace_name", None) + if workspace_name is None or workspace_name == "": + raise Exception("工作区名称不能为空(workspace_name is required)") workspace_dir = os.path.join( folder_paths.output_directory, "mz_train_workspaces", workspace_name) @@ -488,16 +483,88 @@ def MZ_KohyaSSTrain_call(args={}): workspace_config_file = os.path.join(workspace_dir, "config.json") - if not os.path.exists(workspace_config_file): - raise Exception(f"配置文件不存在: {workspace_config_file}") + train_config_template = args.get("train_config_template", None) + train_config_template_dir = os.path.join( + os.path.dirname(__file__), "configs", "kohya_ss_lora") + train_config_template_file = os.path.join( + train_config_template_dir, train_config_template + ".json") + # raise Exception(f"args: {json.dumps(args, indent=4, ensure_ascii=False)}") config = None - with open(workspace_config_file, "r", encoding="utf-8") as f: + with open(train_config_template_file, "r", encoding="utf-8") as f: config = json.load(f) + config["metadata"]["train_type"] = train_config_template + ckpt_name = args.get("ckpt_name", "") + if ckpt_name != "" and ckpt_name is not None: + ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + config["train_config"]["pretrained_model_name_or_path"] = ckpt_path + + # output_dir + output_dir = os.path.join(workspace_dir, "output") + config["train_config"]["output_dir"] = output_dir + + datetime = time.strftime("%Y%m%d_%H%M%S", time.localtime()) + # output_name + config["train_config"]["output_name"] = f"{workspace_name}_{train_config_template}_{datetime}" + + dataset_config_path = os.path.join( + workspace_dir, "dataset.toml") + config["train_config"]["dataset_config"] = dataset_config_path + + config["train_config"]["max_train_steps"] = str( + args.get("max_train_steps")) + + config["train_config"]["max_train_epochs"] = str( + args.get("max_train_epochs")) + if config["train_config"]["max_train_epochs"] == "0": + config["train_config"]["max_train_epochs"] = False + + config["train_config"]["save_every_n_epochs"] = str( + args.get("save_every_n_epochs")) + + config["train_config"]["learning_rate"] = str( + args.get("learning_rate")) + + for k in advanced_config: + print(f"{k}= {advanced_config[k]}") + + if type(advanced_config[k]) == str and advanced_config[k] == "": + if k in config["train_config"]: + del config["train_config"][k] + continue + elif advanced_config[k] == "enable": + advanced_config[k] = True + elif advanced_config[k] == "disable": + advanced_config[k] = False + else: + advanced_config[k] = str(advanced_config[k]) + config["train_config"][k] = advanced_config[k] + + # raise Exception(f"args: {json.dumps(config, indent=4)}") if config is None: raise Exception(f"读取配置文件失败: {workspace_config_file}") + # raise Exception(f"MZ_KohyaSSUseConfig_call: {args}") + return config + + +def MZ_KohyaSSTrain_call(args={}): + args = args.copy() + workspace_config = args.get("workspace_config").copy() + base_lora = args.get("base_lora", "empty") + sample_generate = args.get("sample_generate", "enable") + sample_prompt = args.get("sample_prompt", "") + + workspace_name = workspace_config.get("workspace_name") + workspace_dir = os.path.join( + folder_paths.output_directory, "mz_train_workspaces", workspace_name) + + if not os.path.exists(workspace_dir): + raise Exception(f"工作区不存在: {workspace_dir}") + + config = generate_kohya_ss_config(args) + branch_local_name = workspace_config.get( "branch_local_name", "kohya_ss_lora") kohya_ss_tool_dir = os.path.join( @@ -507,7 +574,6 @@ def MZ_KohyaSSTrain_call(args={}): sys.path.append(kohya_ss_tool_dir) check_install() - base_lora = args.get("base_lora", "empty") if base_lora == "empty": pass elif base_lora == "latest": @@ -529,17 +595,16 @@ def MZ_KohyaSSTrain_call(args={}): else: pass - train_config = config.get("train_config") if base_lora != "empty" and os.path.exists(base_lora): - train_config["network_weights"] = base_lora - train_config["dim_from_weights"] = True + config["train_config"]["network_weights"] = base_lora + config["train_config"]["dim_from_weights"] = True - if "network_dim" in train_config: - del train_config["network_dim"] - if "network_alpha" in train_config: - del train_config["network_alpha"] - if "network_dropout" in train_config: - del train_config["network_dropout"] + if "network_dim" in config["train_config"]: + del config["train_config"]["network_dim"] + if "network_alpha" in config["train_config"]: + del config["train_config"]["network_alpha"] + if "network_dropout" in config["train_config"]: + del config["train_config"]["network_dropout"] base_controlnet = args.get("base_controlnet", "empty") if base_controlnet == "empty": @@ -565,25 +630,32 @@ def MZ_KohyaSSTrain_call(args={}): pass if base_controlnet != "empty" and os.path.exists(base_controlnet): - train_config["controlnet_model_name_or_path"] = base_controlnet + config["train_config"]["controlnet_model_name_or_path"] = base_controlnet train_type = config.get("metadata").get("train_type") - sample_generate = args.get("sample_generate", "enable") - sample_prompt = args.get("sample_prompt", "") if sample_generate == "enable": - other_config = { + config["other_config"] = { "sample_prompt": sample_prompt, } else: - other_config = {} + config["other_config"] = {} + + workspace_config_file = os.path.join(workspace_dir, "config.json") + with open(workspace_config_file, "w", encoding="utf-8") as f: + json.dump(config, f, indent=4, ensure_ascii=False) + + # raise Exception( + # f"config: {json.dumps(config, indent=4, ensure_ascii=False)}") + + output_name = config["train_config"].get("output_name") if train_type == "lora_sd1_5": run_hook_kohya_ss_run_file( - kohya_ss_tool_dir, train_config, "run_lora_sd1_5", other_config) + workspace_dir, output_name, kohya_ss_tool_dir, "run_lora_sd1_5") elif train_type == "lora_sdxl": run_hook_kohya_ss_run_file( - kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config) + workspace_dir, output_name, kohya_ss_tool_dir, "run_lora_sdxl") elif train_type == "controlnet_sd1_5": conditioning_images_dir = os.path.join( @@ -592,11 +664,13 @@ def MZ_KohyaSSTrain_call(args={}): if os.path.exists(conditioning_images_dir): conditioning_images_onec = Utils.listdir( conditioning_images_dir)[0] - other_config["controlnet_image"] = os.path.join( + config["other_config"]["controlnet_image"] = os.path.join( conditioning_images_dir, conditioning_images_onec) + with open(workspace_config_file, "w", encoding="utf-8") as f: + json.dump(config, f, indent=4, ensure_ascii=False) run_hook_kohya_ss_run_file( - kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config) + workspace_dir, output_name, kohya_ss_tool_dir, "run_controlnet_sd1_5") elif train_type == "lora_hunyuan1_2" or train_type == "lora_hunyuan1_1": hunyuan_models_config = args.get( "hunyuan_models_config", {}) @@ -604,11 +678,14 @@ def MZ_KohyaSSTrain_call(args={}): hunyuan_models_config["version"] = config.get( "metadata").get("version") - other_config["hunyuan_models_config"] = check_model_auto_download( + config["other_config"]["hunyuan_models_config"] = check_model_auto_download( hunyuan_models_config) + with open(workspace_config_file, "w", encoding="utf-8") as f: + json.dump(config, f, indent=4, ensure_ascii=False) + run_hook_kohya_ss_run_file( - kohya_ss_tool_dir, train_config, "run_lora_hunyuan1_2", other_config) + workspace_dir, output_name, kohya_ss_tool_dir, "run_lora_hunyuan1_2") else: raise Exception( f"暂时不支持的训练类型: {train_type}") From 2c348743d0b9305bef46375164b40ce3c9229902 Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 8 Jul 2024 20:55:14 +0800 Subject: [PATCH 10/10] 'update' --- hook_kohya_ss_run.py | 16 +++++--- mz_train_tools_core.py | 19 +++++++--- mz_train_tools_utils.py | 82 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 107 insertions(+), 10 deletions(-) diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index 83ab62a..e52cd94 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -1,4 +1,6 @@ import os + +os.system("title hook_kohya_ss_run") import random import time @@ -381,11 +383,15 @@ import requests def LOG(log): - # 发送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}") + try: + # 发送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}") + print(f"LOG failed: {resp.text}") + except Exception as e: + print(f"LOG failed: {e}") if __name__ == "__main__": diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index f5e04b3..4009ecd 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -392,7 +392,7 @@ def get_sample_images(workspace_dir, output_name): return result -def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, trainer_func): +def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, trainer_func, use_screen=False): train_config_file = os.path.join(workspace_dir, "config.json") @@ -422,6 +422,8 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr comfy.model_management.throw_exception_if_processing_interrupted() except Exception as e: stop_server() + if process_instance is not None: + process_instance.stop() # 强制关闭进程 return is_running try: @@ -442,6 +444,7 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr print(f"LOG: {log} e: {e} ") print(f"stack: {traceback.format_exc()}") return is_running + stop_server, port = Utils.Simple_Server(log_callback) try: cmd_list = [sys.executable, exec_pyfile, "--sys_path", kohya_ss_tool_dir, @@ -454,10 +457,16 @@ def run_hook_kohya_ss_run_file(workspace_dir, output_name, kohya_ss_tool_dir, tr with open(startup_script_path_bat, "w", encoding="utf-8") as f: f.write(" ".join(cmd_list)) - subprocess.run( - cmd_list, - check=True, - ) + from .mz_train_tools_utils import HSubprocess + + screen_name = None + if use_screen: + screen_name = "mz_train_tools_core" + + process_instance = HSubprocess( + cmd_list, screen_name=screen_name) + process_instance.wait() + stop_server() is_running = False except Exception as e: diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index b52a7b6..a08b549 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -1135,3 +1135,85 @@ class CustomizeMT5Embedder(nn.Module): def get_hidden_states(self, input_ids, layer_index=-1): return self.get_text_embeddings(input_ids, layer_index=layer_index) + + +class HSubprocess: + process_instance = None + process_instance_pid = None + screen_name = None + _mswindows = False + + def __init__(self, args, screen_name=None): + self._mswindows = (sys.platform == "win32") + if self._mswindows: + self.args = ["start", "/w"] + self.args.extend(args) + return + + self.screen_name = screen_name + if screen_name is not None: + try: + subprocess.check_call(["screen", "-v"]) + except Exception as e: + raise Exception("Please install screen first.") + + screen_cmd = ["screen", "-R", screen_name, "-m", ] + screen_cmd.extend(args) + + self.args = screen_cmd + else: + self.args = args + + def stop(self): + if self._mswindows: + if self.process_instance_pid is not None: + os.system( + f'taskkill /F /FI "WINDOWTITLE eq hook_kohya_ss_run" /T') + if self.process_instance is not None: + if self.screen_name is not None: + subprocess.run( + ["screen", "-d", self.screen_name]) + + self.process_instance.kill() + self.process_instance = None + + try: + try: + import psutil + except ImportError: + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "psutil"]) + import psutil + psutil.Process(self.process_instance_pid).terminate() + except Exception as e: + print(e) + + self.process_instance_pid = None + + def wait(self): + with subprocess.Popen( + self.args, + stdin=subprocess.PIPE, + shell=True, + ) as process: + self.process_instance = process + self.process_instance_pid = process.pid + print(f"Subprocess PID: {self.process_instance_pid}") + try: + stdout, stderr = process.communicate() + except subprocess.TimeoutExpired as exc: + process.kill() + if self._mswindows: + exc.stdout, exc.stderr = process.communicate() + else: + process.wait() + raise + except Exception as e: + process.wait() + raise + retcode = process.poll() + if retcode != 0: + raise subprocess.CalledProcessError(retcode, process.args) + + self.process_instance = None + self.process_instance_pid = None