diff --git a/__init__.py b/__init__.py index 930b314..9acde5a 100644 --- a/__init__.py +++ b/__init__.py @@ -539,7 +539,7 @@ class MZ_HYDiTTrain: "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}), "text_encoder_path": (["auto"] + folders, {"default": "auto"}), "tokenizer_path": (["auto"] + folders, {"default": "auto"}), - "t5_encoder_path": (["auto"] + folders, {"default": "auto"}), + "t5_encoder_path": (["none", "auto"] + folders, {"default": "auto"}), "resolution": ("INT", {"default": 1024}), "batch_size": ("INT", {"default": 1}), "epochs": ("INT", {"default": 50}), @@ -561,7 +561,7 @@ class MZ_HYDiTTrain: FUNCTION = "start" MZ_DESC = """ -base_size The base resolution (n, n) from which to create multiple resolutions | Recommended values: 256/512/1024 +base_size The base resolution (n, n) from which to create multiple resolutions | Recommended values: 256/512/1024 """ OUTPUT_NODE = True @@ -577,6 +577,66 @@ NODE_CLASS_MAPPINGS["MZ_HYDiTTrain"] = MZ_HYDiTTrain NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTTrain"] = f"{AUTHOR_NAME} - HYDiTTrain" +class MZ_HYDiTSimpleT2I: + @classmethod + def INPUT_TYPES(s): + hunyuan_models_path = os.path.join( + Utils.get_comfyui_models_path(), "hunyuan") + os.makedirs(hunyuan_models_path, exist_ok=True) + + models = Utils.get_models_by_folder(hunyuan_models_path) + folders = Utils.get_folders_by_folder(hunyuan_models_path) + + vae_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "vae")) + unet_models = Utils.get_models_by_folder( + os.path.join(Utils.get_comfyui_models_path(), "unet")) + + comfyui_full_loras = [] + comfyui_loras = folder_paths.get_filename_list("loras") + for lora in comfyui_loras: + lora_path = folder_paths.get_full_path("loras", lora) + comfyui_full_loras.append(lora_path) + + return { + "required": { + "unet_path": (["auto"] + models + unet_models, {"default": "auto"}), + "vae_ema_path": (["auto"] + folders + vae_models, {"default": "auto"}), + "text_encoder_path": (["auto"] + folders, {"default": "auto"}), + "tokenizer_path": (["auto"] + folders, {"default": "auto"}), + "t5_encoder_path": (["none", "auto"] + folders, {"default": "auto"}), + "lora_path": (["none"] + comfyui_full_loras, {"default": "none"}), + "seed": ("INT", {"default": 0}), + "steps": ("INT", {"default": 20}), + "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), + "scheduler": ([ + "ddpm", "ddim", "dpmms", + ], {"default": "ddpm"}), + "prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + "negative_prompt": ("STRING", {"default:": "", "dynamicPrompts": True, "multiline": True}), + "width": ("INT", {"default": 512, "max": 8192, "step": 16}), + "height": ("INT", {"default": 512, "max": 8192, "step": 16}), + "keep_device": (["enable", "disable"], {"default": "enable"}), + } + } + + RETURN_TYPES = ("IMAGE",) + + RETURN_NAMES = ("images",) + + FUNCTION = "start" + + CATEGORY = CATEGORY_NAME + + def start(self, **kwargs): + importlib.reload(mz_train_tools_core_HYDiT) + return mz_train_tools_core_HYDiT.MZ_HYDiTSimpleT2I_call(kwargs) + + +NODE_CLASS_MAPPINGS["MZ_HYDiTSimpleT2I"] = MZ_HYDiTSimpleT2I +NODE_DISPLAY_NAME_MAPPINGS["MZ_HYDiTSimpleT2I"] = f"{AUTHOR_NAME} - HYDiTSimpleT2I" + + class MZ_TrainToolsDebug: @classmethod def INPUT_TYPES(s): diff --git a/hook_HYDiT_main_train_deepspeed.py b/hook_HYDiT_main_train_deepspeed.py index e494616..521d3fc 100644 --- a/hook_HYDiT_main_train_deepspeed.py +++ b/hook_HYDiT_main_train_deepspeed.py @@ -287,9 +287,9 @@ def Core(args): ) try: - from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar + from .hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds except: - from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar + from hook_HYDiT_utils import VAE_EMA_PATH, TEXT_ENCODER, TOKENIZER, T5_ENCODER, easy_sample_images, model_resume, PBar, CustomizeEmbeds # Setup VAE logger.info(f" Loading vae from {VAE_EMA_PATH}") @@ -304,8 +304,11 @@ def Core(args): # Setup T5 text encoder from hydit.modules.text_encoder import MT5Embedder mt5_path = T5_ENCODER['MT5'] - embedder_t5 = MT5Embedder( - mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5) + if mt5_path is None: + embedder_t5 = CustomizeEmbeds() + else: + embedder_t5 = MT5Embedder( + mt5_path, torch_dtype=T5_ENCODER['torch_dtype'], max_length=args.text_len_t5) tokenizer_t5 = embedder_t5.tokenizer text_encoder_t5 = embedder_t5.model @@ -550,10 +553,10 @@ def Core(args): gc.collect() pbar.step( - f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {running_loss / step:.4f}", epoch * len(loader), train_steps) + f"Epoch {epoch}, step {step}, loss {loss.item():.4f}, mean_loss {running_loss / step:.4f}", args.epochs * len(loader), train_steps) if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0 # or train_steps == args.max_training_steps - ) and train_steps > 0: + ) and train_steps > 0: easy_sample_images(args, vae, text_encoder, tokenizer, model, embedder_t5, target_height=768, target_width=1280, train_steps=train_steps) save_checkpoint(args, rank, logger, model, ema, diff --git a/hook_HYDiT_run.py b/hook_HYDiT_run.py index 20f9422..e5300e3 100644 --- a/hook_HYDiT_run.py +++ b/hook_HYDiT_run.py @@ -49,7 +49,6 @@ if __name__ == "__main__": print("HYDi run hook") - try: from . import hook_HYDiT_main_train_deepspeed except Exception as e: @@ -108,14 +107,20 @@ if __name__ == "__main__": margs.index_file = train_config.get("index_file", None) margs.lr = train_config.get("lr", 1e-5) + + margs.rank = train_config.get("rank", 8) - margs.noise_offset = train_config.get("noise_offset", 0.1) + margs.noise_offset = train_config.get("noise_offset", 0.0) margs.log_every = train_config.get("log_every", 99999999999999) - margs.use_zero_stage = 2 + + margs.use_zero_stage = train_config.get("use_zero_stage", 2) + + margs.global_batch_size = train_config.get("global_batch_size", 1) margs.deepspeed = True margs.results_dir = train_config.get("results_dir") + margs.mse_loss_weight_type = train_config.get( "mse_loss_weight_type", "constant") diff --git a/hook_HYDiT_utils.py b/hook_HYDiT_utils.py index a5f1a70..eadbe30 100644 --- a/hook_HYDiT_utils.py +++ b/hook_HYDiT_utils.py @@ -2,6 +2,7 @@ import json import os +import time import torch @@ -38,7 +39,7 @@ def set_text_encoder_path(path): T5_ENCODER = { - 'MT5': 'ckpts/t2i/mt5', + 'MT5': None, 'attention_mask': True, 'layer_index': -1, 'attention_pool': True, @@ -82,6 +83,8 @@ def easy_sample_images( from hydit.constants import SAMPLER_FACTORY from hydit.modules.posemb_layers import get_fill_resize_and_crop, get_2d_rotary_pos_embed from hydit.modules.models import HUNYUAN_DIT_CONFIG + + import traceback with torch.no_grad(): workspace_dir = TRAIN_CONFIG.get("workspace_dir") @@ -130,6 +133,19 @@ def easy_sample_images( requires_safety_checker=False, embedder_t5=embedder_t5, ).to("cuda") + # attr _execution_device is not defined + if not hasattr(pipeline, "_execution_device"): + time.sleep(1) + pipeline = StableDiffusionPipeline(vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=model.module, + scheduler=scheduler, + feature_extractor=None, + safety_checker=None, + requires_safety_checker=False, + embedder_t5=embedder_t5, + ).to("cuda") style = torch.as_tensor([0, 0] * batch_size, device="cuda") @@ -152,6 +168,7 @@ def easy_sample_images( freqs_cis_img = calc_rope(height, width) try: + samples = pipeline( height=height, width=width, @@ -170,7 +187,9 @@ def easy_sample_images( pass except Exception as e: - print(f"Failed to sample images: {e}") + print(f"Failed to sample images: {e} ") + # 打印堆栈信息 + traceback.print_exc() continue # print("samples:",type(samples),) @@ -276,3 +295,52 @@ class PBar: "total_steps": total_steps, # "latent": noise_pred_latent_path, }) + + +from torch import nn + + +class CustomizeEmbedsModel(nn.Module): + dtype = torch.float16 + x = torch.zeros(1, 1, 256, 2048) + + def __init__(self, *args, **kwargs): + super().__init__() + + def to(self, *args, **kwargs): + self.dtype = torch.float16 + return self + + def forward(self, *args, **kwargs): + if self.output_hidden_states or kwargs.get("output_hidden_states", False): + return { + "hidden_states": self.x.to("cuda"), + "input_ids": torch.zeros(1, 1), + } + return self.x + + +class CustomizeTokenizer(dict): + + added_tokens_encoder = [] + input_ids = torch.zeros(1, 256) + attention_mask = torch.zeros(1, 256) + + def __init__(self): + self['added_tokens_encoder'] = self.added_tokens_encoder + self['input_ids'] = self.input_ids + self['attention_mask'] = self.attention_mask + + def tokenize(self, text): + return text + + def __call__(self, *args, **kwargs): + return self + + +class CustomizeEmbeds(): + def __init__(self): + super().__init__() + self.tokenizer = CustomizeTokenizer() + self.model = CustomizeEmbedsModel().to("cuda") + self.max_length = 256 diff --git a/hook_kohya_ss_utils.py b/hook_kohya_ss_utils.py index f347a1f..e1f6031 100644 --- a/hook_kohya_ss_utils.py +++ b/hook_kohya_ss_utils.py @@ -216,10 +216,6 @@ 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 @@ -354,7 +350,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, controlnet=None): +def generate_image(pipe_class, cmd_args, accelerator, vae, tokenizer, text_encoder, unet, epoch, prompt_dict_list): import library.train_util # for multi gpu distributed inference. this is a singleton, so it's safe to use it here @@ -373,10 +369,6 @@ 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, @@ -402,6 +394,7 @@ 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() @@ -413,9 +406,6 @@ 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 ) @@ -427,42 +417,3 @@ 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_HYDiT.py b/mz_train_tools_core_HYDiT.py index b304d98..dd34ae2 100644 --- a/mz_train_tools_core_HYDiT.py +++ b/mz_train_tools_core_HYDiT.py @@ -10,8 +10,10 @@ import json import subprocess import time +import safetensors.torch import torch + from .mz_train_tools_utils import Utils import folder_paths import nodes @@ -519,7 +521,7 @@ def MZ_HYDiTTrain_call(args={}): "vae_ema_path": args.get("vae_ema_path"), "text_encoder_path": args.get("text_encoder_path"), "tokenizer_path": args.get("tokenizer_path"), - "t5_encoder_path": args.get("t5_encoder_path"), + "t5_encoder_path": args.get("t5_encoder_path") if args.get("t5_encoder_path") != "none" else None, "results_dir": output_dir, "task_flag": output_name, }) @@ -641,3 +643,353 @@ def get_sample_images(train_config): ) return result return None + + +def get_model_from_path(model_path, lora_path, width, height): + + from hydit.modules.models import HunYuanDiT, HUNYUAN_DIT_CONFIG + + from types import SimpleNamespace + + # self.text_states_dim = args.text_states_dim + # self.text_states_dim_t5 = args.text_states_dim_t5 + # self.text_len = args.text_len + # self.text_len_t5 = args.text_len_t5 + # self.norm = args.norm + # use_flash_attn = args.infer_mode == 'fa' or args.use_flash_attn + + args = SimpleNamespace() + args.learn_sigma = True + args.text_states_dim = 1024 + args.text_states_dim_t5 = 2048 + args.text_len = 77 + args.text_len_t5 = 256 + args.norm = "layer" + args.infer_mode = "torch" + args.use_fp16 = True + args.qk_norm = True + args.use_flash_attn = False + + model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] + + latent_size = (width // 8, height // 8) + + model = HunYuanDiT(args, + input_size=latent_size, + **model_config, + ).half().to("cuda") + state_dict = torch.load( + model_path, map_location=lambda storage, loc: storage) + model.load_state_dict(state_dict) + + if lora_path is not None: + from peft import PeftConfig + + # 查询后缀 + if lora_path.lower().endswith(".safetensors"): + from safetensors.torch import load_file + adapter_state_dict = load_file(lora_path) + else: + adapter_state_dict = torch.load(lora_path) + + peft_config = PeftConfig() + + # 相同目录下同名的json文件 + lora_json_path = os.path.splitext(lora_path)[0] + ".json" + if not os.path.exists(lora_json_path): + raise Exception(f"未找到对应的json文件: {lora_json_path}") + with open(lora_json_path, "r", encoding="utf-8") as f: + loaded_attributes = json.load(f) + + for key, value in loaded_attributes.items(): + setattr(peft_config, key, value) + + model.load_adapter( + adapter_state_dict=adapter_state_dict, + peft_config=peft_config, + ) + model.merge_and_unload() + + return model + + +def clean_unet(): + unet_data = Utils.cache_get("HYDiT_UNET") + if unet_data is not None: + unet_data["model"].cpu() + Utils.cache_set("HYDiT_UNET", None) + torch.cuda.empty_cache() + + +def clean_vae(): + vae_data = Utils.cache_get("HYDiT_VAE_EMA") + if vae_data is not None: + vae_data["model"].cpu() + Utils.cache_set("HYDiT_VAE_EM", None) + torch.cuda.empty_cache() + + +def clean_clip_text_encoder(): + clip_text_encoder = Utils.cache_get("HYDiT_CLIP_TEXT_ENCODER") + if clip_text_encoder is not None: + clip_text_encoder["model"].cpu() + Utils.cache_set("HYDiT_CLIP_TEXT_ENCODER", None) + torch.cuda.empty_cache() + + +def clean_t5_encoder(): + t5_encoder = Utils.cache_get("HYDiT_T5_ENCODER") + if t5_encoder is not None: + t5_encoder["model"].cpu() + Utils.cache_set("HYDiT_T5_ENCODER", None) + torch.cuda.empty_cache() + + +from torch import nn + + +class CustomizeEmbedsModel(nn.Module): + dtype = torch.float16 + x = torch.zeros(1, 1, 256, 2048) + + def __init__(self, *args, **kwargs): + super().__init__() + + def to(self, *args, **kwargs): + self.dtype = torch.float16 + return self + + def forward(self, *args, **kwargs): + if self.output_hidden_states or kwargs.get("output_hidden_states", False): + return { + "hidden_states": self.x.to("cuda"), + "input_ids": torch.zeros(1, 1), + } + return self.x + + +class CustomizeTokenizer(dict): + + added_tokens_encoder = [] + input_ids = torch.zeros(1, 256) + attention_mask = torch.zeros(1, 256) + + def __init__(self): + self['added_tokens_encoder'] = self.added_tokens_encoder + self['input_ids'] = self.input_ids + self['attention_mask'] = self.attention_mask + + def tokenize(self, text): + return text + + def __call__(self, *args, **kwargs): + return self + + +class CustomizeEmbeds(): + def __init__(self): + super().__init__() + self.tokenizer = CustomizeTokenizer() + self.model = CustomizeEmbedsModel() + self.max_length = 256 + + +def MZ_HYDiTSimpleT2I_call(args={}): + check_required() + args = check_model_auto_download(args) + + HYDiT_tool_dir = os.path.join( + Utils.get_minus_zone_models_path(), "train_tools", "HunyuanDiT") + if HYDiT_tool_dir not in sys.path: + sys.path.append(HYDiT_tool_dir) + + import hydit.inference + from hydit.inference import SAMPLER_FACTORY, HUNYUAN_DIT_CONFIG, get_fill_resize_and_crop, get_2d_rotary_pos_embed + from hydit.modules.text_encoder import MT5Embedder + from hydit.utils.tools import set_seeds + from PIL import Image + from diffusers import schedulers, AutoencoderKL + from transformers import BertModel, BertTokenizer + + width = args.get("width") + height = args.get("height") + + unet_path = args.get("unet_path") + lora_path = args.get("lora_path") + unet_cache_key = f"HYDiT_UNET" + unet_data = Utils.cache_get(unet_cache_key) + if unet_data is None or unet_data.get("unet_path") != unet_path or unet_data.get("lora_path") != lora_path: + clean_unet() + _lora_path = lora_path if lora_path != "none" else None + unet = get_model_from_path( + args.get("unet_path"), _lora_path, width, height) + Utils.cache_set(unet_cache_key, { + "model": unet, + "unet_path": unet_path, + "lora_path": lora_path, + }) + else: + unet = unet_data["model"] + + text_encoder_path = args.get("text_encoder_path") + clip_text_encoder_data = Utils.cache_get("HYDiT_CLIP_TEXT_ENCODER") + + if clip_text_encoder_data is None or clip_text_encoder_data.get("text_encoder_path") != text_encoder_path: + clean_clip_text_encoder() + clip_text_encoder = BertModel.from_pretrained( + text_encoder_path, False, revision=None, local_files_only=True).to("cuda") + Utils.cache_set("HYDiT_CLIP_TEXT_ENCODER", { + "model": clip_text_encoder, + "text_encoder_path": text_encoder_path, + }) + else: + clip_text_encoder = clip_text_encoder_data["model"] + + tokenizer_path = args.get("tokenizer_path") + text_tokenizer = BertTokenizer.from_pretrained( + tokenizer_path, revision=None, local_files_only=True) + + vae_ema_path = args.get("vae_ema_path") + vae_ema_data = Utils.cache_get("HYDiT_VAE_EMA") + if vae_ema_data is None or vae_ema_data.get("vae_ema_path") != vae_ema_path: + clean_vae() + if os.path.isdir(vae_ema_path): + vae = AutoencoderKL.from_pretrained( + vae_ema_path, local_files_only=True).to("cuda") + else: + vae = AutoencoderKL.from_single_file( + vae_ema_path, local_files_only=True).to("cuda") + + Utils.cache_set("HYDiT_VAE_EMA", { + "model": vae, + "vae_ema_path": vae_ema_path, + }) + else: + vae = vae_ema_data["model"] + + t5_encoder_path = args.get("t5_encoder_path") + if t5_encoder_path == "none": + t5_encoder = CustomizeEmbeds() + else: + t5_encoder_data = Utils.cache_get("HYDiT_T5_ENCODER") + if t5_encoder_data is None or t5_encoder_data.get("t5_encoder_path") != t5_encoder_path: + clean_t5_encoder() + t5_encoder = MT5Embedder( + t5_encoder_path, torch_dtype=torch.float16, max_length=256) + Utils.cache_set("HYDiT_T5_ENCODER", { + "model": t5_encoder, + "t5_encoder_path": t5_encoder_path, + }) + else: + t5_encoder = t5_encoder_data["model"] + + sampler = args.get("scheduler") + scheduler = SAMPLER_FACTORY[sampler]['scheduler'] + kwargs = SAMPLER_FACTORY[sampler]['kwargs'] + # Build scheduler according to the sampler. + scheduler_class = getattr(schedulers, scheduler) + scheduler = scheduler_class(**kwargs) + + scheduler.set_timesteps(args.get("steps"), "cuda") + + from hydit.diffusion.pipeline import StableDiffusionPipeline + pipeline = StableDiffusionPipeline(vae=vae, + text_encoder=clip_text_encoder, + tokenizer=text_tokenizer, + unet=unet, + scheduler=scheduler, + feature_extractor=None, + safety_checker=None, + requires_safety_checker=False, + progress_bar_config=None, + embedder_t5=t5_encoder, + infer_mode="torch", + ) + + pipeline = pipeline.to("cuda") + + seed = args.get("seed", 0) + target_height = int((height // 16) * 16) + target_width = int((width // 16) * 16) + prompt = args.get("prompt", "").strip() + negative_prompt = args.get("negative_prompt", "").strip() + + batch_size = 1 + style = torch.as_tensor([0, 0] * batch_size, device="cuda") + + src_size_cond = (target_width, target_height) + size_cond = list(src_size_cond) + [target_width, target_height, 0, 0] + image_meta_size = torch.as_tensor( + [size_cond] * 2 * batch_size, device="cuda",) + + def calc_rope(height, width): + model_config = HUNYUAN_DIT_CONFIG["DiT-g/2"] + patch_size = model_config['patch_size'] + head_size = model_config['hidden_size'] // model_config['num_heads'] + th = height // 8 // patch_size + tw = width // 8 // patch_size + base_size = 512 // 8 // patch_size + start, stop = get_fill_resize_and_crop((th, tw), base_size) + sub_args = [start, stop, (th, tw)] + rope = get_2d_rotary_pos_embed(head_size, *sub_args) + return rope + + freqs_cis_img = calc_rope(target_height, target_width) + guidance_scale = args.get("cfg", 7.0) + infer_steps = args.get("steps", 20) + + with torch.inference_mode(): + with torch.cuda.amp.autocast(): + + pbar = Utils.progress_bar(0, "sdxl") + preview = pbar.get_previewer() + + def callback(step: int, timestep: int, latents: torch.FloatTensor, pred_x0): + try: + pil_img = preview.decode_latent_to_preview_image( + None, + latents, + )[1] + pbar.update(step, infer_steps, pil_img) + except Exception as e: + print("decode_tensors error:", e) + generator = set_seeds(seed, device="cuda") + samples = pipeline( + height=height, + width=width, + prompt=prompt, + negative_prompt=negative_prompt, + num_images_per_prompt=batch_size, + guidance_scale=guidance_scale, + num_inference_steps=infer_steps, + style=style, + return_dict=False, + generator=generator, + use_fp16=True, + learn_sigma=True, + freqs_cis_img=freqs_cis_img, + image_meta_size=image_meta_size, + callback=callback, + callback_steps=2, + )[0] + + if type(samples) == list: + pil_image = samples[0] + else: + pil_image = samples + + tensor_image = Utils.pil2tensor(pil_image) + + keep_device = args.get("keep_device", "disable") == "enable" + if not keep_device: + del pipeline + del unet + del clip_text_encoder + del vae + del t5_encoder + clean_unet() + clean_clip_text_encoder() + clean_vae() + torch.cuda.empty_cache() + + return (Utils.list_tensor2tensor([tensor_image]),) diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py index d6f2136..b21b147 100644 --- a/mz_train_tools_utils.py +++ b/mz_train_tools_utils.py @@ -710,12 +710,12 @@ class Utils: return previewer def update(self, step, total_steps, pil_img=None): - pil_img_info = ("JPEG", pil_img, 512) - if pil_img is None: - pil_img_info = None - if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple: - pil_img_info = pil_img try: + pil_img_info = ("JPEG", pil_img, 512) + if pil_img is None: + pil_img_info = None + if type(pil_img) == Tuple or type(pil_img) == list or type(pil_img) == tuple: + pil_img_info = pil_img # print("pil_img_info:", type(pil_img), pil_img_info) # print("step:", step, "total_steps:", total_steps) self.pbar.update_absolute(