'update'

__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):

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,

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")
 

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

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

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]),)

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(