1138 lines
42 KiB
Python
1138 lines
42 KiB
Python
|
|
|
|
import argparse
|
|
import csv
|
|
import hashlib
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import json
|
|
import subprocess
|
|
import time
|
|
|
|
import safetensors.torch
|
|
import torch
|
|
|
|
|
|
from .mz_train_tools_utils import Utils
|
|
import folder_paths
|
|
import nodes
|
|
|
|
|
|
git_accelerate_urls = {
|
|
"githubfast": "githubfast.com",
|
|
"521github": "521github.com",
|
|
"kkgithub": "kkgithub.com",
|
|
}
|
|
|
|
# 初始化工具仓库和工作区
|
|
|
|
|
|
def MZ_HYDiTInitClone_call(args={}):
|
|
mz_dir = Utils.get_minus_zone_models_path()
|
|
git_url = "https://github.com/Tencent/HunyuanDiT"
|
|
source = args.get("source", "github")
|
|
hunyuan_lora_dir = os.path.join(mz_dir, "train_tools", "HunyuanDiT")
|
|
if git_accelerate_urls.get(source, None) is not None:
|
|
git_url = f"https://{git_accelerate_urls[source]}/Tencent/HunyuanDiT"
|
|
try:
|
|
if not os.path.exists(hunyuan_lora_dir) or not os.path.exists(os.path.join(hunyuan_lora_dir, ".git")):
|
|
subprocess.run(
|
|
["git", "clone", "--depth", "1", git_url, hunyuan_lora_dir], check=True)
|
|
|
|
# 切换远程分支 git remote set-branches origin 'main'
|
|
branch = args.get("branch", "main")
|
|
|
|
# 查看本地分支是否一致
|
|
short_result = subprocess.run(
|
|
["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=hunyuan_lora_dir, stdout=subprocess.PIPE, check=True)
|
|
result = subprocess.run(
|
|
["git", "rev-parse", "HEAD"], cwd=hunyuan_lora_dir, stdout=subprocess.PIPE, check=True)
|
|
|
|
short_current_branch = short_result.stdout.decode().strip()
|
|
long_current_branch = result.stdout.decode().strip()
|
|
print(
|
|
f"当前分支(current branch): {long_current_branch}({short_current_branch})")
|
|
print(f"目标分支(target branch): {branch}")
|
|
time.sleep(1)
|
|
if branch != long_current_branch and branch != short_current_branch:
|
|
subprocess.run(
|
|
["git", "remote", "set-branches", "origin", branch], cwd=hunyuan_lora_dir, check=True)
|
|
subprocess.run(
|
|
["git", "fetch", "--depth", "1", "origin", branch], cwd=hunyuan_lora_dir, check=True)
|
|
|
|
# 恢复所有文件
|
|
subprocess.run(
|
|
["git", "checkout", "."], cwd=hunyuan_lora_dir, check=True)
|
|
|
|
subprocess.run(
|
|
["git", "checkout", branch], cwd=hunyuan_lora_dir, check=True)
|
|
|
|
content = None
|
|
with open(os.path.join(hunyuan_lora_dir, "hydit/diffusion/pipeline.py"), "r", encoding="utf-8") as f:
|
|
pre_replace = "device = self._execution_device"
|
|
content = f.read()
|
|
content = content.replace(
|
|
pre_replace, "device = torch.device('cuda')")
|
|
with open(os.path.join(hunyuan_lora_dir, "hydit/diffusion/pipeline.py"), "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
|
|
except Exception as e:
|
|
raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...")
|
|
|
|
|
|
def MZ_HYDiTInitWorkspace_call(args={}):
|
|
MZ_HYDiTInitClone_call(args)
|
|
|
|
workspace_name = args.get("train_name", None)
|
|
workspace_name = workspace_name.strip()
|
|
|
|
if workspace_name is None or workspace_name == "":
|
|
raise Exception("训练名称不能为空(train_name is required)")
|
|
|
|
args["workspace_name"] = workspace_name
|
|
workspaces_dir = os.path.join(
|
|
folder_paths.output_directory, "mz_train_workspaces")
|
|
|
|
os.makedirs(workspaces_dir, exist_ok=True)
|
|
|
|
workspace_dir = os.path.join(workspaces_dir, workspace_name)
|
|
if not os.path.exists(workspace_dir):
|
|
os.makedirs(workspace_dir)
|
|
|
|
args["workspace_dir"] = workspace_dir
|
|
|
|
return (
|
|
args,
|
|
)
|
|
|
|
|
|
def MZ_HYDiTDatasetConfig_call(args={}):
|
|
workspace_config = args.get("workspace_config", {})
|
|
workspace_name = workspace_config.get("workspace_name", None)
|
|
workspace_dir = os.path.join(
|
|
folder_paths.output_directory, "mz_train_workspaces", workspace_name)
|
|
train_images_dir = os.path.join(workspace_dir, "train_images")
|
|
os.makedirs(train_images_dir, exist_ok=True)
|
|
|
|
images = args.get("images", None)
|
|
|
|
if images is None:
|
|
print(f"训练数据位于:{train_images_dir},未检测到传入图片默认已有训练数据,将直接跳过数据准备步骤...")
|
|
same_caption_generate = args.get("same_caption_generate") == "enable"
|
|
if same_caption_generate:
|
|
saved_images_path = Utils.listdir(train_images_dir)
|
|
same_caption = args.get("same_caption").strip()
|
|
if same_caption != "":
|
|
# 循环已经保存的图片
|
|
for i, filename in enumerate(saved_images_path):
|
|
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:
|
|
f.write(same_caption)
|
|
|
|
|
|
return (
|
|
train_images_dir,
|
|
)
|
|
|
|
pil_images = Utils.tensors2pil_list(images)
|
|
|
|
resolution = args.get("resolution", 1024)
|
|
|
|
if workspace_name is None or workspace_name == "":
|
|
raise Exception("训练名称不能为空(workspace_name is required)")
|
|
|
|
force_clear = args.get("force_clear") == "enable"
|
|
force_clear_only_images = args.get("force_clear_only_images") == "enable"
|
|
if force_clear:
|
|
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") 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)
|
|
|
|
saved_images_path = []
|
|
|
|
for i, pil_image in enumerate(pil_images):
|
|
pil_image = Utils.resize_max(pil_image, resolution, resolution)
|
|
width, height = pil_image.size
|
|
filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png"
|
|
pil_image.save(os.path.join(train_images_dir, filename))
|
|
saved_images_path.append(filename)
|
|
|
|
same_caption_generate = args.get("same_caption_generate") == "enable"
|
|
if same_caption_generate:
|
|
same_caption = args.get("same_caption").strip()
|
|
if same_caption != "":
|
|
# 循环已经保存的图片
|
|
for i, filename in enumerate(saved_images_path):
|
|
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:
|
|
f.write(same_caption)
|
|
|
|
return (
|
|
train_images_dir,
|
|
)
|
|
|
|
|
|
HYDiT_MODEL = {
|
|
"HunyuanDiT-v1.2/t2i/model/pytorch_model_module.pt": {
|
|
"url": "https://www.modelscope.cn/models/wailovet/hy1.2/resolve/master/t2i/model/pytorch_model_module.pt",
|
|
},
|
|
"HunyuanDiT/t2i/model/pytorch_model_ema.pt": {
|
|
"url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_ema.pt",
|
|
},
|
|
"HunyuanDiT/t2i/model/pytorch_model_module.pt": {
|
|
"url": "https://www.modelscope.cn/models/wailovet/hy1.1/resolve/master/t2i/model/pytorch_model_module.pt",
|
|
},
|
|
"HunyuanDiT/t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fsdxl-vae-fp16-fix%2Fdiffusion_pytorch_model.safetensors",
|
|
},
|
|
"HunyuanDiT/t2i/sdxl-vae-fp16-fix/config.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fsdxl-vae-fp16-fix%2Fconfig.json",
|
|
},
|
|
"HunyuanDiT/t2i/clip_text_encoder/pytorch_model.bin": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fclip_text_encoder%2Fpytorch_model.bin",
|
|
},
|
|
"HunyuanDiT/t2i/clip_text_encoder/config.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fclip_text_encoder%2Fconfig.json",
|
|
},
|
|
"HunyuanDiT/t2i/tokenizer/special_tokens_map.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Fspecial_tokens_map.json",
|
|
},
|
|
"HunyuanDiT/t2i/tokenizer/tokenizer_config.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Ftokenizer_config.json",
|
|
},
|
|
"HunyuanDiT/t2i/tokenizer/vocab.txt": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Fvocab.txt",
|
|
},
|
|
"HunyuanDiT/t2i/tokenizer/vocab_org.txt": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Ftokenizer%2Fvocab_org.txt",
|
|
},
|
|
"HunyuanDiT/t2i/mt5/config.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fconfig.json",
|
|
},
|
|
"HunyuanDiT/t2i/mt5/generation_config.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fgeneration_config.json",
|
|
},
|
|
"HunyuanDiT/t2i/mt5/special_tokens_map.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fspecial_tokens_map.json",
|
|
},
|
|
"HunyuanDiT/t2i/mt5/spiece.model": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fspiece.model",
|
|
},
|
|
"HunyuanDiT/t2i/mt5/tokenizer_config.json": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Ftokenizer_config.json",
|
|
},
|
|
"HunyuanDiT/t2i/mt5/pytorch_model.bin": {
|
|
"url": "https://www.modelscope.cn/api/v1/models/modelscope/HunyuanDiT/repo?Revision=master&FilePath=t2i%2Fmt5%2Fpytorch_model.bin",
|
|
},
|
|
}
|
|
|
|
|
|
def check_model_auto_download(args):
|
|
unet_path = args.get("unet_path", "auto")
|
|
hunyuan_base_path = os.path.join(
|
|
Utils.get_comfyui_models_path(), "hunyuan")
|
|
if unet_path == "auto":
|
|
version = args.get("version", None)
|
|
if version is None:
|
|
ema_to_module = args.get("ema_to_module", "enable") == "enable"
|
|
if ema_to_module:
|
|
download_file = "HunyuanDiT/t2i/model/pytorch_model_ema.pt"
|
|
else:
|
|
download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt"
|
|
else:
|
|
if version == "1.2":
|
|
download_file = "HunyuanDiT-v1.2/t2i/model/pytorch_model_module.pt"
|
|
elif version == "1.1":
|
|
download_file = "HunyuanDiT/t2i/model/pytorch_model_module.pt"
|
|
else:
|
|
raise Exception(f"未知的版本号: {version}")
|
|
|
|
download_fullpath = os.path.join(
|
|
hunyuan_base_path, download_file)
|
|
if not os.path.exists(download_fullpath):
|
|
os.makedirs(os.path.dirname(download_fullpath), exist_ok=True)
|
|
|
|
if os.path.exists(download_fullpath):
|
|
args["unet_path"] = download_fullpath
|
|
else:
|
|
success_path = Utils.download_file(
|
|
HYDiT_MODEL[download_file]["url"],
|
|
download_fullpath,
|
|
)
|
|
if os.path.exists(success_path):
|
|
args["unet_path"] = download_fullpath
|
|
|
|
vae_ema_path = args.get("vae_ema_path", "auto")
|
|
if vae_ema_path == "auto":
|
|
download_files = []
|
|
prefix = "HunyuanDiT/t2i/sdxl-vae-fp16-fix/"
|
|
for key in HYDiT_MODEL:
|
|
if key.startswith(prefix):
|
|
download_files.append(key)
|
|
|
|
for download_file in download_files:
|
|
download_fullpath = os.path.join(
|
|
hunyuan_base_path, download_file)
|
|
if not os.path.exists(download_fullpath):
|
|
os.makedirs(os.path.dirname(download_fullpath), exist_ok=True)
|
|
|
|
if not os.path.exists(download_fullpath):
|
|
Utils.download_file(
|
|
HYDiT_MODEL[download_file]["url"],
|
|
download_fullpath,
|
|
)
|
|
|
|
args["vae_ema_path"] = os.path.join(
|
|
hunyuan_base_path, prefix)
|
|
|
|
text_encoder_path = args.get("text_encoder_path", "auto")
|
|
if text_encoder_path == "auto":
|
|
download_files = []
|
|
prefix = "HunyuanDiT/t2i/clip_text_encoder/"
|
|
for key in HYDiT_MODEL:
|
|
if key.startswith(prefix):
|
|
download_files.append(key)
|
|
|
|
for download_file in download_files:
|
|
download_fullpath = os.path.join(
|
|
hunyuan_base_path, download_file)
|
|
if not os.path.exists(download_fullpath):
|
|
os.makedirs(os.path.dirname(download_fullpath), exist_ok=True)
|
|
|
|
if not os.path.exists(download_fullpath):
|
|
Utils.download_file(
|
|
HYDiT_MODEL[download_file]["url"],
|
|
download_fullpath,
|
|
)
|
|
|
|
args["text_encoder_path"] = os.path.join(
|
|
hunyuan_base_path, prefix)
|
|
|
|
tokenizer_path = args.get("tokenizer_path", "auto")
|
|
if tokenizer_path == "auto":
|
|
download_files = []
|
|
prefix = "HunyuanDiT/t2i/tokenizer/"
|
|
for key in HYDiT_MODEL:
|
|
if key.startswith(prefix):
|
|
download_files.append(key)
|
|
|
|
for download_file in download_files:
|
|
download_fullpath = os.path.join(
|
|
hunyuan_base_path, download_file)
|
|
if not os.path.exists(download_fullpath):
|
|
os.makedirs(os.path.dirname(download_fullpath), exist_ok=True)
|
|
|
|
if not os.path.exists(download_fullpath):
|
|
Utils.download_file(
|
|
HYDiT_MODEL[download_file]["url"],
|
|
download_fullpath,
|
|
)
|
|
|
|
args["tokenizer_path"] = os.path.join(
|
|
hunyuan_base_path, prefix)
|
|
|
|
t5_encoder_path = args.get("t5_encoder_path", "auto")
|
|
if t5_encoder_path == "auto":
|
|
download_files = []
|
|
prefix = "HunyuanDiT/t2i/mt5/"
|
|
for key in HYDiT_MODEL:
|
|
if key.startswith(prefix):
|
|
download_files.append(key)
|
|
|
|
for download_file in download_files:
|
|
download_fullpath = os.path.join(
|
|
hunyuan_base_path, download_file)
|
|
if not os.path.exists(download_fullpath):
|
|
os.makedirs(os.path.dirname(download_fullpath), exist_ok=True)
|
|
|
|
if not os.path.exists(download_fullpath):
|
|
Utils.download_file(
|
|
HYDiT_MODEL[download_file]["url"],
|
|
download_fullpath,
|
|
)
|
|
|
|
args["t5_encoder_path"] = os.path.join(
|
|
hunyuan_base_path, prefix)
|
|
|
|
return args
|
|
|
|
|
|
def check_required():
|
|
packages = [
|
|
"pandas", "pyarrow", "diffusers", "transformers",
|
|
"timm", "peft", "accelerate", "loguru", "einops", "sentencepiece",
|
|
"polygraphy", "protobuf"
|
|
]
|
|
for package in packages:
|
|
try:
|
|
__import__(package)
|
|
except ImportError:
|
|
subprocess.run([sys.executable, "-m", "pip",
|
|
"install", package], check=True)
|
|
|
|
try:
|
|
import deepspeed
|
|
except ImportError:
|
|
raise Exception("请手动安装合适的deepspeed版本")
|
|
|
|
|
|
def MZ_HYDiTTrain_call(args={}):
|
|
|
|
check_required()
|
|
|
|
advanced_config = args.get("advanced_config", {})
|
|
fix_float_keys = ["lr", "warmup_min_lr",
|
|
"warmup_max_lr", "weight_decay", "warmup_num_steps",]
|
|
for key in fix_float_keys:
|
|
if key in advanced_config:
|
|
if isinstance(advanced_config[key], str):
|
|
advanced_config[key] = float(advanced_config[key])
|
|
|
|
args = check_model_auto_download(args)
|
|
# raise Exception(args)
|
|
resolution = args.get("resolution")
|
|
|
|
workspace_config = args.get("workspace_config", {})
|
|
workspace_name = workspace_config.get("workspace_name", None)
|
|
workspace_dir = workspace_config.get("workspace_dir", None)
|
|
|
|
workspace_images_dir = args.get("workspace_images_dir", None)
|
|
if workspace_images_dir is None:
|
|
workspace_images_dir = os.path.join(workspace_dir, "train_images")
|
|
|
|
full_filenames = Utils.listdir(workspace_images_dir)
|
|
|
|
# 创建image_text.csv文件
|
|
csv_filename = os.path.join(workspace_dir, "image_text.csv")
|
|
|
|
if os.path.exists(csv_filename):
|
|
os.remove(csv_filename)
|
|
|
|
max_width, max_height, min_width, min_height = 0, 0, 99999, 99999
|
|
|
|
from PIL import Image
|
|
|
|
with open(csv_filename, "w", encoding="utf-8", newline="") as file:
|
|
writer = csv.writer(file)
|
|
|
|
writer.writerow(["image_path", "text_zh"])
|
|
for filename in full_filenames:
|
|
print(f"处理文件: {filename}")
|
|
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)
|
|
width, height = pil_image.size
|
|
if width > max_width:
|
|
max_width = width
|
|
if height > max_height:
|
|
max_height = height
|
|
if width < min_width:
|
|
min_width = width
|
|
if height < min_height:
|
|
min_height = height
|
|
|
|
text_zh = ""
|
|
|
|
# 去除后缀
|
|
_caption_filename = filename.split(".")[0]
|
|
caption_filename = ""
|
|
if os.path.exists(os.path.join(workspace_images_dir, _caption_filename + ".txt")):
|
|
caption_filename = _caption_filename + ".txt"
|
|
elif os.path.exists(os.path.join(workspace_images_dir, _caption_filename + ".caption")):
|
|
caption_filename = _caption_filename + ".caption"
|
|
else:
|
|
caption_filename = ""
|
|
if caption_filename != "":
|
|
with open(os.path.join(workspace_images_dir, caption_filename), "r") as f:
|
|
caption = f.read()
|
|
text_zh = caption
|
|
|
|
writer.writerow([image_path, text_zh])
|
|
|
|
if resolution > max_width and resolution > max_height:
|
|
print("resolution 大于所有图片的宽高,强制使用最大宽高作为resolution")
|
|
resolution = max(max_width, max_height) // 8 * 8
|
|
|
|
print(f"csv文件已生成: {csv_filename}")
|
|
|
|
HYDiT_tool_dir = os.path.join(
|
|
Utils.get_minus_zone_models_path(), "train_tools", "HunyuanDiT")
|
|
|
|
# python ./hydit/data_loader/csv2arrow.py ./dataset/porcelain/csvfile/image_text.csv ./dataset/porcelain/arrows
|
|
arrows_dir = os.path.join(workspace_dir, "arrows")
|
|
if os.path.exists(arrows_dir):
|
|
shutil.rmtree(arrows_dir)
|
|
|
|
os.makedirs(arrows_dir, exist_ok=True)
|
|
csv2arrow_exec = os.path.join(
|
|
HYDiT_tool_dir, "hydit", "data_loader", "csv2arrow.py")
|
|
try:
|
|
subprocess.run([
|
|
sys.executable, csv2arrow_exec, csv_filename, arrows_dir
|
|
], check=True)
|
|
except Exception as e:
|
|
raise Exception(f"生成arrow文件时出现异常,详细信息请查看控制台...")
|
|
|
|
arrows_files = Utils.listdir(arrows_dir)
|
|
|
|
dataset_yaml_path = os.path.join(workspace_dir, "dataset.yaml")
|
|
|
|
import yaml
|
|
with open(dataset_yaml_path, "w", encoding="utf-8") as f:
|
|
yaml.dump({
|
|
"source": [os.path.join(arrows_dir, arrow_file) for arrow_file in arrows_files],
|
|
}, f)
|
|
|
|
print(f"dataset.yaml文件已生成: {dataset_yaml_path}")
|
|
|
|
idk_exec = os.path.join(
|
|
os.path.dirname(__file__), "hook_HYDiT_idk_run.py")
|
|
|
|
try:
|
|
from index_kits import __version__
|
|
except Exception as e:
|
|
subprocess.run([
|
|
sys.executable, "-m", "pip", "install", "-e", os.path.join(
|
|
HYDiT_tool_dir, "IndexKits")
|
|
], check=True)
|
|
# idk base -c dataset/yamls/porcelain.yaml -t dataset/porcelain/jsons/porcelain.json
|
|
|
|
dataset_json = os.path.join(workspace_dir, "dataset.json")
|
|
try:
|
|
subprocess.run([
|
|
sys.executable, idk_exec, "base", "-c", dataset_yaml_path, "-t", dataset_json,
|
|
], check=True)
|
|
except Exception as e:
|
|
raise Exception(f"生成index文件时出现异常,详细信息请查看控制台...")
|
|
|
|
print(f"index文件已生成: {dataset_json}")
|
|
|
|
dataset_mt_yaml_path = os.path.join(workspace_dir, "dataset_mt.yaml")
|
|
|
|
with open(dataset_mt_yaml_path, "w", encoding="utf-8") as f:
|
|
yaml.dump({
|
|
"src": [dataset_json],
|
|
"base_size": resolution,
|
|
"target_ratios": ["1:1", "3:4", "4:3", "16:9", "9:16"],
|
|
}, f)
|
|
|
|
print(f"dataset_mt.yaml文件已生成: {dataset_mt_yaml_path}")
|
|
|
|
# idk multireso -c dataset/yamls/porcelain_mt.yaml -t dataset/porcelain/jsons/porcelain_mt.json
|
|
|
|
dataset_mt_json = os.path.join(workspace_dir, "dataset_mt.json")
|
|
try:
|
|
subprocess.run([
|
|
sys.executable, idk_exec, "multireso", "-c", dataset_mt_yaml_path, "-t", dataset_mt_json,
|
|
], check=True)
|
|
except Exception as e:
|
|
raise Exception(f"生成多分辨率index文件时出现异常,详细信息请查看控制台...")
|
|
|
|
print(f"多分辨率index文件已生成: {dataset_mt_json}")
|
|
|
|
train_config = {}
|
|
|
|
train_config_path = os.path.join(workspace_dir, "mz_train_config.json")
|
|
|
|
if os.path.exists(train_config_path):
|
|
with open(train_config_path, "r", encoding="utf-8") as f:
|
|
train_config = json.load(f)
|
|
|
|
# "unet": (["auto"] + models + unet_models, {"default": "auto"}),
|
|
# "vae": (["auto"] + folders + vae_models, {"default": "auto"}),
|
|
# "text_encoder": (["auto"] + folders, {"default": "auto"}),
|
|
# "tokenizer": (["auto"] + folders, {"default": "auto"}),
|
|
# "t5_encoder": (["auto"] + folders, {"default": "auto"}),
|
|
|
|
epochs = args.get("epochs", 50)
|
|
output_dir = os.path.join(workspace_dir, "output")
|
|
|
|
datetime = time.strftime("%Y%m%d_%H%M%S", time.localtime())
|
|
output_name = f"{workspace_name}_{datetime}"
|
|
|
|
lora_ckpt = args.get("base_lora")
|
|
if lora_ckpt == "empty":
|
|
lora_ckpt = None
|
|
elif lora_ckpt == "latest":
|
|
lora_ckpt = None
|
|
loras = search_loras([workspace_dir])
|
|
if len(loras) > 0:
|
|
lora_ckpt = loras[0]
|
|
elif not os.path.exists(lora_ckpt):
|
|
raise Exception(f"未找到指定的lora文件: {lora_ckpt}")
|
|
|
|
print(f"使用指定的lora文件: {lora_ckpt}")
|
|
train_config.update({
|
|
"workspace_name": workspace_name,
|
|
"workspace_dir": workspace_dir,
|
|
"batch_size": args.get("batch_size", 1),
|
|
"ema_to_module": args.get("ema_to_module", "enable") == "enable",
|
|
"target_ratios": train_config.get("target_ratios", ["1:1", "3:4", "4:3", "16:9", "9:16"]),
|
|
"index_file": dataset_mt_json,
|
|
"epochs": epochs,
|
|
"ckpt_every": int(args.get("ckpt_every")),
|
|
"rope_img": f"base{resolution}",
|
|
"rope_real": args.get("rope_real", "enable") == "enable",
|
|
"image_size": resolution,
|
|
"unet_path": args.get("unet_path"),
|
|
"vae_ema_path": args.get("vae_ema_path"),
|
|
"text_encoder_path": args.get("text_encoder_path"),
|
|
"tokenizer_path": args.get("tokenizer_path"),
|
|
"t5_encoder_path": args.get("t5_encoder_path") if args.get("t5_encoder_path") != "none" else None,
|
|
"results_dir": output_dir,
|
|
"task_flag": output_name,
|
|
"lora_ckpt": lora_ckpt,
|
|
"rank": int(args.get("rank")),
|
|
})
|
|
|
|
if "target_modules" in train_config:
|
|
del train_config["target_modules"]
|
|
for key in advanced_config:
|
|
if key.startswith("target_modules_"):
|
|
if advanced_config[key] == "enable":
|
|
if "target_modules" not in train_config:
|
|
train_config["target_modules"] = []
|
|
train_config["target_modules"].append(
|
|
key.replace("target_modules_", ""))
|
|
elif advanced_config[key] == "enable" or advanced_config[key] == "disable":
|
|
train_config[key] = advanced_config[key] == "enable"
|
|
else:
|
|
train_config[key] = advanced_config[key]
|
|
|
|
sample_generate = args.get("sample_generate", "enable")
|
|
sample_prompt = args.get("sample_prompt", "")
|
|
sample_config_file = os.path.join(workspace_dir, "sample_config.json")
|
|
if sample_generate == "enable":
|
|
resolution = args.get("resolution")
|
|
sample_config = {
|
|
"prompt": sample_prompt,
|
|
"negative_prompt": "",
|
|
"cfg": 5.0,
|
|
"steps": 20,
|
|
"width": resolution,
|
|
"height": resolution,
|
|
}
|
|
with open(sample_config_file, "w", encoding="utf-8") as f:
|
|
json.dump(sample_config, f, indent=4, ensure_ascii=False)
|
|
|
|
train_config["sample_config_file"] = sample_config_file
|
|
else:
|
|
train_config["sample_config_file"] = None
|
|
|
|
with open(train_config_path, "w", encoding="utf-8") as f:
|
|
json.dump(train_config, f, indent=4, ensure_ascii=False)
|
|
|
|
run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir)
|
|
return (
|
|
args,
|
|
)
|
|
|
|
|
|
def run_hook_HYDiT_pyexec(train_config_path, train_config, HYDiT_tool_dir):
|
|
pb = Utils.progress_bar(0, "sdxl")
|
|
import traceback
|
|
import comfy.model_management
|
|
stop_server = None
|
|
|
|
is_running = True
|
|
|
|
def log_callback(log):
|
|
try:
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
except Exception as e:
|
|
stop_server()
|
|
return is_running
|
|
|
|
try:
|
|
resp = log
|
|
if resp.get("type") == "sample_images":
|
|
total_steps = resp.get("total_steps")
|
|
global_step = resp.get("global_step")
|
|
xy_img = get_sample_images(train_config)
|
|
if xy_img is None:
|
|
pb.update(
|
|
int(global_step), int(total_steps), None)
|
|
return is_running
|
|
max_side = max(xy_img.width, xy_img.height)
|
|
|
|
pb.update(
|
|
int(global_step), int(total_steps), ("JPEG", xy_img, max_side))
|
|
else:
|
|
print(f"LOG: {log}")
|
|
except Exception as e:
|
|
print(f"LOG: {log} e: {e} ")
|
|
print(f"stack: {traceback.format_exc()}")
|
|
return is_running
|
|
stop_server, port = Utils.Simple_Server(log_callback)
|
|
|
|
exec_pyfile = os.path.join(os.path.dirname(
|
|
__file__), "hook_HYDiT_run.py",)
|
|
|
|
try:
|
|
subprocess.run(
|
|
[sys.executable, "-m", "torch.distributed.launch", "--use_env", exec_pyfile, "--sys_path", HYDiT_tool_dir,
|
|
"--mz_master_port", str(port), "--train_config_file", train_config_path],
|
|
check=True,
|
|
)
|
|
stop_server()
|
|
is_running = False
|
|
except Exception as e:
|
|
stop_server()
|
|
is_running = False
|
|
raise Exception(f"训练失败!!! 具体报错信息请查看控制台...")
|
|
|
|
|
|
def get_sample_images(train_config):
|
|
if train_config.get("sample_config_file") is None:
|
|
return None
|
|
|
|
from PIL import Image
|
|
output_name = train_config.get("task_flag")
|
|
|
|
workspace_dir = train_config.get("workspace_dir")
|
|
|
|
sample_images_dir = os.path.join(workspace_dir, "sample_images")
|
|
|
|
pil_images = []
|
|
pre_render_texts_x = []
|
|
if os.path.exists(sample_images_dir):
|
|
image_files = Utils.listdir(sample_images_dir)
|
|
|
|
image_files = list(
|
|
filter(lambda x: x.lower().endswith(".png"), image_files))
|
|
|
|
# 筛选 output_name 前缀
|
|
image_files = list(
|
|
filter(lambda x: x.startswith(output_name), image_files))
|
|
|
|
image_files = sorted(image_files, key=lambda x: x)
|
|
|
|
for image_file in image_files:
|
|
pil_image = Image.open(os.path.join(sample_images_dir, image_file))
|
|
pil_images.append([pil_image])
|
|
pre_render_texts_x.append(image_file)
|
|
result = Utils.xy_image(
|
|
pre_render_images=pil_images,
|
|
pre_render_texts_x=pre_render_texts_x,
|
|
pre_render_texts_y=[""],
|
|
)
|
|
return result
|
|
return None
|
|
|
|
|
|
def get_HunYuanDiT_model_from_path(model_path, lora_path, width, height):
|
|
|
|
from hydit.modules.models import HUNYUAN_DIT_CONFIG
|
|
import importlib
|
|
import hydit.modules.models
|
|
importlib.reload(hydit.modules.models)
|
|
HunYuanDiT = hydit.modules.models.HunYuanDiT
|
|
|
|
from types import SimpleNamespace
|
|
|
|
# self.text_states_dim = args.text_states_dim
|
|
# self.text_states_dim_t5 = args.text_states_dim_t5
|
|
# self.text_len = args.text_len
|
|
# self.text_len_t5 = args.text_len_t5
|
|
# self.norm = args.norm
|
|
# use_flash_attn = args.infer_mode == 'fa' or args.use_flash_attn
|
|
|
|
args = SimpleNamespace()
|
|
args.learn_sigma = True
|
|
args.text_states_dim = 1024
|
|
args.text_states_dim_t5 = 2048
|
|
args.text_len = 77
|
|
args.text_len_t5 = 256
|
|
args.norm = "layer"
|
|
args.infer_mode = "torch"
|
|
args.use_fp16 = True
|
|
args.qk_norm = True
|
|
args.use_flash_attn = False
|
|
|
|
model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"]
|
|
|
|
latent_size = (width // 8, height // 8)
|
|
|
|
model = HunYuanDiT(args,
|
|
input_size=latent_size,
|
|
**model_config,
|
|
).half().to("cuda")
|
|
state_dict = torch.load(
|
|
model_path, map_location=lambda storage, loc: storage)
|
|
model.load_state_dict(state_dict, strict=False)
|
|
|
|
if lora_path is not None:
|
|
|
|
lora_path_files = Utils.listdir(lora_path)
|
|
safetensors_file = None
|
|
config_json_file = None
|
|
for file in lora_path_files:
|
|
if file.lower().endswith(".safetensors"):
|
|
safetensors_file = os.path.join(lora_path, file)
|
|
if file.lower().endswith(".json"):
|
|
config_json_file = os.path.join(lora_path, file)
|
|
|
|
if safetensors_file is None or config_json_file is None:
|
|
raise Exception("未找到对应的lora文件")
|
|
|
|
from peft import PeftConfig
|
|
|
|
# 查询后缀
|
|
from safetensors.torch import load_file
|
|
adapter_state_dict = load_file(safetensors_file)
|
|
|
|
peft_config = PeftConfig()
|
|
|
|
with open(config_json_file, "r", encoding="utf-8") as f:
|
|
loaded_attributes = json.load(f)
|
|
|
|
for key, value in loaded_attributes.items():
|
|
setattr(peft_config, key, value)
|
|
|
|
model.load_adapter(
|
|
adapter_state_dict=adapter_state_dict,
|
|
peft_config=peft_config,
|
|
)
|
|
model.merge_and_unload()
|
|
|
|
return model
|
|
|
|
|
|
def clean_unet():
|
|
unet_data = Utils.cache_get("HYDiT_UNET")
|
|
if unet_data is not None:
|
|
unet_data["model"].cpu()
|
|
Utils.cache_set("HYDiT_UNET", None)
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def clean_vae():
|
|
vae_data = Utils.cache_get("HYDiT_VAE_EMA")
|
|
if vae_data is not None:
|
|
vae_data["model"].cpu()
|
|
Utils.cache_set("HYDiT_VAE_EM", None)
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def clean_clip_text_encoder():
|
|
clip_text_encoder = Utils.cache_get("HYDiT_CLIP_TEXT_ENCODER")
|
|
if clip_text_encoder is not None:
|
|
clip_text_encoder["model"].cpu()
|
|
Utils.cache_set("HYDiT_CLIP_TEXT_ENCODER", None)
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def clean_t5_encoder():
|
|
t5_encoder = Utils.cache_get("HYDiT_T5_ENCODER")
|
|
if t5_encoder is not None:
|
|
t5_encoder["model"].cpu()
|
|
Utils.cache_set("HYDiT_T5_ENCODER", None)
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
from torch import nn
|
|
|
|
|
|
class CustomizeEmbedsModel(nn.Module):
|
|
dtype = torch.float16
|
|
x = torch.zeros(1, 1, 256, 2048)
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__()
|
|
|
|
def to(self, *args, **kwargs):
|
|
self.dtype = torch.float16
|
|
return self
|
|
|
|
def forward(self, *args, **kwargs):
|
|
if kwargs.get("output_hidden_states", False):
|
|
return {
|
|
"hidden_states": self.x.to("cuda"),
|
|
"input_ids": torch.zeros(1, 1),
|
|
}
|
|
return self.x
|
|
|
|
|
|
class CustomizeTokenizer(dict):
|
|
|
|
added_tokens_encoder = []
|
|
input_ids = torch.zeros(1, 256)
|
|
attention_mask = torch.zeros(1, 256)
|
|
|
|
def __init__(self):
|
|
self['added_tokens_encoder'] = self.added_tokens_encoder
|
|
self['input_ids'] = self.input_ids
|
|
self['attention_mask'] = self.attention_mask
|
|
|
|
def tokenize(self, text):
|
|
return text
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self
|
|
|
|
|
|
class CustomizeEmbeds():
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.tokenizer = CustomizeTokenizer()
|
|
self.model = CustomizeEmbedsModel()
|
|
self.max_length = 256
|
|
|
|
|
|
def MZ_HYDiTSimpleT2I_call(args={}):
|
|
check_required()
|
|
MZ_HYDiTInitClone_call(args)
|
|
args = check_model_auto_download(args)
|
|
|
|
HYDiT_tool_dir = os.path.join(
|
|
Utils.get_minus_zone_models_path(), "train_tools", "HunyuanDiT")
|
|
if HYDiT_tool_dir not in sys.path:
|
|
print(f"add {HYDiT_tool_dir} to sys.path")
|
|
sys.path.append(HYDiT_tool_dir)
|
|
|
|
from hydit.inference import SAMPLER_FACTORY, HUNYUAN_DIT_CONFIG, get_fill_resize_and_crop, get_2d_rotary_pos_embed
|
|
from hydit.modules.text_encoder import MT5Embedder
|
|
from hydit.utils.tools import set_seeds
|
|
from PIL import Image
|
|
from diffusers import schedulers, AutoencoderKL
|
|
from transformers import BertModel, BertTokenizer
|
|
|
|
SAMPLER_FACTORY["uni_pc"] = {
|
|
'scheduler': 'UniPCMultistepScheduler',
|
|
'name': 'UniPCMultistepScheduler',
|
|
'kwargs': {
|
|
'beta_schedule': 'scaled_linear',
|
|
'beta_start': 0.00085,
|
|
'beta_end': 0.03,
|
|
'prediction_type': 'v_prediction',
|
|
'trained_betas': None,
|
|
'solver_order': 2,
|
|
# 'algorithm_type': 'dpmsolver++',
|
|
}
|
|
}
|
|
SAMPLER_FACTORY["dpmpp_2m_karras"] = {
|
|
'scheduler': 'DPMSolverMultistepScheduler',
|
|
'name': 'DPMSolverMultistepScheduler',
|
|
'kwargs': {
|
|
'beta_schedule': 'scaled_linear',
|
|
'beta_start': 0.00085,
|
|
'beta_end': 0.03,
|
|
'prediction_type': 'v_prediction',
|
|
'trained_betas': None,
|
|
'solver_order': 2,
|
|
'algorithm_type': 'dpmsolver++',
|
|
"use_karras_sigmas": True,
|
|
}
|
|
}
|
|
|
|
width = args.get("width")
|
|
height = args.get("height")
|
|
|
|
unet_path = args.get("unet_path")
|
|
lora_path = args.get("lora_path")
|
|
unet_cache_key = f"HYDiT_UNET"
|
|
unet_data = Utils.cache_get(unet_cache_key)
|
|
if unet_data is None or unet_data.get("unet_path") != unet_path or unet_data.get("lora_path") != lora_path:
|
|
clean_unet()
|
|
_lora_path = lora_path if lora_path != "none" else None
|
|
unet = get_HunYuanDiT_model_from_path(
|
|
args.get("unet_path"), _lora_path, width, height)
|
|
Utils.cache_set(unet_cache_key, {
|
|
"model": unet,
|
|
"unet_path": unet_path,
|
|
"lora_path": lora_path,
|
|
})
|
|
else:
|
|
unet = unet_data["model"]
|
|
|
|
text_encoder_path = args.get("text_encoder_path")
|
|
clip_text_encoder_data = Utils.cache_get("HYDiT_CLIP_TEXT_ENCODER")
|
|
|
|
if clip_text_encoder_data is None or clip_text_encoder_data.get("text_encoder_path") != text_encoder_path:
|
|
clean_clip_text_encoder()
|
|
clip_text_encoder = BertModel.from_pretrained(
|
|
text_encoder_path, False, revision=None, local_files_only=True).to("cuda")
|
|
Utils.cache_set("HYDiT_CLIP_TEXT_ENCODER", {
|
|
"model": clip_text_encoder,
|
|
"text_encoder_path": text_encoder_path,
|
|
})
|
|
else:
|
|
clip_text_encoder = clip_text_encoder_data["model"]
|
|
|
|
tokenizer_path = args.get("tokenizer_path")
|
|
text_tokenizer = BertTokenizer.from_pretrained(
|
|
tokenizer_path, revision=None, local_files_only=True)
|
|
|
|
vae_ema_path = args.get("vae_ema_path")
|
|
vae_ema_data = Utils.cache_get("HYDiT_VAE_EMA")
|
|
if vae_ema_data is None or vae_ema_data.get("vae_ema_path") != vae_ema_path:
|
|
clean_vae()
|
|
if os.path.isdir(vae_ema_path):
|
|
vae = AutoencoderKL.from_pretrained(
|
|
vae_ema_path, local_files_only=True).to("cuda")
|
|
else:
|
|
vae = AutoencoderKL.from_single_file(
|
|
vae_ema_path, local_files_only=True).to("cuda")
|
|
|
|
Utils.cache_set("HYDiT_VAE_EMA", {
|
|
"model": vae,
|
|
"vae_ema_path": vae_ema_path,
|
|
})
|
|
else:
|
|
vae = vae_ema_data["model"]
|
|
|
|
t5_encoder_path = args.get("t5_encoder_path")
|
|
if t5_encoder_path == "none":
|
|
t5_encoder = CustomizeEmbeds()
|
|
else:
|
|
t5_encoder_data = Utils.cache_get("HYDiT_T5_ENCODER")
|
|
if t5_encoder_data is None or t5_encoder_data.get("t5_encoder_path") != t5_encoder_path:
|
|
clean_t5_encoder()
|
|
t5_encoder = MT5Embedder(
|
|
t5_encoder_path, torch_dtype=torch.float16, max_length=256)
|
|
Utils.cache_set("HYDiT_T5_ENCODER", {
|
|
"model": t5_encoder,
|
|
"t5_encoder_path": t5_encoder_path,
|
|
})
|
|
else:
|
|
t5_encoder = t5_encoder_data["model"]
|
|
|
|
sampler = args.get("scheduler")
|
|
scheduler = SAMPLER_FACTORY[sampler]['scheduler']
|
|
kwargs = SAMPLER_FACTORY[sampler]['kwargs']
|
|
# Build scheduler according to the sampler.
|
|
scheduler_class = getattr(schedulers, scheduler)
|
|
scheduler = scheduler_class(**kwargs)
|
|
|
|
scheduler.set_timesteps(args.get("steps"), "cuda")
|
|
|
|
from hydit.diffusion.pipeline import StableDiffusionPipeline
|
|
pipeline = StableDiffusionPipeline(vae=vae,
|
|
text_encoder=clip_text_encoder,
|
|
tokenizer=text_tokenizer,
|
|
unet=unet,
|
|
scheduler=scheduler,
|
|
feature_extractor=None,
|
|
safety_checker=None,
|
|
requires_safety_checker=False,
|
|
progress_bar_config=None,
|
|
embedder_t5=t5_encoder,
|
|
infer_mode="torch",
|
|
)
|
|
|
|
pipeline = pipeline.to("cuda")
|
|
|
|
seed = args.get("seed", 0)
|
|
target_height = int((height // 16) * 16)
|
|
target_width = int((width // 16) * 16)
|
|
prompt = args.get("prompt", "").strip()
|
|
negative_prompt = args.get("negative_prompt", "").strip()
|
|
|
|
batch_size = 1
|
|
style = torch.as_tensor([0, 0] * batch_size, device="cuda")
|
|
|
|
src_size_cond = (target_width, target_height)
|
|
size_cond = list(src_size_cond) + [target_width, target_height, 0, 0]
|
|
image_meta_size = torch.as_tensor(
|
|
[size_cond] * 2 * batch_size, device="cuda",)
|
|
|
|
def calc_rope(height, width):
|
|
model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"]
|
|
patch_size = model_config['patch_size']
|
|
head_size = model_config['hidden_size'] // model_config['num_heads']
|
|
th = height // 8 // patch_size
|
|
tw = width // 8 // patch_size
|
|
base_size = 512 // 8 // patch_size
|
|
start, stop = get_fill_resize_and_crop((th, tw), base_size)
|
|
sub_args = [start, stop, (th, tw)]
|
|
rope = get_2d_rotary_pos_embed(head_size, *sub_args)
|
|
return rope
|
|
|
|
freqs_cis_img = calc_rope(target_height, target_width)
|
|
guidance_scale = args.get("cfg", 7.0)
|
|
infer_steps = args.get("steps", 20)
|
|
|
|
with torch.inference_mode():
|
|
with torch.cuda.amp.autocast():
|
|
|
|
pbar = Utils.progress_bar(0, "sdxl")
|
|
preview = pbar.get_previewer()
|
|
|
|
def callback(step: int, timestep: int, latents: torch.FloatTensor, pred_x0):
|
|
try:
|
|
pil_img = preview.decode_latent_to_preview_image(
|
|
None,
|
|
latents,
|
|
)[1]
|
|
pbar.update(step, infer_steps, pil_img)
|
|
except Exception as e:
|
|
print("decode_tensors error:", e)
|
|
generator = set_seeds(seed, device="cuda")
|
|
samples = pipeline(
|
|
height=height,
|
|
width=width,
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
num_images_per_prompt=batch_size,
|
|
guidance_scale=guidance_scale,
|
|
num_inference_steps=infer_steps,
|
|
style=style,
|
|
return_dict=False,
|
|
generator=generator,
|
|
use_fp16=True,
|
|
learn_sigma=True,
|
|
freqs_cis_img=freqs_cis_img,
|
|
image_meta_size=image_meta_size,
|
|
callback=callback,
|
|
callback_steps=2,
|
|
)[0]
|
|
|
|
if type(samples) == list:
|
|
pil_image = samples[0]
|
|
else:
|
|
pil_image = samples
|
|
|
|
tensor_image = Utils.pil2tensor(pil_image)
|
|
|
|
keep_device = args.get("keep_device", "disable") == "enable"
|
|
if not keep_device:
|
|
del pipeline
|
|
del unet
|
|
del clip_text_encoder
|
|
del vae
|
|
del t5_encoder
|
|
clean_unet()
|
|
clean_clip_text_encoder()
|
|
clean_vae()
|
|
torch.cuda.empty_cache()
|
|
|
|
return (Utils.list_tensor2tensor([tensor_image]),)
|
|
|
|
|
|
def search_loras(wdirs):
|
|
loras = []
|
|
for wdir in wdirs:
|
|
for root, dirs, files in os.walk(wdir):
|
|
# 排除隐藏文件夹
|
|
dirs[:] = [d for d in dirs if not d.startswith('.')]
|
|
# 存在adapter_config.json文件并且存在adapter_model.safetensors文件
|
|
for file in files:
|
|
if file == "adapter_config.json":
|
|
lora_path = os.path.join(root, "adapter_model.safetensors")
|
|
if os.path.exists(lora_path):
|
|
loras.append(root)
|
|
|
|
# 按时间倒序
|
|
loras = sorted(loras, key=lambda x: os.path.getmtime(x), reverse=True)
|
|
return loras
|