'update'
This commit is contained in:
wailovet
+5
-1
@@ -159,4 +159,8 @@ cython_debug/
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
#.idea/
|
||||
kohya_ss
|
||||
copytoww.bat
|
||||
exclude.txt
|
||||
.vscode/settings.json
|
||||
|
||||
+141
@@ -0,0 +1,141 @@
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import folder_paths
|
||||
import importlib
|
||||
from .mz_train_tools_utils import Utils
|
||||
from . import mz_train_tools_core
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
}
|
||||
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
}
|
||||
|
||||
AUTHOR_NAME = "MinusZone"
|
||||
CATEGORY_NAME = f"{AUTHOR_NAME} - TrainTools"
|
||||
|
||||
|
||||
class MZ_KohyaSSInitWorkspace:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"workspace_name": ("STRING", {"default": ""}),
|
||||
"branch": ("STRING", {"default": "main"}),
|
||||
"seed": ("INT", {"default": 0}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MZ_TT_SS_WorkspaceConfig",)
|
||||
RETURN_NAMES = ("workspace_config",)
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
return mz_train_tools_core.MZ_KohyaSSInitWorkspace_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSInitWorkspace"] = MZ_KohyaSSInitWorkspace
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSInitWorkspace"] = f"{AUTHOR_NAME} - KohyaSSInitWorkspace"
|
||||
|
||||
|
||||
class MZ_ImagesCopyWorkspace:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
|
||||
"images": ("IMAGE",),
|
||||
"enable_bucket": (["enable", "disable"], {"default": "enable"}),
|
||||
"resolution": ("INT", {"default": 512}),
|
||||
"num_repeats": ("INT", {"default": 1}),
|
||||
"batch_size": ("INT", {"default": 1}),
|
||||
# "class_name": ("STRING", {"default": "girl", "dynamicPrompts": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (f"STRING",)
|
||||
RETURN_NAMES = ("workspace_images_dir",)
|
||||
|
||||
# OUTPUT_NODE = True
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
return mz_train_tools_core.MZ_ImageSelecter_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_ImagesCopyWorkspace"] = MZ_ImagesCopyWorkspace
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_ImagesCopyWorkspace"] = f"{AUTHOR_NAME} - ImagesCopyWorkspace"
|
||||
|
||||
|
||||
class MZ_KohyaSSUseConfig:
|
||||
|
||||
train_config_template_dir = os.path.join(
|
||||
os.path.dirname(__file__), "configs", "kohya_ss_lora"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
train_config_templates = os.listdir(s.train_config_template_dir)
|
||||
# 去掉json后缀
|
||||
train_config_templates = [os.path.splitext(x)[0]
|
||||
for x in train_config_templates]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"workspace_config": ("MZ_TT_SS_WorkspaceConfig",),
|
||||
"workspace_images_dir": ("STRING", {"forceInput": True}),
|
||||
"train_config_template": (train_config_templates,),
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), )
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (f"MZ_TT_SS_TrainConfig",)
|
||||
RETURN_NAMES = ("train_config",)
|
||||
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
kwargs["train_config_template_dir"] = self.train_config_template_dir
|
||||
return mz_train_tools_core.MZ_KohyaSSUseConfig_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSUseConfig"] = MZ_KohyaSSUseConfig
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSUseConfig"] = f"{AUTHOR_NAME} - KohyaSSUseConfig"
|
||||
|
||||
|
||||
class MZ_KohyaSSTrain:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"train_config": ("MZ_TT_SS_TrainConfig",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (f"STRING",)
|
||||
RETURN_NAMES = ("train_output_dir",)
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "start"
|
||||
CATEGORY = CATEGORY_NAME
|
||||
|
||||
def start(self, **kwargs):
|
||||
importlib.reload(mz_train_tools_core)
|
||||
return mz_train_tools_core.MZ_KohyaSSTrain_call(kwargs)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSTrain
|
||||
NODE_DISPLAY_NAME_MAPPINGS["MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain"
|
||||
@@ -0,0 +1,35 @@
|
||||
{
|
||||
"metadata": {
|
||||
"train_type": "lora_sd1_5"
|
||||
},
|
||||
"train_config": {
|
||||
"pretrained_model_name_or_path": "",
|
||||
"max_train_steps": "4500",
|
||||
"xformers": true,
|
||||
"sdpa": false,
|
||||
"fp8_base": false,
|
||||
"mixed_precision": "fp16",
|
||||
"cache_latents": true,
|
||||
"cache_latents_to_disk": true,
|
||||
"network_dim": "16",
|
||||
"network_alpha": "8",
|
||||
"network_module": "networks.lora",
|
||||
"network_train_unet_only": true,
|
||||
"lr_scheduler": "cosine_with_restarts",
|
||||
"optimizer_type": "AdamW",
|
||||
"unet_lr": "1e-5",
|
||||
"text_encoder_lr": "1e-5",
|
||||
"save_every_n_epochs": "20",
|
||||
"shuffle_caption": true,
|
||||
"lr_warmup_steps": "0",
|
||||
"save_precision": "fp16",
|
||||
"lr_scheduler_num_cycles": "1",
|
||||
"persistent_data_loader_workers": true,
|
||||
"no_metadata": true,
|
||||
"noise_offset": "0.1",
|
||||
"output_dir": "",
|
||||
"output_name": "",
|
||||
"no_half_vae": true,
|
||||
"lowram": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
{
|
||||
"metadata": {
|
||||
"train_type": "lora_sdxl"
|
||||
},
|
||||
"train_config": {
|
||||
"pretrained_model_name_or_path": "",
|
||||
"max_train_steps": "4500",
|
||||
"xformers": true,
|
||||
"sdpa": false,
|
||||
"fp8_base": false,
|
||||
"mixed_precision": "fp16",
|
||||
"cache_latents": true,
|
||||
"cache_latents_to_disk": true,
|
||||
"network_dim": "16",
|
||||
"network_alpha": "8",
|
||||
"network_module": "networks.lora",
|
||||
"network_train_unet_only": true,
|
||||
"lr_scheduler": "cosine_with_restarts",
|
||||
"optimizer_type": "AdamW",
|
||||
"unet_lr": "1e-5",
|
||||
"text_encoder_lr": "1e-5",
|
||||
"save_every_n_epochs": "20",
|
||||
"shuffle_caption": true,
|
||||
"lr_warmup_steps": "0",
|
||||
"save_precision": "fp16",
|
||||
"lr_scheduler_num_cycles": "1",
|
||||
"persistent_data_loader_workers": true,
|
||||
"no_metadata": true,
|
||||
"noise_offset": "0.1",
|
||||
"output_dir": "",
|
||||
"output_name": "",
|
||||
"no_half_vae": true,
|
||||
"lowram": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"unk_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"single_word": false,
|
||||
"lstrip": false,
|
||||
"rstrip": false,
|
||||
"normalized": true,
|
||||
"__type": "AddedToken"
|
||||
},
|
||||
"bos_token": {
|
||||
"content": "<|startoftext|>",
|
||||
"single_word": false,
|
||||
"lstrip": false,
|
||||
"rstrip": false,
|
||||
"normalized": true,
|
||||
"__type": "AddedToken"
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"single_word": false,
|
||||
"lstrip": false,
|
||||
"rstrip": false,
|
||||
"normalized": true,
|
||||
"__type": "AddedToken"
|
||||
},
|
||||
"pad_token": "<|endoftext|>",
|
||||
"add_prefix_space": false,
|
||||
"errors": "replace",
|
||||
"do_lower_case": true,
|
||||
"name_or_path": "openai/clip-vit-base-patch32",
|
||||
"model_max_length": 77,
|
||||
"special_tokens_map_file": "./special_tokens_map.json",
|
||||
"tokenizer_class": "CLIPTokenizer"
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
model:
|
||||
base_learning_rate: 1.0e-04
|
||||
target: ldm.models.diffusion.ddpm.LatentDiffusion
|
||||
params:
|
||||
linear_start: 0.00085
|
||||
linear_end: 0.0120
|
||||
num_timesteps_cond: 1
|
||||
log_every_t: 200
|
||||
timesteps: 1000
|
||||
first_stage_key: "jpg"
|
||||
cond_stage_key: "txt"
|
||||
image_size: 64
|
||||
channels: 4
|
||||
cond_stage_trainable: false # Note: different from the one we trained before
|
||||
conditioning_key: crossattn
|
||||
monitor: val/loss_simple_ema
|
||||
scale_factor: 0.18215
|
||||
use_ema: False
|
||||
|
||||
scheduler_config: # 10000 warmup steps
|
||||
target: ldm.lr_scheduler.LambdaLinearScheduler
|
||||
params:
|
||||
warm_up_steps: [ 10000 ]
|
||||
cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
|
||||
f_start: [ 1.e-6 ]
|
||||
f_max: [ 1. ]
|
||||
f_min: [ 1. ]
|
||||
|
||||
unet_config:
|
||||
target: ldm.modules.diffusionmodules.openaimodel.UNetModel
|
||||
params:
|
||||
image_size: 32 # unused
|
||||
in_channels: 4
|
||||
out_channels: 4
|
||||
model_channels: 320
|
||||
attention_resolutions: [ 4, 2, 1 ]
|
||||
num_res_blocks: 2
|
||||
channel_mult: [ 1, 2, 4, 4 ]
|
||||
num_heads: 8
|
||||
use_spatial_transformer: True
|
||||
transformer_depth: 1
|
||||
context_dim: 768
|
||||
use_checkpoint: True
|
||||
legacy: False
|
||||
|
||||
first_stage_config:
|
||||
target: ldm.models.autoencoder.AutoencoderKL
|
||||
params:
|
||||
embed_dim: 4
|
||||
monitor: val/rec_loss
|
||||
ddconfig:
|
||||
double_z: true
|
||||
z_channels: 4
|
||||
resolution: 256
|
||||
in_channels: 3
|
||||
out_ch: 3
|
||||
ch: 128
|
||||
ch_mult:
|
||||
- 1
|
||||
- 2
|
||||
- 4
|
||||
- 4
|
||||
num_res_blocks: 2
|
||||
attn_resolutions: []
|
||||
dropout: 0.0
|
||||
lossconfig:
|
||||
target: torch.nn.Identity
|
||||
|
||||
cond_stage_config:
|
||||
target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
|
||||
@@ -0,0 +1,93 @@
|
||||
model:
|
||||
target: sgm.models.diffusion.DiffusionEngine
|
||||
params:
|
||||
scale_factor: 0.13025
|
||||
disable_first_stage_autocast: True
|
||||
|
||||
denoiser_config:
|
||||
target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser
|
||||
params:
|
||||
num_idx: 1000
|
||||
|
||||
scaling_config:
|
||||
target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling
|
||||
discretization_config:
|
||||
target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization
|
||||
|
||||
network_config:
|
||||
target: sgm.modules.diffusionmodules.openaimodel.UNetModel
|
||||
params:
|
||||
adm_in_channels: 2816
|
||||
num_classes: sequential
|
||||
use_checkpoint: True
|
||||
in_channels: 4
|
||||
out_channels: 4
|
||||
model_channels: 320
|
||||
attention_resolutions: [4, 2]
|
||||
num_res_blocks: 2
|
||||
channel_mult: [1, 2, 4]
|
||||
num_head_channels: 64
|
||||
use_linear_in_transformer: True
|
||||
transformer_depth: [1, 2, 10]
|
||||
context_dim: 2048
|
||||
spatial_transformer_attn_type: softmax-xformers
|
||||
|
||||
conditioner_config:
|
||||
target: sgm.modules.GeneralConditioner
|
||||
params:
|
||||
emb_models:
|
||||
- is_trainable: False
|
||||
input_key: txt
|
||||
target: sgm.modules.encoders.modules.FrozenCLIPEmbedder
|
||||
params:
|
||||
layer: hidden
|
||||
layer_idx: 11
|
||||
|
||||
- is_trainable: False
|
||||
input_key: txt
|
||||
target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2
|
||||
params:
|
||||
arch: ViT-bigG-14
|
||||
version: laion2b_s39b_b160k
|
||||
freeze: True
|
||||
layer: penultimate
|
||||
always_return_pooled: True
|
||||
legacy: False
|
||||
|
||||
- is_trainable: False
|
||||
input_key: original_size_as_tuple
|
||||
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
|
||||
params:
|
||||
outdim: 256
|
||||
|
||||
- is_trainable: False
|
||||
input_key: crop_coords_top_left
|
||||
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
|
||||
params:
|
||||
outdim: 256
|
||||
|
||||
- is_trainable: False
|
||||
input_key: target_size_as_tuple
|
||||
target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
|
||||
params:
|
||||
outdim: 256
|
||||
|
||||
first_stage_config:
|
||||
target: sgm.models.autoencoder.AutoencoderKL
|
||||
params:
|
||||
embed_dim: 4
|
||||
monitor: val/rec_loss
|
||||
ddconfig:
|
||||
attn_type: vanilla-xformers
|
||||
double_z: true
|
||||
z_channels: 4
|
||||
resolution: 256
|
||||
in_channels: 3
|
||||
out_ch: 3
|
||||
ch: 128
|
||||
ch_mult: [1, 2, 4, 4]
|
||||
num_res_blocks: 2
|
||||
attn_resolutions: []
|
||||
dropout: 0.0
|
||||
lossconfig:
|
||||
target: torch.nn.Identity
|
||||
@@ -0,0 +1,28 @@
|
||||
prompts_file: ""
|
||||
pretrained_model:
|
||||
name_or_path: "" # you can also use .ckpt or .safetensors models
|
||||
v2: false # true if model is v2.x
|
||||
v_pred: false # true if model uses v-prediction
|
||||
network:
|
||||
type: "c3lier" # or "c3lier" or "lierla"
|
||||
rank: 4
|
||||
alpha: 1.0
|
||||
training_method: "noxattn"
|
||||
train:
|
||||
precision: "float16"
|
||||
noise_scheduler: "ddim" # or "ddpm", "lms", "euler_a"
|
||||
iterations: 1000
|
||||
lr: 0.0002
|
||||
optimizer: "AdamW"
|
||||
lr_scheduler: "constant"
|
||||
max_denoising_steps: 50
|
||||
save:
|
||||
name: "temp"
|
||||
path: "./models"
|
||||
per_steps: 500
|
||||
precision: "float16"
|
||||
logging:
|
||||
use_wandb: false
|
||||
verbose: false
|
||||
other:
|
||||
use_xformers: true
|
||||
@@ -0,0 +1,9 @@
|
||||
- target: "" # 什么词可以消除积极的概念
|
||||
positive: "" # 要擦除的概念
|
||||
unconditional: "" # 与积极概念区别的词
|
||||
neutral: "" # 调节目标的起点
|
||||
action: "enhance" # 擦除或增强
|
||||
guidance_scale: 4
|
||||
resolution: 768
|
||||
dynamic_resolution: true
|
||||
batch_size: 1
|
||||
@@ -0,0 +1,129 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
os.environ["HTTP_PROXY"] = "http://127.0.0.1:0"
|
||||
os.environ["HTTPS_PROXY"] = "http://127.0.0.1:0"
|
||||
import logging
|
||||
import sys
|
||||
import json
|
||||
import importlib
|
||||
import argparse
|
||||
import toml
|
||||
|
||||
|
||||
def config2args(train_parser: argparse.ArgumentParser, config):
|
||||
|
||||
config_args_list = []
|
||||
for key, value in config.items():
|
||||
if type(value) == bool:
|
||||
if value:
|
||||
config_args_list.append(f"--{key}")
|
||||
else:
|
||||
config_args_list.append(f"--{key}")
|
||||
config_args_list.append(str(value))
|
||||
args = train_parser.parse_args(config_args_list)
|
||||
return args
|
||||
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
import numpy as np
|
||||
import tempfile
|
||||
import safetensors.torch
|
||||
|
||||
import hook_kohya_ss_utils
|
||||
|
||||
|
||||
def sample_images(self, *args, **kwargs):
|
||||
# accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet
|
||||
global_step = args[3]
|
||||
last_noise_pred = hook_kohya_ss_utils.running_info.get(
|
||||
"last_noise_pred", None)
|
||||
noise_pred_latent_path = None
|
||||
if last_noise_pred is not None:
|
||||
noise_pred_latent_path = os.path.join(
|
||||
tempfile.gettempdir(), f"sample_images_{global_step % 5}.latent")
|
||||
|
||||
output = {}
|
||||
output["latent_tensor"] = last_noise_pred
|
||||
output["latent_format_version_0"] = torch.tensor([])
|
||||
safetensors.torch.save_file(output, noise_pred_latent_path)
|
||||
|
||||
LOG(json.dumps({
|
||||
"type": "sample_images",
|
||||
"global_step": global_step,
|
||||
"latent": noise_pred_latent_path,
|
||||
}))
|
||||
|
||||
|
||||
def run_lora_sd1_5():
|
||||
hook_kohya_ss_utils.hook_kohya_ss()
|
||||
import train_network
|
||||
train_network.NetworkTrainer.sample_images = sample_images
|
||||
trainer = train_network.NetworkTrainer()
|
||||
train_config = json.loads(train_config_json)
|
||||
train_args = config2args(train_network.setup_parser(), train_config)
|
||||
|
||||
trainer.train(train_args)
|
||||
|
||||
|
||||
def run_lora_sdxl():
|
||||
hook_kohya_ss_utils.hook_kohya_ss()
|
||||
import hook_kohya_ss_utils
|
||||
import sdxl_train_network
|
||||
sdxl_train_network.SdxlNetworkTrainer.sample_images = sample_images
|
||||
trainer = sdxl_train_network.SdxlNetworkTrainer()
|
||||
train_config = json.loads(train_config_json)
|
||||
train_args = config2args(sdxl_train_network.setup_parser(), train_config)
|
||||
trainer.train(train_args)
|
||||
|
||||
|
||||
func_map = {
|
||||
"run_lora_sd1_5": run_lora_sd1_5,
|
||||
"run_lora_sdxl": run_lora_sdxl,
|
||||
}
|
||||
|
||||
|
||||
import socket
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
udp_port = 0
|
||||
|
||||
|
||||
def is_connected():
|
||||
try:
|
||||
sock.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR)
|
||||
except socket.error:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def LOG(log):
|
||||
sock.sendto(log.encode(), ('127.0.0.1', udp_port))
|
||||
if not is_connected():
|
||||
raise Exception("Connection closed")
|
||||
|
||||
|
||||
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("--train_func", type=str, default="")
|
||||
parser.add_argument("--udp_port", type=int, default=0)
|
||||
args = parser.parse_args()
|
||||
|
||||
udp_port = args.udp_port
|
||||
|
||||
print(f"udp_port = {udp_port}")
|
||||
|
||||
sys_path = args.sys_path
|
||||
if sys_path != "":
|
||||
sys.path.append(sys_path)
|
||||
|
||||
train_config_json = args.train_config_json
|
||||
|
||||
train_func = args.train_func
|
||||
if train_func == "":
|
||||
raise Exception("train_func is empty")
|
||||
|
||||
func_map[train_func]()
|
||||
@@ -0,0 +1,340 @@
|
||||
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from typing import *
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
|
||||
from transformers import CLIPTokenizer
|
||||
import requests
|
||||
|
||||
_requests_get = requests.get
|
||||
|
||||
source_replacement_table = {
|
||||
"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml": os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-v1.5", "v1-inference.yaml"),
|
||||
"https://raw.githubusercontent.com/Stability-AI/generative-models/main/configs/inference/sd_xl_base.yaml": os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-xl", "sd_xl_base.yaml"),
|
||||
"https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/tokenizer_config.json": os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "clip-vit-large-patch14", "tokenizer_config.json"),
|
||||
"https://huggingface.co/api/models/stabilityai/stable-diffusion-3-medium-diffusers/revision/main": os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-3-medium-diffusers", "revision.json"),
|
||||
}
|
||||
|
||||
source_replacement_dir = {
|
||||
"https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers/resolve/b1148b4028b9ec56ebd36444c193d56aeff7ab56": os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-3-medium-diffusers"),
|
||||
}
|
||||
|
||||
|
||||
class DictWrapper:
|
||||
def __init__(self, d):
|
||||
self.d = d
|
||||
|
||||
def __getattribute__(self, name: str):
|
||||
if name == "content":
|
||||
return self.d["content"]
|
||||
if name == "raise_for_status":
|
||||
return lambda: None
|
||||
if name == "json":
|
||||
return lambda: json.loads(self.d["content"])
|
||||
if name == "status_code":
|
||||
return self.d["status_code"]
|
||||
if name == "headers":
|
||||
return {
|
||||
"Location": self.d["Location"],
|
||||
"Content-Length": len(self.d["content"]),
|
||||
}
|
||||
if name == "request":
|
||||
return None
|
||||
return super().__getattribute__(name)
|
||||
|
||||
|
||||
def request_wrapper(*args, **kwargs):
|
||||
url = args[1]
|
||||
print(f"request_wrapper requesting {url}")
|
||||
if url in source_replacement_table:
|
||||
with open(source_replacement_table[url], "rb") as f:
|
||||
return DictWrapper({
|
||||
"Location": url,
|
||||
"content": f.read(),
|
||||
"status_code": 200,
|
||||
})
|
||||
|
||||
print(f"request_wrapper requesting {url} from original requests")
|
||||
return _requests_get(*args, **kwargs)
|
||||
|
||||
|
||||
from requests import api
|
||||
from requests import Session
|
||||
last_request = api.request
|
||||
api.request = request_wrapper
|
||||
|
||||
|
||||
def Session_request_wrapper(cls, method, url, **kwargs):
|
||||
# print(f"Session_request_wrapper requesting {url}")
|
||||
# print(f"Session_request_wrapper requesting kwargs: {kwargs}")
|
||||
if url in source_replacement_table:
|
||||
with open(source_replacement_table[url], "rb") as f:
|
||||
return DictWrapper({
|
||||
"Location": url,
|
||||
"content": f.read(),
|
||||
"status_code": 200,
|
||||
})
|
||||
|
||||
for k, v in source_replacement_dir.items():
|
||||
if url.startswith(k):
|
||||
file_path = source_replacement_dir[k] + url[len(k):]
|
||||
# print(
|
||||
# f"source_replacement_dir:{k}||||||||||||||||||||| {source_replacement_dir[k]} ||||||||||||||||||| {file_path}")
|
||||
with open(file_path, "rb") as f:
|
||||
return DictWrapper({
|
||||
"Location": url,
|
||||
"content": f.read(),
|
||||
"status_code": 200,
|
||||
})
|
||||
raise NotImplementedError("Session.request is not supported")
|
||||
|
||||
|
||||
Session.request = Session_request_wrapper
|
||||
|
||||
import huggingface_hub.file_download
|
||||
|
||||
|
||||
def _hf_hub_download_to_cache_dir(repo_id, filename, *args, **kwargs):
|
||||
print(f"_hf_hub_download_to_cache_dir: {args}")
|
||||
print(f"_hf_hub_download_to_cache_dir: {kwargs}")
|
||||
if repo_id == "stabilityai/stable-diffusion-3-medium-diffusers":
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-3-medium-diffusers", filename)
|
||||
raise NotImplementedError("_hf_hub_download_to_cache_dir is not supported")
|
||||
|
||||
|
||||
huggingface_hub.file_download._hf_hub_download_to_cache_dir = _hf_hub_download_to_cache_dir
|
||||
|
||||
import diffusers.loaders.single_file
|
||||
original_snapshot_download = diffusers.loaders.single_file.snapshot_download
|
||||
|
||||
|
||||
def _snapshot_download(repo_id, *args, **kwargs):
|
||||
if repo_id == "stabilityai/stable-diffusion-3-medium-diffusers":
|
||||
return os.path.join(
|
||||
os.path.dirname(__file__), "configs", "models_config", "stable-diffusion-3-medium-diffusers",)
|
||||
return original_snapshot_download(*args, **kwargs)
|
||||
|
||||
|
||||
diffusers.loaders.single_file.snapshot_download = _snapshot_download
|
||||
|
||||
original_load_target_model = None
|
||||
|
||||
|
||||
def setup_logging(*args, **kwargs):
|
||||
pass
|
||||
|
||||
|
||||
clip_large_tokenizer = None
|
||||
clip_big_tokenizer = None
|
||||
|
||||
|
||||
class TokenizersWrapper:
|
||||
typed = None
|
||||
model_max_length = 77
|
||||
|
||||
def __init__(self, t):
|
||||
self.model_max_length = 77
|
||||
self.typed = t
|
||||
|
||||
def __getattribute__(self, name: str):
|
||||
# print(f"TokenizersWrapper.__getattribute__ {name}")
|
||||
if name == "model_max_length":
|
||||
return 77
|
||||
try:
|
||||
typed = object.__getattribute__(self, "typed")
|
||||
if typed == "clip_large" and clip_large_tokenizer is not None:
|
||||
return clip_large_tokenizer.__getattribute__(name)
|
||||
if typed == "clip_big" and clip_big_tokenizer is not None:
|
||||
return clip_big_tokenizer.__getattribute__(name)
|
||||
except:
|
||||
pass
|
||||
|
||||
return object.__getattribute__(self, name)
|
||||
|
||||
def __call__(self, *args, **kargs):
|
||||
if self.typed == "clip_large":
|
||||
return clip_large_tokenizer(*args, **kargs)
|
||||
if self.typed == "clip_big":
|
||||
return clip_big_tokenizer(*args, **kargs)
|
||||
|
||||
raise NotImplementedError(
|
||||
f"TokenizersWrapper: {self.typed} is not supported")
|
||||
|
||||
|
||||
def load_tokenizers(*args, **kwargs):
|
||||
return TokenizersWrapper("clip_large")
|
||||
|
||||
|
||||
def load_sdxl_tokenizers(*args, **kwargs):
|
||||
return [TokenizersWrapper("clip_large"), TokenizersWrapper("clip_big")]
|
||||
|
||||
|
||||
original_conditional_loss = None
|
||||
|
||||
|
||||
running_info = {}
|
||||
|
||||
|
||||
def conditional_loss(*args, **kwargs):
|
||||
running_info["last_noise_pred"] = args[0]
|
||||
return original_conditional_loss(*args, **kwargs)
|
||||
|
||||
|
||||
def decode_latents(vae, latents):
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
latents = latents.to(dtype=vae.dtype).to(device)
|
||||
vae = vae.to(device)
|
||||
latents = 1 / 0.18215 * latents
|
||||
image = vae.decode(latents).sample
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
|
||||
image = image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
|
||||
return image
|
||||
|
||||
|
||||
def hook_kohya_ss():
|
||||
import library.utils
|
||||
import library.train_util
|
||||
import library.sdxl_train_util
|
||||
library.utils.setup_logging = setup_logging
|
||||
|
||||
library.train_util.load_tokenizer = load_tokenizers
|
||||
library.sdxl_train_util.load_tokenizers = load_sdxl_tokenizers
|
||||
|
||||
global original_load_target_model
|
||||
if original_load_target_model is None:
|
||||
original_load_target_model = library.train_util._load_target_model
|
||||
|
||||
library.train_util._load_target_model = _load_target_model
|
||||
|
||||
library.sdxl_train_util._load_target_model = _sdxl_load_target_model
|
||||
|
||||
global original_conditional_loss
|
||||
if original_conditional_loss is None:
|
||||
original_conditional_loss = library.train_util.conditional_loss
|
||||
library.train_util.conditional_loss = conditional_loss
|
||||
|
||||
|
||||
def _sdxl_load_target_model(
|
||||
name_or_path: str, vae_path: Optional[str], model_version: str, weight_dtype, device="cpu", model_dtype=None, *args, **kwargs
|
||||
):
|
||||
import library.sdxl_model_util as sdxl_model_util
|
||||
import library.model_util as model_util
|
||||
import library.sdxl_original_unet as sdxl_original_unet
|
||||
import library.sdxl_train_util
|
||||
init_empty_weights = library.sdxl_train_util.init_empty_weights
|
||||
|
||||
# model_dtype only work with full fp16/bf16
|
||||
name_or_path = os.readlink(name_or_path) if os.path.islink(
|
||||
name_or_path) else name_or_path
|
||||
load_stable_diffusion_format = False
|
||||
|
||||
if True:
|
||||
# Diffusers model is loaded to CPU
|
||||
variant = "fp16" if weight_dtype == torch.float16 else None
|
||||
print(
|
||||
f"load Diffusers pretrained models: {name_or_path}, variant={variant}")
|
||||
try:
|
||||
try:
|
||||
pipe = StableDiffusionXLPipeline.from_single_file(
|
||||
name_or_path, local_files_only=True, safety_checker=None)
|
||||
except EnvironmentError as ex:
|
||||
raise ex
|
||||
except EnvironmentError as ex:
|
||||
print(
|
||||
f"model is not found as a file or in Hugging Face, perhaps file name is wrong? / 指定したモデル名のファイル、またはHugging Faceのモデルが見つかりません。ファイル名が誤っているかもしれません: {name_or_path}"
|
||||
)
|
||||
raise ex
|
||||
|
||||
text_encoder1 = pipe.text_encoder
|
||||
text_encoder2 = pipe.text_encoder_2
|
||||
|
||||
# convert to fp32 for cache text_encoders outputs
|
||||
if text_encoder1.dtype != torch.float32:
|
||||
text_encoder1 = text_encoder1.to(dtype=torch.float32)
|
||||
if text_encoder2.dtype != torch.float32:
|
||||
text_encoder2 = text_encoder2.to(dtype=torch.float32)
|
||||
|
||||
vae = pipe.vae
|
||||
unet = pipe.unet
|
||||
global clip_large_tokenizer, clip_big_tokenizer
|
||||
clip_large_tokenizer = pipe.tokenizer
|
||||
clip_big_tokenizer = pipe.tokenizer_2
|
||||
del pipe
|
||||
|
||||
# Diffusers U-Net to original U-Net
|
||||
state_dict = sdxl_model_util.convert_diffusers_unet_state_dict_to_sdxl(
|
||||
unet.state_dict())
|
||||
with init_empty_weights():
|
||||
unet = sdxl_original_unet.SdxlUNet2DConditionModel() # overwrite unet
|
||||
sdxl_model_util._load_state_dict_on_device(
|
||||
unet, state_dict, device=device, dtype=model_dtype)
|
||||
print("U-Net converted to original U-Net")
|
||||
|
||||
logit_scale = None
|
||||
ckpt_info = None
|
||||
|
||||
# VAEを読み込む
|
||||
if vae_path is not None:
|
||||
vae = model_util.load_vae(vae_path, weight_dtype)
|
||||
print("additional VAE loaded")
|
||||
|
||||
return load_stable_diffusion_format, text_encoder1, text_encoder2, vae, unet, logit_scale, ckpt_info
|
||||
|
||||
|
||||
def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", unet_use_linear_projection_in_v2=False):
|
||||
import library.model_util as model_util
|
||||
from library.original_unet import UNet2DConditionModel
|
||||
|
||||
name_or_path = args.pretrained_model_name_or_path
|
||||
name_or_path = os.path.realpath(name_or_path) if os.path.islink(
|
||||
name_or_path) else name_or_path
|
||||
load_stable_diffusion_format = False
|
||||
if True:
|
||||
# Diffusers model is loaded to CPU
|
||||
try:
|
||||
pipe = StableDiffusionPipeline.from_single_file(
|
||||
name_or_path, local_files_only=True, safety_checker=None)
|
||||
except EnvironmentError as ex:
|
||||
print(
|
||||
f"model is not found as a file or in Hugging Face, perhaps file name is wrong? / 指定したモデル名のファイル、またはHugging Faceのモデルが見つかりません。ファイル名が誤っているかもしれません: {name_or_path}"
|
||||
)
|
||||
raise ex
|
||||
|
||||
text_encoder = pipe.text_encoder
|
||||
vae = pipe.vae
|
||||
unet = pipe.unet
|
||||
global clip_large_tokenizer
|
||||
clip_large_tokenizer = pipe.tokenizer
|
||||
del pipe
|
||||
|
||||
# Diffusers U-Net to original U-Net
|
||||
# TODO *.ckpt/*.safetensorsのv2と同じ形式にここで変換すると良さそう
|
||||
# print(f"unet config: {unet.config}")
|
||||
original_unet = UNet2DConditionModel(
|
||||
unet.config.sample_size,
|
||||
unet.config.attention_head_dim,
|
||||
unet.config.cross_attention_dim,
|
||||
unet.config.use_linear_projection,
|
||||
unet.config.upcast_attention,
|
||||
)
|
||||
original_unet.load_state_dict(unet.state_dict())
|
||||
unet = original_unet
|
||||
unet = unet.train()
|
||||
print("U-Net converted to original U-Net")
|
||||
|
||||
# VAEを読み込む
|
||||
if args.vae is not None:
|
||||
vae = model_util.load_vae(args.vae, weight_dtype)
|
||||
print("additional VAE loaded")
|
||||
|
||||
return text_encoder, vae, unet, load_stable_diffusion_format
|
||||
@@ -0,0 +1,341 @@
|
||||
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import json
|
||||
import subprocess
|
||||
|
||||
import torch
|
||||
|
||||
from .mz_train_tools_utils import Utils
|
||||
import folder_paths
|
||||
import nodes
|
||||
|
||||
|
||||
# 初始化工具仓库和工作区
|
||||
def MZ_KohyaSSInitWorkspace_call(args={}):
|
||||
mz_dir = Utils.get_minus_zone_models_path()
|
||||
git_url = "https://github.com/kohya-ss/sd-scripts"
|
||||
kohya_ss_lora_dir = os.path.join(mz_dir, "train_tools", "kohya_ss_lora")
|
||||
try:
|
||||
if not os.path.exists(kohya_ss_lora_dir) or not os.path.exists(os.path.join(kohya_ss_lora_dir, ".git")):
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", git_url, kohya_ss_lora_dir], check=True)
|
||||
|
||||
# 切换远程分支 git remote set-branches origin 'main'
|
||||
branch = args.get("branch", "main")
|
||||
subprocess.run(
|
||||
["git", "remote", "set-branches", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
subprocess.run(
|
||||
["git", "fetch", "--depth", "1", "origin", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
# 恢复所有文件
|
||||
subprocess.run(
|
||||
["git", "checkout", "."], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
subprocess.run(
|
||||
["git", "checkout", branch], cwd=kohya_ss_lora_dir, check=True)
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"克隆kohya-ss/sd-scripts或者切换分支时出现异常,详细信息请查看控制台...")
|
||||
|
||||
workspace_name = args.get("workspace_name", None)
|
||||
|
||||
if workspace_name is None or workspace_name == "":
|
||||
raise Exception("工作区名称不能为空(workspace_name is required)")
|
||||
|
||||
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)
|
||||
|
||||
return (
|
||||
args,
|
||||
)
|
||||
|
||||
|
||||
def MZ_ImageSelecter_call(args={}):
|
||||
images = args.get("images")
|
||||
pil_images = Utils.tensors2pil_list(images)
|
||||
|
||||
resolution = args.get("resolution", 512)
|
||||
|
||||
workspace_config = args.get("workspace_config", {})
|
||||
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)
|
||||
train_images_dir = os.path.join(workspace_dir, "train_images")
|
||||
os.makedirs(train_images_dir, exist_ok=True)
|
||||
|
||||
for i, pil_image in enumerate(pil_images):
|
||||
pil_image = Utils.resize_max(pil_image, resolution, resolution)
|
||||
filename = hashlib.md5(pil_image.tobytes()).hexdigest() + ".png"
|
||||
pil_image.save(os.path.join(train_images_dir, filename))
|
||||
|
||||
dataset_config_path = os.path.join(workspace_dir, "dataset.toml")
|
||||
generate_toml_config(
|
||||
dataset_config_path,
|
||||
enable_bucket=args.get("enable_bucket") == "enable",
|
||||
resolution=args.get("resolution"),
|
||||
batch_size=args.get("batch_size"),
|
||||
image_dir=train_images_dir,
|
||||
caption_extension=".caption",
|
||||
num_repeats=args.get("num_repeats"),
|
||||
)
|
||||
return (
|
||||
train_images_dir,
|
||||
)
|
||||
|
||||
|
||||
def MZ_KohyaSSUseConfig_call(args={}):
|
||||
# raise Exception(f"MZ_KohyaSSUseConfig_call: {args}")
|
||||
workspace_config = args.get("workspace_config", {})
|
||||
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)
|
||||
|
||||
if not os.path.exists(workspace_dir):
|
||||
raise Exception(f"工作区不存在: {workspace_dir}")
|
||||
|
||||
workspace_config_file = os.path.join(workspace_dir, "config.json")
|
||||
if not os.path.exists(workspace_config_file):
|
||||
train_config_template_dir = args.get("train_config_template_dir", None)
|
||||
train_config_template = args.get("train_config_template", None)
|
||||
|
||||
train_config_template_file = os.path.join(
|
||||
train_config_template_dir, train_config_template + ".json")
|
||||
|
||||
shutil.copy(train_config_template_file, workspace_config_file)
|
||||
|
||||
config = None
|
||||
with open(workspace_config_file, "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
ckpt_name = args.get("ckpt_name", "")
|
||||
if ckpt_name == "":
|
||||
raise Exception("未选择模型文件(ckpt_name is required)")
|
||||
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
|
||||
|
||||
# output_name
|
||||
config["train_config"]["output_name"] = workspace_name
|
||||
|
||||
dataset_config_path = os.path.join(
|
||||
workspace_dir, "dataset.toml")
|
||||
config["train_config"]["dataset_config"] = dataset_config_path
|
||||
|
||||
if config is None:
|
||||
raise Exception(f"读取配置文件失败: {workspace_config_file}")
|
||||
|
||||
with open(workspace_config_file, "w", encoding="utf-8") as f:
|
||||
json.dump(config, f, indent=4, ensure_ascii=False)
|
||||
|
||||
# raise Exception(f"MZ_KohyaSSUseConfig_call: {args}")
|
||||
return (
|
||||
args,
|
||||
)
|
||||
|
||||
|
||||
def config2args(train_parser: argparse.ArgumentParser, config):
|
||||
try:
|
||||
config_args_list = []
|
||||
for key, value in config.items():
|
||||
if type(value) == bool:
|
||||
if value:
|
||||
config_args_list.append(f"--{key}")
|
||||
else:
|
||||
config_args_list.append(f"--{key}")
|
||||
config_args_list.append(str(value))
|
||||
args = train_parser.parse_args(config_args_list)
|
||||
return args
|
||||
except Exception as e:
|
||||
raise Exception(f"config2args: {e}")
|
||||
|
||||
|
||||
def check_install():
|
||||
try:
|
||||
import toml
|
||||
except ImportError:
|
||||
os.system(f"{sys.executable} -m pip install toml")
|
||||
|
||||
# imagesize
|
||||
try:
|
||||
import imagesize
|
||||
except ImportError:
|
||||
os.system(f"{sys.executable} -m pip install imagesize")
|
||||
|
||||
# voluptuous
|
||||
try:
|
||||
import voluptuous
|
||||
except ImportError:
|
||||
os.system(f"{sys.executable} -m pip install voluptuous")
|
||||
|
||||
|
||||
import logging
|
||||
|
||||
|
||||
def generate_toml_config(output_path, enable_bucket=True, resolution=512, batch_size=1, image_dir=None, caption_extension=".caption", num_repeats=10, ):
|
||||
import toml
|
||||
config = {
|
||||
'general': {
|
||||
'enable_bucket': enable_bucket,
|
||||
},
|
||||
'datasets': [
|
||||
{
|
||||
'resolution': resolution,
|
||||
'batch_size': batch_size,
|
||||
'subsets': [
|
||||
{
|
||||
'image_dir': image_dir,
|
||||
'caption_extension': caption_extension,
|
||||
'num_repeats': num_repeats,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
toml.dump(config, f)
|
||||
|
||||
|
||||
def latent2image(previews, latent_file):
|
||||
samples = nodes.LoadLatent().load(latent_file)[0]
|
||||
with torch.no_grad():
|
||||
|
||||
samples = samples["samples"].to("cuda")
|
||||
preview_bytes = previews.decode_latent_to_preview_image(
|
||||
"", samples)
|
||||
return preview_bytes[1]
|
||||
|
||||
|
||||
def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func):
|
||||
|
||||
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")
|
||||
is_running = True
|
||||
|
||||
taesd_type = "sd1_5"
|
||||
if trainer_func.find("sd1_5") != -1:
|
||||
taesd_type = "sd1_5"
|
||||
if trainer_func.find("sdxl") != -1:
|
||||
taesd_type = "sdxl"
|
||||
pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type)
|
||||
|
||||
previewer = pb.get_previewer()
|
||||
import traceback
|
||||
|
||||
def log_callback(log):
|
||||
try:
|
||||
resp = json.loads(log)
|
||||
if resp.get("type") == "sample_images":
|
||||
global_step = resp.get("global_step")
|
||||
latent_path = resp.get("latent")
|
||||
preview_bytes = None
|
||||
if latent_path is not None:
|
||||
preview_bytes = latent2image(previewer, latent_path)
|
||||
pb.update(
|
||||
int(global_step), int(max_train_steps), preview_bytes)
|
||||
else:
|
||||
print(f"LOG: {log}")
|
||||
except Exception as e:
|
||||
print(f"LOG: {log} e: {e} ")
|
||||
print(f"stack: {traceback.format_exc()}")
|
||||
return is_running
|
||||
server, port = Utils.UDP_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, "--udp_port", str(port)],
|
||||
check=True,
|
||||
)
|
||||
is_running = False
|
||||
server.close()
|
||||
except Exception as e:
|
||||
is_running = False
|
||||
server.close()
|
||||
raise Exception(f"训练失败!!! 具体报错信息请查看控制台 {e}")
|
||||
|
||||
|
||||
def MZ_KohyaSSTrain_call(args={}):
|
||||
|
||||
# pb = Utils.progress_bar(1)
|
||||
# previewer = pb.get_previewer()
|
||||
|
||||
# img = latent2image(previewer, "/tmp/sample_images_4.latent")
|
||||
|
||||
# pb.update(1, 1, img)
|
||||
# raise Exception(f"MZ_KohyaSSTrain_call: {args}")
|
||||
|
||||
train_config = args.get("train_config", {})
|
||||
workspace_config = train_config.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)
|
||||
|
||||
workspace_config_file = os.path.join(workspace_dir, "config.json")
|
||||
|
||||
if not os.path.exists(workspace_config_file):
|
||||
raise Exception(f"配置文件不存在: {workspace_config_file}")
|
||||
|
||||
config = None
|
||||
with open(workspace_config_file, "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
if config is None:
|
||||
raise Exception(f"读取配置文件失败: {workspace_config_file}")
|
||||
|
||||
kohya_ss_tool_dir = os.path.join(
|
||||
Utils.get_minus_zone_models_path(), "train_tools", "kohya_ss_lora")
|
||||
|
||||
if kohya_ss_tool_dir not in sys.path:
|
||||
sys.path.append(kohya_ss_tool_dir)
|
||||
check_install()
|
||||
|
||||
train_type = config.get("metadata").get("train_type")
|
||||
# 在ComfyUI中无法运行, 梯度有问题, 原因不清楚
|
||||
# import train_network
|
||||
# train_network.logger = logging.getLogger()
|
||||
# trainer = train_network.NetworkTrainer()
|
||||
# train_args = config2args(train_network.setup_parser(), train_config)
|
||||
# trainer.train(train_args)
|
||||
# return
|
||||
|
||||
if train_type == "lora_sd1_5":
|
||||
run_hook_kohya_ss_run_file(kohya_ss_tool_dir, config.get(
|
||||
"train_config"), "run_lora_sd1_5")
|
||||
elif train_type == "lora_sdxl":
|
||||
run_hook_kohya_ss_run_file(kohya_ss_tool_dir, config.get(
|
||||
"train_config"), "run_lora_sdxl")
|
||||
elif train_type == "lora_sd3":
|
||||
run_hook_kohya_ss_run_file(kohya_ss_tool_dir, config.get(
|
||||
"train_config"), "run_lora_sd3")
|
||||
else:
|
||||
raise Exception(
|
||||
f"暂时不支持的训练类型: {train_type}")
|
||||
|
||||
return (
|
||||
"训练完成",
|
||||
)
|
||||
@@ -0,0 +1,785 @@
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from typing import Tuple
|
||||
import warnings
|
||||
import numpy as np
|
||||
import folder_paths
|
||||
import base64
|
||||
from PIL import Image, ImageFilter
|
||||
import io
|
||||
import torch
|
||||
import re
|
||||
import hashlib
|
||||
import cv2
|
||||
# sys.path.append(os.path.join(os.path.dirname(__file__)))
|
||||
temp_directory = folder_paths.get_temp_directory()
|
||||
from tqdm import tqdm
|
||||
import requests
|
||||
import comfy.utils
|
||||
|
||||
|
||||
CACHE_POOL = {}
|
||||
|
||||
|
||||
class Utils:
|
||||
def Md5(str):
|
||||
return hashlib.md5(str.encode('utf-8')).hexdigest()
|
||||
|
||||
def check_frames_path(frames_path):
|
||||
|
||||
if frames_path == "" or frames_path.startswith(".") or frames_path.startswith("/") or frames_path.endswith("/") or frames_path.endswith("\\"):
|
||||
return "frames_path不能为空"
|
||||
|
||||
frames_path = os.path.join(
|
||||
folder_paths.get_output_directory(), frames_path)
|
||||
|
||||
if frames_path == folder_paths.get_output_directory():
|
||||
return "frames_path不能为output目录"
|
||||
|
||||
return ""
|
||||
|
||||
def base64_to_pil_image(base64_str):
|
||||
if base64_str is None:
|
||||
return None
|
||||
if len(base64_str) == 0:
|
||||
return None
|
||||
if type(base64_str) not in [str, bytes]:
|
||||
return None
|
||||
if base64_str.startswith("data:image/png;base64,"):
|
||||
base64_str = base64_str.split(",")[-1]
|
||||
base64_str = base64_str.encode("utf-8")
|
||||
base64_str = base64.b64decode(base64_str)
|
||||
return Image.open(io.BytesIO(base64_str))
|
||||
|
||||
def pil_image_to_base64(pil_image):
|
||||
buffered = io.BytesIO()
|
||||
pil_image.save(buffered, format="PNG")
|
||||
img_str = base64.b64encode(buffered.getvalue())
|
||||
img_str = str(img_str, encoding="utf-8")
|
||||
return f"data:image/png;base64,{img_str}"
|
||||
|
||||
def listdir_png(path):
|
||||
try:
|
||||
files = os.listdir(path)
|
||||
new_files = []
|
||||
for file in files:
|
||||
if file.endswith(".png"):
|
||||
new_files.append(file)
|
||||
files = new_files
|
||||
files.sort(key=lambda x: int(os.path.basename(x).split(".")[0]))
|
||||
return files
|
||||
except Exception as e:
|
||||
return []
|
||||
|
||||
def tensor2pil(image):
|
||||
return Image.fromarray(np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
||||
|
||||
def tensors2pil_list(images):
|
||||
return [Utils.tensor2pil(image) for image in images]
|
||||
|
||||
# Convert PIL to Tensor
|
||||
|
||||
def pil2tensor(image):
|
||||
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)[0]
|
||||
|
||||
def pil2cv(image):
|
||||
return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
||||
|
||||
def cv2pil(image):
|
||||
return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
def list_tensor2tensor(data):
|
||||
result_tensor = torch.stack(data)
|
||||
return result_tensor
|
||||
|
||||
def loadImage(path):
|
||||
img = Image.open(path)
|
||||
img = img.convert("RGB")
|
||||
return img
|
||||
|
||||
def vae_encode_crop_pixels(pixels):
|
||||
x = (pixels.shape[1] // 8) * 8
|
||||
y = (pixels.shape[2] // 8) * 8
|
||||
if pixels.shape[1] != x or pixels.shape[2] != y:
|
||||
x_offset = (pixels.shape[1] % 8) // 2
|
||||
y_offset = (pixels.shape[2] % 8) // 2
|
||||
pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
|
||||
return pixels
|
||||
|
||||
def native_vae_encode(vae, image):
|
||||
pixels = Utils.vae_encode_crop_pixels(image)
|
||||
t = vae.encode(pixels[:, :, :, :3])
|
||||
return {"samples": t}
|
||||
|
||||
def native_vae_encode_for_inpaint(vae, pixels, mask):
|
||||
x = (pixels.shape[1] // 8) * 8
|
||||
y = (pixels.shape[2] // 8) * 8
|
||||
mask = torch.nn.functional.interpolate(mask.reshape(
|
||||
(-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
|
||||
|
||||
pixels = pixels.clone()
|
||||
if pixels.shape[1] != x or pixels.shape[2] != y:
|
||||
x_offset = (pixels.shape[1] % 8) // 2
|
||||
y_offset = (pixels.shape[2] % 8) // 2
|
||||
pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
|
||||
mask = mask[:, :, x_offset:x + x_offset, y_offset:y + y_offset]
|
||||
|
||||
# grow mask by a few pixels to keep things seamless in latent space
|
||||
|
||||
mask_erosion = mask
|
||||
|
||||
m = (1.0 - mask.round()).squeeze(1)
|
||||
for i in range(3):
|
||||
pixels[:, :, :, i] -= 0.5
|
||||
pixels[:, :, :, i] *= m
|
||||
pixels[:, :, :, i] += 0.5
|
||||
t = vae.encode(pixels)
|
||||
|
||||
return {"samples": t, "noise_mask": (mask_erosion[:, :, :x, :y].round())}
|
||||
|
||||
def native_vae_decode(vae, samples):
|
||||
return vae.decode(samples["samples"])
|
||||
|
||||
def native_clip_text_encode(clip, text):
|
||||
tokens = clip.tokenize(text)
|
||||
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
||||
return [[cond, {"pooled_output": pooled}]]
|
||||
|
||||
def cache_get(key):
|
||||
return CACHE_POOL.get(key, None)
|
||||
|
||||
def cache_set(key, value):
|
||||
global CACHE_POOL
|
||||
CACHE_POOL[key] = value
|
||||
return True
|
||||
|
||||
def get_minus_zone_models_path():
|
||||
models_path = os.path.join(
|
||||
folder_paths.models_dir, "minus_zone_models")
|
||||
os.makedirs(models_path, exist_ok=True)
|
||||
return models_path
|
||||
|
||||
def get_comfyui_models_path():
|
||||
return folder_paths.models_dir
|
||||
|
||||
def translate_text(text, from_code, to_code):
|
||||
try:
|
||||
import argostranslate
|
||||
from argostranslate import translate
|
||||
except ImportError:
|
||||
subprocess.run([
|
||||
sys.executable, "-m",
|
||||
"pip", "install", "argostranslate"], check=True)
|
||||
import argostranslate
|
||||
from argostranslate import translate
|
||||
|
||||
try:
|
||||
translation = translate.get_translation_from_codes(
|
||||
from_code, to_code)
|
||||
if translation is None:
|
||||
raise Exception("Translation not found")
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
argostranslate.package.update_package_index()
|
||||
available_packages = argostranslate.package.get_available_packages()
|
||||
package_to_install = next(
|
||||
filter(
|
||||
lambda x: (x.from_code == from_code and x.to_code ==
|
||||
to_code), available_packages,
|
||||
)
|
||||
)
|
||||
download_path = package_to_install.download()
|
||||
print("package_to_install.download():", download_path)
|
||||
argostranslate.package.install_from_path(download_path)
|
||||
|
||||
translation = translate.get_translation_from_codes(
|
||||
from_code, to_code)
|
||||
if translation is None:
|
||||
return text
|
||||
|
||||
# Translate
|
||||
translatedText = translation.translate(
|
||||
text)
|
||||
|
||||
return translatedText
|
||||
|
||||
def zh2en(text):
|
||||
return Utils.translate_text(text, "zh", "en")
|
||||
|
||||
def en2zh(text):
|
||||
return Utils.translate_text(text, "en", "zh")
|
||||
|
||||
def prompt_zh_to_en(prompt):
|
||||
prompt = prompt.replace("，", ",")
|
||||
prompt = prompt.replace("。", ",")
|
||||
prompt = prompt.replace("\n", ",")
|
||||
tags = prompt.split(",")
|
||||
# 判断是否有中文
|
||||
for i, tag in enumerate(tags):
|
||||
if re.search(u'[\u4e00-\u9fff]', tag):
|
||||
tags[i] = Utils.zh2en(tag)
|
||||
# 如果第一个字母是大写,转为小写
|
||||
if tags[i][0].isupper():
|
||||
tags[i] = tags[i].lower().replace(".", "")
|
||||
|
||||
return ",".join(tags)
|
||||
|
||||
def mask_resize(mask, width, height):
|
||||
mask = mask.unsqueeze(0).unsqueeze(0)
|
||||
mask = torch.nn.functional.interpolate(
|
||||
mask, size=(height, width), mode="bilinear")
|
||||
mask = mask.squeeze(0).squeeze(0)
|
||||
return mask
|
||||
|
||||
def mask_threshold(interested_mask):
|
||||
mask_image = Utils.tensor2pil(interested_mask)
|
||||
mask_image_cv2 = Utils.pil2cv(mask_image)
|
||||
ret, thresh1 = cv2.threshold(
|
||||
mask_image_cv2, 127, 255, cv2.THRESH_BINARY)
|
||||
thresh1 = Utils.cv2pil(thresh1)
|
||||
thresh1 = np.array(thresh1)
|
||||
thresh1 = thresh1[:, :, 0]
|
||||
return Utils.pil2tensor(thresh1)
|
||||
|
||||
def mask_erode(interested_mask, value):
|
||||
value = int(value)
|
||||
mask_image = Utils.tensor2pil(interested_mask)
|
||||
mask_image_cv2 = Utils.pil2cv(mask_image)
|
||||
kernel = np.ones((5, 5), np.uint8)
|
||||
erosion = cv2.erode(mask_image_cv2, kernel, iterations=value)
|
||||
erosion = Utils.cv2pil(erosion)
|
||||
erosion = np.array(erosion)
|
||||
erosion = erosion[:, :, 0]
|
||||
return Utils.pil2tensor(erosion)
|
||||
|
||||
def mask_dilate(interested_mask, value):
|
||||
value = int(value)
|
||||
mask_image = Utils.tensor2pil(interested_mask)
|
||||
mask_image_cv2 = Utils.pil2cv(mask_image)
|
||||
kernel = np.ones((5, 5), np.uint8)
|
||||
dilation = cv2.dilate(mask_image_cv2, kernel, iterations=value)
|
||||
dilation = Utils.cv2pil(dilation)
|
||||
dilation = np.array(dilation)
|
||||
dilation = dilation[:, :, 0]
|
||||
return Utils.pil2tensor(dilation)
|
||||
|
||||
def mask_edge_opt(interested_mask, edge_feathering):
|
||||
|
||||
mask_image = Utils.tensor2pil(interested_mask)
|
||||
mask_image_cv2 = Utils.pil2cv(mask_image)
|
||||
|
||||
# 高斯模糊
|
||||
dilation2 = Utils.cv2pil(mask_image_cv2)
|
||||
dilation2 = mask_image.filter(
|
||||
ImageFilter.GaussianBlur(edge_feathering))
|
||||
|
||||
# mask_image dilation2 图片蒙版叠加
|
||||
dilation2 = Utils.pil2cv(dilation2)
|
||||
# dilation2[mask_image_cv2 < 127] = 0
|
||||
dilation2 = Utils.cv2pil(dilation2)
|
||||
# to RGB
|
||||
dilation2 = np.array(dilation2)
|
||||
dilation2 = dilation2[:, :, 0]
|
||||
return Utils.pil2tensor(dilation2)
|
||||
|
||||
def mask_composite(destination, source, x, y, mask=None, multiplier=8, resize_source=False):
|
||||
source = source.to(destination.device)
|
||||
if resize_source:
|
||||
source = torch.nn.functional.interpolate(source, size=(
|
||||
destination.shape[2], destination.shape[3]), mode="bilinear")
|
||||
|
||||
source = comfy.utils.repeat_to_batch_size(source, destination.shape[0])
|
||||
|
||||
x = max(-source.shape[3] * multiplier,
|
||||
min(x, destination.shape[3] * multiplier))
|
||||
y = max(-source.shape[2] * multiplier,
|
||||
min(y, destination.shape[2] * multiplier))
|
||||
|
||||
left, top = (x // multiplier, y // multiplier)
|
||||
right, bottom = (left + source.shape[3], top + source.shape[2],)
|
||||
|
||||
if mask is None:
|
||||
mask = torch.ones_like(source)
|
||||
else:
|
||||
mask = mask.to(destination.device, copy=True)
|
||||
mask = torch.nn.functional.interpolate(mask.reshape(
|
||||
(-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
|
||||
mask = comfy.utils.repeat_to_batch_size(mask, source.shape[0])
|
||||
|
||||
# calculate the bounds of the source that will be overlapping the destination
|
||||
# this prevents the source trying to overwrite latent pixels that are out of bounds
|
||||
# of the destination
|
||||
visible_width, visible_height = (
|
||||
destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)
|
||||
|
||||
mask = mask[:, :, :visible_height, :visible_width]
|
||||
inverse_mask = torch.ones_like(mask) - mask
|
||||
|
||||
source_portion = mask * source[:, :, :visible_height, :visible_width]
|
||||
destination_portion = inverse_mask * \
|
||||
destination[:, :, top:bottom, left:right]
|
||||
|
||||
destination[:, :, top:bottom,
|
||||
left:right] = source_portion + destination_portion
|
||||
return destination
|
||||
|
||||
def latent_upscale_by(samples, scale_by):
|
||||
s = samples.copy()
|
||||
width = round(samples["samples"].shape[3] * scale_by)
|
||||
height = round(samples["samples"].shape[2] * scale_by)
|
||||
s["samples"] = comfy.utils.common_upscale(
|
||||
samples["samples"], width, height, "nearest-exact", "disabled")
|
||||
return s
|
||||
|
||||
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 add_watermark(image, watermark):
|
||||
if watermark == "":
|
||||
return image
|
||||
|
||||
try:
|
||||
import PIL
|
||||
from PIL import ImageDraw, ImageFont
|
||||
except ImportError:
|
||||
subprocess.run([
|
||||
sys.executable, "-m",
|
||||
"pip", "install", "Pillow"], check=True)
|
||||
import PIL
|
||||
from PIL import ImageDraw, ImageFont
|
||||
|
||||
# 获取PIL版本号
|
||||
pil_version = PIL.__version__
|
||||
|
||||
if pil_version >= "10.0.0":
|
||||
def textsize(self, text, font):
|
||||
left, top, right, bottom = self.textbbox((0, 0), text, font)
|
||||
return right - left, bottom - top
|
||||
ImageDraw.ImageDraw.textsize = textsize
|
||||
|
||||
font_fullpath = Utils.download_model(
|
||||
{
|
||||
"url": "https://www.modelscope.cn/api/v1/models/wailovet/MinusZoneAIModels/repo?Revision=master&FilePath=font%2FAlibabaPuHuiTi-2-75-SemiBold.ttf",
|
||||
"output": "font/AlibabaPuHuiTi-2-75-SemiBold.ttf",
|
||||
}
|
||||
)
|
||||
|
||||
watermarks = watermark.split("\n")
|
||||
|
||||
width, height = image.size
|
||||
short_edge = min(width, height)
|
||||
font_size = short_edge // 12
|
||||
|
||||
font = ImageFont.truetype(font_fullpath, font_size)
|
||||
|
||||
# print("pil_version:", pil_version)
|
||||
|
||||
draw = ImageDraw.Draw(image)
|
||||
|
||||
text = watermarks[0]
|
||||
textwidth, textheight = draw.textsize(text, font)
|
||||
|
||||
x = (width - textwidth) // 2
|
||||
|
||||
bottom = 10
|
||||
|
||||
y = height - textheight - (textheight * 0.4 + bottom + 8)
|
||||
draw.text((x, y), text, font=font)
|
||||
|
||||
if len(watermarks) > 1:
|
||||
y1 = y + textheight
|
||||
text = watermarks[1]
|
||||
font_size = int(font_size * 0.4)
|
||||
font = ImageFont.truetype(font_fullpath, font_size)
|
||||
textwidth, textheight = draw.textsize(text, font)
|
||||
x = (width - textwidth) // 2
|
||||
y = y1 - bottom + 4
|
||||
draw.text((x, y), text, font=font)
|
||||
|
||||
return image
|
||||
|
||||
def get_device():
|
||||
return comfy.model_management.get_torch_device()
|
||||
|
||||
def download_file(url, filepath, threads=8, retries=6):
|
||||
|
||||
get_size_tmp = requests.get(url, stream=True)
|
||||
total_size = int(get_size_tmp.headers.get("content-length", 0))
|
||||
|
||||
print(f"Downloading {url} to {filepath} with size {total_size} bytes")
|
||||
|
||||
base_filename = os.path.basename(filepath)
|
||||
cache_dir = os.path.join(os.path.dirname(
|
||||
filepath), f"{base_filename}.t_{threads}_cache")
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
|
||||
def get_total_existing_size():
|
||||
fs = os.listdir(cache_dir)
|
||||
existing_size = 0
|
||||
for f in fs:
|
||||
if f.startswith("block_"):
|
||||
existing_size += os.path.getsize(
|
||||
os.path.join(cache_dir, f))
|
||||
return existing_size
|
||||
|
||||
total_existing_size = get_total_existing_size()
|
||||
|
||||
if total_size != 0 and total_existing_size != total_size:
|
||||
|
||||
with tqdm(total=total_size, initial=total_existing_size, unit="B", unit_scale=True) as progress_bar:
|
||||
all_threads = []
|
||||
|
||||
for i in range(threads):
|
||||
cache_filepath = os.path.join(cache_dir, f"block_{i}")
|
||||
|
||||
start = total_size // threads * i
|
||||
end = total_size // threads * (i + 1) - 1
|
||||
|
||||
if i == threads - 1:
|
||||
end = total_size
|
||||
|
||||
# Check if the file already exists
|
||||
if os.path.exists(cache_filepath):
|
||||
# Get the size of the existing file
|
||||
existing_size = os.path.getsize(cache_filepath)
|
||||
else:
|
||||
existing_size = 0
|
||||
|
||||
headers = {"Range": f"bytes={start + existing_size}-{end}"}
|
||||
if end == total_size:
|
||||
headers = {"Range": f"bytes={start + existing_size}-"}
|
||||
if start + existing_size >= end:
|
||||
continue
|
||||
# print(f"Downloading {cache_filepath} with headers bytes={start + existing_size}-{end}")
|
||||
|
||||
# Streaming, so we can iterate over the response.
|
||||
response = requests.get(url, stream=True, headers=headers)
|
||||
|
||||
def download_file_thread(response, cache_filepath):
|
||||
block_size = 1024
|
||||
if end - (start + existing_size) < block_size:
|
||||
block_size = end - (start + existing_size)
|
||||
with open(cache_filepath, "ab") as file:
|
||||
for data in response.iter_content(block_size):
|
||||
file.write(data)
|
||||
progress_bar.update(
|
||||
len(data)
|
||||
)
|
||||
|
||||
t = threading.Thread(
|
||||
target=download_file_thread, args=(response, cache_filepath))
|
||||
|
||||
all_threads.append(t)
|
||||
|
||||
t.start()
|
||||
|
||||
for t in all_threads:
|
||||
t.join()
|
||||
|
||||
if total_size != 0 and get_total_existing_size() > total_size:
|
||||
# 文件下载失败
|
||||
shutil.rmtree(cache_dir)
|
||||
raise RuntimeError("Download failed, file is incomplete")
|
||||
|
||||
if total_size != 0 and total_size != get_total_existing_size():
|
||||
if retries > 0:
|
||||
retries -= 1
|
||||
print(
|
||||
f"Download failed: {total_size} != {get_total_existing_size()}, retrying... {retries} retries left")
|
||||
return Utils.download_file(url, filepath, threads, retries)
|
||||
|
||||
# 文件损坏
|
||||
raise RuntimeError(
|
||||
f"Download failed: {total_size} != {get_total_existing_size()}")
|
||||
|
||||
if os.path.exists(filepath):
|
||||
shutil.move(filepath, filepath + ".old." +
|
||||
time.strftime("%Y%m%d%H%M%S"))
|
||||
|
||||
# merge the files
|
||||
with open(filepath, "wb") as f:
|
||||
for i in range(threads):
|
||||
cache_filepath = os.path.join(cache_dir, f"block_{i}")
|
||||
with open(cache_filepath, "rb") as cf:
|
||||
f.write(cf.read())
|
||||
|
||||
shutil.rmtree(cache_dir)
|
||||
return filepath
|
||||
|
||||
def hf_download_model(url, only_get_path=False):
|
||||
if not url.startswith("https://"):
|
||||
raise ValueError("URL must start with https://")
|
||||
if url.startswith("https://huggingface.co/") or url.startswith("https://hf-mirror.com/"):
|
||||
base_model_path = os.path.abspath(os.path.join(
|
||||
Utils.get_models_path(), "transformers_models"))
|
||||
# https://huggingface.co/FaradayDotDev/llama-3-8b-Instruct-GGUF/resolve/main/llama-3-8b-Instruct.Q2_K.gguf?download=true
|
||||
texts = url.split("?")[0].split("/")
|
||||
file_name = texts[-1]
|
||||
zone_path = f"{texts[3]}/{texts[4]}"
|
||||
|
||||
save_path = os.path.join(base_model_path, zone_path, file_name)
|
||||
|
||||
if os.path.exists(save_path) is False:
|
||||
if only_get_path:
|
||||
return None
|
||||
os.makedirs(os.path.join(
|
||||
base_model_path, zone_path), exist_ok=True)
|
||||
Utils.download_file(url, save_path)
|
||||
|
||||
Utils.print_log(
|
||||
f"File {save_path} => {os.path.getsize(save_path)} ")
|
||||
|
||||
# 获取大小
|
||||
if os.path.getsize(save_path) == 0:
|
||||
if only_get_path:
|
||||
return None
|
||||
os.remove(save_path)
|
||||
raise ValueError(f"Download failed: {url}")
|
||||
return save_path
|
||||
else:
|
||||
texts = url.split("?")[0].split("/")
|
||||
host = texts[2].replace(".", "_")
|
||||
base_model_path = os.path.abspath(os.path.join(
|
||||
Utils.get_models_path(), f"{host}_models"))
|
||||
|
||||
file_name = texts[-1]
|
||||
file_name_no_ext = os.path.splitext(file_name)[0]
|
||||
file_ext = os.path.splitext(file_name)[1]
|
||||
md5_hash = Utils.Md5(url)
|
||||
|
||||
save_path = os.path.join(
|
||||
base_model_path, f"{file_name_no_ext}.{md5_hash}{file_ext}")
|
||||
|
||||
if os.path.exists(save_path) is False:
|
||||
if only_get_path:
|
||||
return None
|
||||
os.makedirs(base_model_path, exist_ok=True)
|
||||
Utils.download_file(url, save_path)
|
||||
|
||||
return save_path
|
||||
|
||||
def print_log(*args):
|
||||
if os.environ.get("MZ_DEV", None) is not None:
|
||||
print(*args)
|
||||
|
||||
def download_model(model_info, only_get_path=False):
|
||||
|
||||
url = model_info["url"]
|
||||
output = model_info["output"]
|
||||
save_path = os.path.abspath(
|
||||
os.path.join(Utils.get_comfyui_models_path(), output))
|
||||
if not os.path.exists(save_path):
|
||||
if only_get_path:
|
||||
return None
|
||||
save_path = Utils.download_file(url, save_path)
|
||||
return save_path
|
||||
|
||||
def load_lora(model, lora_path, strength_model):
|
||||
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
model_lora, _ = comfy.sd.load_lora_for_models(
|
||||
model, None, lora, strength_model, 0)
|
||||
return model_lora
|
||||
|
||||
def load_checkpoint(ckpt_name):
|
||||
cache_data = Utils.cache_get(ckpt_name)
|
||||
if cache_data is not None:
|
||||
print("load from cache: ", ckpt_name)
|
||||
return cache_data
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
out = comfy.sd.load_checkpoint_guess_config(
|
||||
ckpt_path, output_vae=True, output_clip=True)
|
||||
model, clip, vae = out[:3]
|
||||
Utils.cache_set(ckpt_name, (model, clip, vae))
|
||||
return model, clip, vae
|
||||
|
||||
def file_hash(file_path, hash_method):
|
||||
if not os.path.isfile(file_path):
|
||||
return ''
|
||||
h = hash_method()
|
||||
with open(file_path, 'rb') as f:
|
||||
while b := f.read(8192):
|
||||
h.update(b)
|
||||
return h.hexdigest()
|
||||
|
||||
def file_sha256(file_path):
|
||||
return Utils.file_hash(file_path, hashlib.sha256)
|
||||
|
||||
def get_auto_model_fullpath(model_name):
|
||||
find_paths = []
|
||||
target_sha256 = ""
|
||||
file_path = ""
|
||||
download_url = ""
|
||||
for model in MODEL_ZOO:
|
||||
if model["model"] == model_name:
|
||||
find_paths = model["find_path"]
|
||||
target_sha256 = model["SHA256"]
|
||||
file_path = model["file_path"]
|
||||
download_url = model["url"]
|
||||
break
|
||||
|
||||
if target_sha256 == "":
|
||||
raise ValueError(f"Model {model_name} not found in MODEL_ZOO")
|
||||
|
||||
if os.path.exists(file_path):
|
||||
if Utils.file_sha256(file_path) != target_sha256:
|
||||
print(f"Model {model_name} file hash not match...")
|
||||
return file_path
|
||||
|
||||
for find_path in find_paths:
|
||||
find_fullpath = os.path.join(
|
||||
folder_paths.get_output_directory(), find_path)
|
||||
|
||||
if os.path.exists(find_fullpath):
|
||||
for root, dirs, files in os.walk(find_fullpath):
|
||||
for file in files:
|
||||
if target_sha256 == Utils.file_sha256(os.path.join(root, file)):
|
||||
return os.path.join(root, file)
|
||||
|
||||
return Utils.download_model({"url": download_url, "output": file_path})
|
||||
|
||||
def is_sd15_model(model):
|
||||
type_str = str(type(model.model.model_config).__name__)
|
||||
return "SD15" in type_str
|
||||
|
||||
def is_sdxl_model(model):
|
||||
type_str = str(type(model.model.model_config).__name__)
|
||||
return "SDXL" in type_str
|
||||
|
||||
def progress_bar(steps,taesd_type="sd1_5"):
|
||||
class pb:
|
||||
if taesd_type == "sd1_5":
|
||||
taesd_decoder_name = "taesd_decoder"
|
||||
elif taesd_type == "sdxl":
|
||||
taesd_decoder_name = "taesdxl_decoder"
|
||||
latent_rgb_factors = None
|
||||
latent_channels = 4
|
||||
def __init__(self, steps):
|
||||
self.steps = steps
|
||||
self.pbar = comfy.utils.ProgressBar(steps)
|
||||
|
||||
def get_previewer(self):
|
||||
import latent_preview
|
||||
previewer = latent_preview.get_previewer(
|
||||
'cuda', self)
|
||||
return previewer
|
||||
|
||||
def update(self, step, total_steps, pil_img=None):
|
||||
print(f"step: {step}/{total_steps}")
|
||||
pil_img_info = ("JPEG", pil_img, 512)
|
||||
if pil_img is None:
|
||||
pil_img_info = None
|
||||
if type(pil_img) == Tuple:
|
||||
pil_img_info = pil_img
|
||||
try:
|
||||
self.pbar.update_absolute(
|
||||
step, total_steps, pil_img_info)
|
||||
except Exception as e:
|
||||
print("progress_bar:", e)
|
||||
|
||||
return pb(steps)
|
||||
|
||||
def UDP_Server(reader):
|
||||
import socket
|
||||
import threading
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
s.bind(("127.0.0.1", 0))
|
||||
port = s.getsockname()[1]
|
||||
|
||||
def recv():
|
||||
while True:
|
||||
data, addr = s.recvfrom(1024)
|
||||
if s.fileno() == -1:
|
||||
break
|
||||
try:
|
||||
data = str(data, encoding="utf-8")
|
||||
except:
|
||||
pass
|
||||
if reader(data) == False:
|
||||
break
|
||||
|
||||
threading.Thread(target=recv).start()
|
||||
|
||||
return s, port
|
||||
|
||||
|
||||
MODEL_ZOO = [
|
||||
{
|
||||
"model": "ip-adapter-plus_sd15",
|
||||
"find_path": [
|
||||
"ipadapter",
|
||||
],
|
||||
"file_path": "ipadapter/ip-adapter-plus_sd15.safetensors",
|
||||
"url": "https://www.modelscope.cn/api/v1/models/dannysun85/IP-Adapter/repo?Revision=master&FilePath=models%2Fip-adapter-plus_sd15.safetensors",
|
||||
"SHA256": "a1c250be40455cc61a43da1201ec3f1edaea71214865fb47f57927e06cbe4996",
|
||||
},
|
||||
{
|
||||
"model": "ip-adapter-plus_sdxl_vit-h",
|
||||
"find_path": [
|
||||
"ipadapter",
|
||||
],
|
||||
"file_path": "ipadapter/ip-adapter-plus_sdxl_vit-h.safetensors",
|
||||
"url": "https://www.modelscope.cn/api/v1/models/dannysun85/IP-Adapter/repo?Revision=master&FilePath=sdxl_models%2Fip-adapter-plus_sdxl_vit-h.safetensors",
|
||||
"SHA256": "a1c250be40455cc61a43da1201ec3f1edaea71214865fb47f57927e06cbe4996",
|
||||
},
|
||||
{
|
||||
"model": "CLIP-ViT-H-14-laion2B-s32B-b79K",
|
||||
"find_path": [
|
||||
"clip_vision",
|
||||
],
|
||||
"file_path": "clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors",
|
||||
"url": "https://www.modelscope.cn/api/v1/models/dannysun85/IP-Adapter/repo?Revision=master&FilePath=models%2Fimage_encoder%2Fmodel.safetensors",
|
||||
"SHA256": "6ca9667da1ca9e0b0f75e46bb030f7e011f44f86cbfb8d5a36590fcd7507b030",
|
||||
},
|
||||
{
|
||||
"model": "iclight_sd15_fc",
|
||||
"find_path": [
|
||||
"unet",
|
||||
],
|
||||
"file_path": "unet/iclight_sd15_fc.safetensors",
|
||||
"url": "https://www.modelscope.cn/api/v1/models/AI-ModelScope/ic-light/repo?Revision=master&FilePath=iclight_sd15_fc.safetensors",
|
||||
"SHA256": "a033fbaaa2f3f7859fa6a4477ee63ebbf9c116bf3569d5811856d2807f3468cd",
|
||||
},
|
||||
]
|
||||
Reference in New Issue
Block a user