From 09154267c90ab411db57ec7065f10457d6b906be Mon Sep 17 00:00:00 2001 From: wailovet Date: Mon, 17 Jun 2024 23:30:23 +0800 Subject: [PATCH] 'update' --- .gitignore | 6 +- __init__.py | 141 ++++ configs/kohya_ss_lora/lora_sd1_5.json | 35 + configs/kohya_ss_lora/lora_sdxl.json | 35 + .../tokenizer_config.json | 34 + .../stable-diffusion-v1.5/v1-inference.yaml | 70 ++ .../stable-diffusion-xl/sd_xl_base.yaml | 93 +++ configs/sliders/config-xl.yaml | 28 + configs/sliders/prompts-xl.yaml | 9 + hook_kohya_ss_run.py | 129 +++ hook_kohya_ss_utils.py | 340 ++++++++ mz_train_tools_core.py | 341 ++++++++ mz_train_tools_utils.py | 785 ++++++++++++++++++ 13 files changed, 2045 insertions(+), 1 deletion(-) create mode 100644 __init__.py create mode 100644 configs/kohya_ss_lora/lora_sd1_5.json create mode 100644 configs/kohya_ss_lora/lora_sdxl.json create mode 100644 configs/models_config/clip-vit-large-patch14/tokenizer_config.json create mode 100644 configs/models_config/stable-diffusion-v1.5/v1-inference.yaml create mode 100644 configs/models_config/stable-diffusion-xl/sd_xl_base.yaml create mode 100644 configs/sliders/config-xl.yaml create mode 100644 configs/sliders/prompts-xl.yaml create mode 100644 hook_kohya_ss_run.py create mode 100644 hook_kohya_ss_utils.py create mode 100644 mz_train_tools_core.py create mode 100644 mz_train_tools_utils.py diff --git a/.gitignore b/.gitignore index 82f9275..b079f6c 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e56bd8a --- /dev/null +++ b/__init__.py @@ -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" diff --git a/configs/kohya_ss_lora/lora_sd1_5.json b/configs/kohya_ss_lora/lora_sd1_5.json new file mode 100644 index 0000000..bff6a15 --- /dev/null +++ b/configs/kohya_ss_lora/lora_sd1_5.json @@ -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 + } +} diff --git a/configs/kohya_ss_lora/lora_sdxl.json b/configs/kohya_ss_lora/lora_sdxl.json new file mode 100644 index 0000000..d655abf --- /dev/null +++ b/configs/kohya_ss_lora/lora_sdxl.json @@ -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 + } +} diff --git a/configs/models_config/clip-vit-large-patch14/tokenizer_config.json b/configs/models_config/clip-vit-large-patch14/tokenizer_config.json new file mode 100644 index 0000000..721c704 --- /dev/null +++ b/configs/models_config/clip-vit-large-patch14/tokenizer_config.json @@ -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" +} diff --git a/configs/models_config/stable-diffusion-v1.5/v1-inference.yaml b/configs/models_config/stable-diffusion-v1.5/v1-inference.yaml new file mode 100644 index 0000000..d4effe5 --- /dev/null +++ b/configs/models_config/stable-diffusion-v1.5/v1-inference.yaml @@ -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 diff --git a/configs/models_config/stable-diffusion-xl/sd_xl_base.yaml b/configs/models_config/stable-diffusion-xl/sd_xl_base.yaml new file mode 100644 index 0000000..6047379 --- /dev/null +++ b/configs/models_config/stable-diffusion-xl/sd_xl_base.yaml @@ -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 diff --git a/configs/sliders/config-xl.yaml b/configs/sliders/config-xl.yaml new file mode 100644 index 0000000..7a30879 --- /dev/null +++ b/configs/sliders/config-xl.yaml @@ -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 \ No newline at end of file diff --git a/configs/sliders/prompts-xl.yaml b/configs/sliders/prompts-xl.yaml new file mode 100644 index 0000000..2a6e5c7 --- /dev/null +++ b/configs/sliders/prompts-xl.yaml @@ -0,0 +1,9 @@ +- target: "" # 什么词可以消除积极的概念 + positive: "" # 要擦除的概念 + unconditional: "" # 与积极概念区别的词 + neutral: "" # 调节目标的起点 + action: "enhance" # 擦除或增强 + guidance_scale: 4 + resolution: 768 + dynamic_resolution: true + batch_size: 1 \ No newline at end of file diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py new file mode 100644 index 0000000..fc26155 --- /dev/null +++ b/hook_kohya_ss_run.py @@ -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]() diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py new file mode 100644 index 0000000..d515950 --- /dev/null +++ b/hook_kohya_ss_utils.py @@ -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 diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py new file mode 100644 index 0000000..6746a63 --- /dev/null +++ b/mz_train_tools_core.py @@ -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 ( + "训练完成", + ) diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py new file mode 100644 index 0000000..aed4ead --- /dev/null +++ b/mz_train_tools_utils.py @@ -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", + }, +]