File size: 8,764 Bytes
7bc5051 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 |
import argparse
import os
from typing import Optional, Union
import numpy as np
import torch
from tqdm import tqdm
from dataset import config_utils
from dataset.config_utils import BlueprintGenerator, ConfigSanitizer
import accelerate
from dataset.image_video_dataset import ARCHITECTURE_HUNYUAN_VIDEO, BaseDataset, ItemInfo, save_text_encoder_output_cache
from hunyuan_model import text_encoder as text_encoder_module
from hunyuan_model.text_encoder import TextEncoder
import logging
from utils.model_utils import str_to_dtype
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
def encode_prompt(text_encoder: TextEncoder, prompt: Union[str, list[str]]):
data_type = "video" # video only, image is not supported
text_inputs = text_encoder.text2tokens(prompt, data_type=data_type)
with torch.no_grad():
prompt_outputs = text_encoder.encode(text_inputs, data_type=data_type)
return prompt_outputs.hidden_state, prompt_outputs.attention_mask
def encode_and_save_batch(
text_encoder: TextEncoder, batch: list[ItemInfo], is_llm: bool, accelerator: Optional[accelerate.Accelerator]
):
prompts = [item.caption for item in batch]
# print(prompts)
# encode prompt
if accelerator is not None:
with accelerator.autocast():
prompt_embeds, prompt_mask = encode_prompt(text_encoder, prompts)
else:
prompt_embeds, prompt_mask = encode_prompt(text_encoder, prompts)
# # convert to fp16 if needed
# if prompt_embeds.dtype == torch.float32 and text_encoder.dtype != torch.float32:
# prompt_embeds = prompt_embeds.to(text_encoder.dtype)
# save prompt cache
for item, embed, mask in zip(batch, prompt_embeds, prompt_mask):
save_text_encoder_output_cache(item, embed, mask, is_llm)
def prepare_cache_files_and_paths(datasets: list[BaseDataset]):
all_cache_files_for_dataset = [] # exisiting cache files
all_cache_paths_for_dataset = [] # all cache paths in the dataset
for dataset in datasets:
all_cache_files = [os.path.normpath(file) for file in dataset.get_all_text_encoder_output_cache_files()]
all_cache_files = set(all_cache_files)
all_cache_files_for_dataset.append(all_cache_files)
all_cache_paths_for_dataset.append(set())
return all_cache_files_for_dataset, all_cache_paths_for_dataset
def process_text_encoder_batches(
num_workers: Optional[int],
skip_existing: bool,
batch_size: int,
datasets: list[BaseDataset],
all_cache_files_for_dataset: list[set],
all_cache_paths_for_dataset: list[set],
encode: callable,
):
num_workers = num_workers if num_workers is not None else max(1, os.cpu_count() - 1)
for i, dataset in enumerate(datasets):
logger.info(f"Encoding dataset [{i}]")
all_cache_files = all_cache_files_for_dataset[i]
all_cache_paths = all_cache_paths_for_dataset[i]
for batch in tqdm(dataset.retrieve_text_encoder_output_cache_batches(num_workers)):
# update cache files (it's ok if we update it multiple times)
all_cache_paths.update([os.path.normpath(item.text_encoder_output_cache_path) for item in batch])
# skip existing cache files
if skip_existing:
filtered_batch = [
item for item in batch if not os.path.normpath(item.text_encoder_output_cache_path) in all_cache_files
]
# print(f"Filtered {len(batch) - len(filtered_batch)} existing cache files")
if len(filtered_batch) == 0:
continue
batch = filtered_batch
bs = batch_size if batch_size is not None else len(batch)
for i in range(0, len(batch), bs):
encode(batch[i : i + bs])
def post_process_cache_files(
datasets: list[BaseDataset], all_cache_files_for_dataset: list[set], all_cache_paths_for_dataset: list[set], keep_cache: bool
):
for i, dataset in enumerate(datasets):
all_cache_files = all_cache_files_for_dataset[i]
all_cache_paths = all_cache_paths_for_dataset[i]
for cache_file in all_cache_files:
if cache_file not in all_cache_paths:
if keep_cache:
logger.info(f"Keep cache file not in the dataset: {cache_file}")
else:
os.remove(cache_file)
logger.info(f"Removed old cache file: {cache_file}")
def main(args):
device = args.device if args.device is not None else "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
# Load dataset config
blueprint_generator = BlueprintGenerator(ConfigSanitizer())
logger.info(f"Load dataset config from {args.dataset_config}")
user_config = config_utils.load_user_config(args.dataset_config)
blueprint = blueprint_generator.generate(user_config, args, architecture=ARCHITECTURE_HUNYUAN_VIDEO)
train_dataset_group = config_utils.generate_dataset_group_by_blueprint(blueprint.dataset_group)
datasets = train_dataset_group.datasets
# define accelerator for fp8 inference
accelerator = None
if args.fp8_llm:
accelerator = accelerate.Accelerator(mixed_precision="fp16")
# prepare cache files and paths: all_cache_files_for_dataset = exisiting cache files, all_cache_paths_for_dataset = all cache paths in the dataset
all_cache_files_for_dataset, all_cache_paths_for_dataset = prepare_cache_files_and_paths(datasets)
# Load Text Encoder 1
text_encoder_dtype = torch.float16 if args.text_encoder_dtype is None else str_to_dtype(args.text_encoder_dtype)
logger.info(f"loading text encoder 1: {args.text_encoder1}")
text_encoder_1 = text_encoder_module.load_text_encoder_1(args.text_encoder1, device, args.fp8_llm, text_encoder_dtype)
text_encoder_1.to(device=device)
# Encode with Text Encoder 1 (LLM)
logger.info("Encoding with Text Encoder 1")
def encode_for_text_encoder_1(batch: list[ItemInfo]):
encode_and_save_batch(text_encoder_1, batch, is_llm=True, accelerator=accelerator)
process_text_encoder_batches(
args.num_workers,
args.skip_existing,
args.batch_size,
datasets,
all_cache_files_for_dataset,
all_cache_paths_for_dataset,
encode_for_text_encoder_1,
)
del text_encoder_1
# Load Text Encoder 2
logger.info(f"loading text encoder 2: {args.text_encoder2}")
text_encoder_2 = text_encoder_module.load_text_encoder_2(args.text_encoder2, device, text_encoder_dtype)
text_encoder_2.to(device=device)
# Encode with Text Encoder 2
logger.info("Encoding with Text Encoder 2")
def encode_for_text_encoder_2(batch: list[ItemInfo]):
encode_and_save_batch(text_encoder_2, batch, is_llm=False, accelerator=None)
process_text_encoder_batches(
args.num_workers,
args.skip_existing,
args.batch_size,
datasets,
all_cache_files_for_dataset,
all_cache_paths_for_dataset,
encode_for_text_encoder_2,
)
del text_encoder_2
# remove cache files not in dataset
post_process_cache_files(datasets, all_cache_files_for_dataset, all_cache_paths_for_dataset, args.keep_cache)
def setup_parser_common():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset_config", type=str, required=True, help="path to dataset config .toml file")
parser.add_argument("--device", type=str, default=None, help="device to use, default is cuda if available")
parser.add_argument(
"--batch_size", type=int, default=None, help="batch size, override dataset config if dataset batch size > this"
)
parser.add_argument("--num_workers", type=int, default=None, help="number of workers for dataset. default is cpu count-1")
parser.add_argument("--skip_existing", action="store_true", help="skip existing cache files")
parser.add_argument("--keep_cache", action="store_true", help="keep cache files not in dataset")
return parser
def hv_setup_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
parser.add_argument("--text_encoder1", type=str, required=True, help="Text Encoder 1 directory")
parser.add_argument("--text_encoder2", type=str, required=True, help="Text Encoder 2 directory")
parser.add_argument("--text_encoder_dtype", type=str, default=None, help="data type for Text Encoder, default is float16")
parser.add_argument("--fp8_llm", action="store_true", help="use fp8 for Text Encoder 1 (LLM)")
return parser
if __name__ == "__main__":
parser = setup_parser_common()
parser = hv_setup_parser(parser)
args = parser.parse_args()
main(args)
|