diff --git a/__init__.py b/__init__.py index 471ff52..289db33 100644 --- a/__init__.py +++ b/__init__.py @@ -221,7 +221,7 @@ class AlwaysEqualProxy(str): return False -class MZ_KohyaSSTrain: +class MZ_KohyaSSLoraTrain: @classmethod def INPUT_TYPES(s): @@ -282,8 +282,61 @@ class MZ_KohyaSSTrain: 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" +NODE_CLASS_MAPPINGS["MZ_KohyaSSTrain"] = MZ_KohyaSSLoraTrain +NODE_DISPLAY_NAME_MAPPINGS[ + "MZ_KohyaSSTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(old version)" + +NODE_CLASS_MAPPINGS["MZ_KohyaSSLoraTrain"] = MZ_KohyaSSLoraTrain +NODE_DISPLAY_NAME_MAPPINGS[ + "MZ_KohyaSSLoraTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(lora)" + +class MZ_KohyaSSControlnetTrain: + + @classmethod + def INPUT_TYPES(s): + models = [ + "latest", + "empty", + ] + + comfyui_full_m_path = [] + comfyui_basemodels = folder_paths.get_filename_list("controlnet") + for b_model in comfyui_basemodels: + m_path = folder_paths.get_full_path("controlnet", b_model) + comfyui_full_m_path.append(m_path) + + # 按创建时间排序 + comfyui_full_m_path = sorted( + comfyui_full_m_path, key=lambda x: os.path.getctime(x), reverse=True) + + models = models + comfyui_full_m_path + + return { + "required": { + "train_config": ("MZ_TT_SS_TrainConfig",), + "base_controlnet": (models, {"default": "latest"}), + "sample_generate": (["enable", "disable"], {"default": "enable"}), + "sample_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + }, + "optional": { + "has_no_effect": (AlwaysEqualProxy("*"),), + }, + } + + RETURN_TYPES = () + RETURN_NAMES = () + 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_KohyaSSControlnetTrain"] = MZ_KohyaSSControlnetTrain +NODE_DISPLAY_NAME_MAPPINGS[ + "MZ_KohyaSSControlnetTrain"] = f"{AUTHOR_NAME} - KohyaSSTrain(controlnet)" class MZ_LoadImagesFromDirectoryPath: diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py index 0454550..73b821d 100644 --- a/hook_kohya_ss_run.py +++ b/hook_kohya_ss_run.py @@ -43,6 +43,10 @@ train_config_json = "{}" sample_images_pipe_class = None +def utils_sample_images(*args, **kwargs): + return sample_images(None, *args, **kwargs) + + def sample_images(self, *args, **kwargs): # accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet accelerator = args[0] @@ -58,27 +62,15 @@ def sample_images(self, *args, **kwargs): # print(f"sample_images: args = {args}") # print(f"sample_images: kwargs = {kwargs}") - # 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) + controlnet = kwargs.get("controlnet", None) if epoch is not None and cmd_args.save_every_n_epochs is not None and epoch % cmd_args.save_every_n_epochs == 0: prompt_dict_list = other_config.get("prompt_dict_list", []) if len(prompt_dict_list) == 0: seed = other_config.get("seed", 0) - - prompt_dict_list.append({ - "controlnet_image": None, + prompt_dict = { + "controlnet_image": other_config.get("controlnet_image", None), "prompt": other_config.get("sample_prompt", ""), "seed": seed, "negative_prompt": "", @@ -86,7 +78,9 @@ def sample_images(self, *args, **kwargs): "sample_sampler": "euler_a", "sample_steps": 20, "scale": 5.0, - }) + } + + prompt_dict_list.append(prompt_dict) else: for i, prompt_dict in enumerate(prompt_dict_list): if prompt_dict.get("controlnet_image", None) is None: @@ -98,6 +92,7 @@ def sample_images(self, *args, **kwargs): if prompt_dict.get("enum", None) is None: prompt_dict["enum"] = i hook_kohya_ss_utils.generate_image( + pipe_class=sample_images_pipe_class, cmd_args=cmd_args, accelerator=accelerator, epoch=epoch, @@ -105,8 +100,8 @@ def sample_images(self, *args, **kwargs): tokenizer=tokenizer, unet=unet, vae=vae, - pipe_class=sample_images_pipe_class, prompt_dict_list=prompt_dict_list, + controlnet=controlnet, ) LOG({ @@ -169,6 +164,7 @@ class SimpleNamespaceCNWarrper(SimpleNamespace): super().__init__(*args, **kwargs) self.__dict__.update(kwargs) # or self.__dict__ = kwargs self.__dict__["mid_block_type"] = "UNetMidBlock2DCrossAttn" + self.__dict__["_diffusers_version"] = "0.6.0" self.__iter__ = lambda: iter(kwargs.keys()) # is not iterable @@ -185,10 +181,13 @@ def run_controlnet_sd1_5(): types.SimpleNamespace = SimpleNamespaceCNWarrper hook_kohya_ss_utils.hook_kohya_ss() # 覆盖sample_images生成函数,包括进度条和生成图片功能 + import train_controlnet # 配置对应的pipeline import library.train_util + library.train_util.sample_images = utils_sample_images + global sample_images_pipe_class sample_images_pipe_class = library.train_util.StableDiffusionLongPromptWeightingPipeline diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py index e1f6031..f347a1f 100644 --- a/hook_kohya_ss_utils.py +++ b/hook_kohya_ss_utils.py @@ -216,6 +216,10 @@ def hook_kohya_ss(): import library.utils import library.train_util import library.sdxl_train_util + import library.model_util + + library.model_util.convert_controlnet_state_dict_to_diffusers = convert_controlnet_state_dict_to_diffusers + library.utils.setup_logging = setup_logging library.train_util.load_tokenizer = load_tokenizers @@ -350,7 +354,7 @@ def _load_target_model(args: argparse.Namespace, weight_dtype, device="cpu", une return text_encoder, vae, unet, load_stable_diffusion_format -def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list): +def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list, controlnet=None): import library.train_util # for multi gpu distributed inference. this is a singleton, so it's safe to use it here @@ -369,6 +373,10 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod v_parameterization=cmd_args.v_parameterization, ) + # print("unet.config._diffusers_version:", unet.config._diffusers_version) + # print("unet.config:", unet.config) + # input("press enter to continue...") + pipeline = pipe_class( text_encoder=text_encoder, vae=vae, @@ -394,7 +402,6 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod save_dir = sample_images_path prompt_replacement = None steps = 0 - controlnet = None # save random state to restore later rng_state = torch.get_rng_state() @@ -406,6 +413,9 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod with torch.no_grad(): for prompt_dict in prompt_dict_list: + print("prompt_dict:", json.dumps( + prompt_dict, indent=4, ensure_ascii=False)) + library.train_util.sample_image_inference( accelerator, cmd_args, pipeline, save_dir, prompt_dict, epoch, steps, prompt_replacement, controlnet=controlnet ) @@ -417,3 +427,42 @@ def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encod if cuda_rng_state is not None: torch.cuda.set_rng_state(cuda_rng_state) vae.to(org_vae_device) + + +def convert_controlnet_state_dict_to_diffusers(controlnet_state_dict): + import library.model_util as model_util + unet_conversion_map, unet_conversion_map_resnet, unet_conversion_map_layer = model_util.controlnet_conversion_map() + + mapping = {k: k for k in controlnet_state_dict.keys()} + + # print("controlnet_state_dict:", controlnet_state_dict.keys()) + keys_first = list(controlnet_state_dict.keys())[0] + + if keys_first.startswith("control_model."): + new_mapping = {} + for k, v in mapping.items(): + new_mapping[k.replace("control_model.", "")] = v.replace( + "control_model.", "") + mapping = new_mapping + + for sd_name, diffusers_name in unet_conversion_map: + mapping[sd_name] = diffusers_name + for k, v in mapping.items(): + for sd_part, diffusers_part in unet_conversion_map_layer: + v = v.replace(sd_part, diffusers_part) + mapping[k] = v + for k, v in mapping.items(): + if "resnets" in v: + for sd_part, diffusers_part in unet_conversion_map_resnet: + v = v.replace(sd_part, diffusers_part) + mapping[k] = v + + if keys_first.startswith("control_model."): + new_state_dict = { + v: controlnet_state_dict[ + f"control_model.{k}"] for k, v in mapping.items()} + else: + new_state_dict = { + v: controlnet_state_dict[k] for k, v in mapping.items()} + + return new_state_dict diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py index dd6d7c5..96b3cb5 100644 --- a/mz_train_tools_core.py +++ b/mz_train_tools_core.py @@ -511,14 +511,38 @@ def MZ_KohyaSSTrain_call(args={}): if "network_dropout" in train_config: del train_config["network_dropout"] + + + + base_controlnet = args.get("base_controlnet", "empty") + if base_controlnet == "empty": + pass + elif base_controlnet == "latest": + workspace_controlnet_dir = os.path.join(workspace_dir, "output") + if os.path.exists(workspace_controlnet_dir): + workspace_controlnet_files = os.listdir(workspace_controlnet_dir) + workspace_controlnet_files = list( + filter(lambda x: x.endswith(".safetensors"), workspace_controlnet_files)) + workspace_controlnet_files = list( + map(lambda x: os.path.join(workspace_controlnet_dir, x), workspace_controlnet_files)) + # 排序 + workspace_controlnet_files = sorted( + workspace_controlnet_files, key=lambda x: os.path.getctime(x), reverse=True) + if len(workspace_controlnet_files) > 0: + base_controlnet = os.path.join( + workspace_controlnet_dir, workspace_controlnet_files[0]) + else: + base_controlnet = "empty" + else: + pass + + if base_controlnet != "empty" and os.path.exists(base_controlnet): + train_config["controlnet_model_name_or_path"] = base_controlnet + + 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 + + sample_generate = args.get("sample_generate", "enable") sample_prompt = args.get("sample_prompt", "") if sample_generate == "enable": @@ -535,6 +559,15 @@ def MZ_KohyaSSTrain_call(args={}): run_hook_kohya_ss_run_file( kohya_ss_tool_dir, train_config, "run_lora_sdxl", other_config) elif train_type == "controlnet_sd1_5": + + conditioning_images_dir = os.path.join( + workspace_dir, "conditioning_images") + conditioning_images_onec = "" + if os.path.exists(conditioning_images_dir): + conditioning_images_onec = os.listdir(conditioning_images_dir)[0] + other_config["controlnet_image"] = os.path.join( + conditioning_images_dir, conditioning_images_onec) + run_hook_kohya_ss_run_file( kohya_ss_tool_dir, train_config, "run_controlnet_sd1_5", other_config) else: