喵哩个咪
+27
-25
@@ -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",),
|
||||
@@ -213,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
|
||||
@@ -375,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)
|
||||
|
||||
|
||||
@@ -573,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}),
|
||||
@@ -604,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:
|
||||
@@ -632,7 +635,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 = []
|
||||
@@ -944,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"]
|
||||
@@ -953,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"
|
||||
|
||||
+41
-19
@@ -1,5 +1,8 @@
|
||||
import os
|
||||
|
||||
os.system("title hook_kohya_ss_run")
|
||||
import random
|
||||
import time
|
||||
|
||||
import torch
|
||||
import logging
|
||||
@@ -43,7 +46,7 @@ other_config = {}
|
||||
original_save_model = None
|
||||
|
||||
|
||||
train_config_json = "{}"
|
||||
train_config = {}
|
||||
|
||||
sample_images_pipe_class = None
|
||||
|
||||
@@ -53,7 +56,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))
|
||||
@@ -83,7 +85,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
|
||||
@@ -106,7 +108,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):
|
||||
@@ -154,7 +156,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({
|
||||
@@ -176,7 +177,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({
|
||||
@@ -220,7 +220,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({
|
||||
@@ -367,7 +366,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}")
|
||||
@@ -390,25 +388,25 @@ 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__":
|
||||
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}")
|
||||
@@ -417,10 +415,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]()
|
||||
|
||||
+156
-47
@@ -129,17 +129,20 @@ 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)
|
||||
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)
|
||||
|
||||
@@ -180,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)
|
||||
|
||||
@@ -246,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]
|
||||
@@ -272,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 = []
|
||||
@@ -351,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 = []
|
||||
@@ -382,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, use_screen=False):
|
||||
|
||||
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"
|
||||
@@ -401,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
|
||||
|
||||
@@ -414,13 +422,15 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
|
||||
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:
|
||||
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}")
|
||||
@@ -434,13 +444,29 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
|
||||
print(f"LOG: {log} e: {e} ")
|
||||
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,
|
||||
)
|
||||
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))
|
||||
|
||||
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:
|
||||
@@ -449,14 +475,15 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
|
||||
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)
|
||||
|
||||
@@ -465,16 +492,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(
|
||||
@@ -484,7 +583,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":
|
||||
@@ -506,17 +604,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":
|
||||
@@ -542,25 +639,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(
|
||||
@@ -569,11 +673,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", {})
|
||||
@@ -581,11 +687,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}")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -1134,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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -1,245 +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="")
|
||||
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
|
||||
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)
|
||||
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()
|
||||
Reference in New Issue
Block a user