415 lines
18 KiB
Python
415 lines
18 KiB
Python
import torch
|
|
import copy
|
|
import math
|
|
from typing import Any
|
|
import argparse
|
|
from .library import flux_models, flux_train_utils, flux_utils, sd3_train_utils, strategy_base, strategy_flux, train_util
|
|
from .train_network import NetworkTrainer, clean_memory_on_device
|
|
|
|
from accelerate import Accelerator
|
|
|
|
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class FluxNetworkTrainer(NetworkTrainer):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.sample_prompts_te_outputs = None
|
|
|
|
def assert_extra_args(self, args, train_dataset_group):
|
|
super().assert_extra_args(args, train_dataset_group)
|
|
|
|
if args.cache_text_encoder_outputs:
|
|
assert (
|
|
train_dataset_group.is_text_encoder_output_cacheable()
|
|
), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used"
|
|
|
|
assert (
|
|
args.network_train_unet_only or not args.cache_text_encoder_outputs
|
|
), "network for Text Encoder cannot be trained with caching Text Encoder outputs"
|
|
|
|
if args.max_token_length is not None:
|
|
logger.warning("max_token_length is not used in Flux training")
|
|
|
|
train_dataset_group.verify_bucket_reso_steps(32) # TODO check this
|
|
|
|
def get_flux_model_name(self, args):
|
|
return "schnell" if "schnell" in args.pretrained_model_name_or_path else "dev"
|
|
|
|
def load_target_model(self, args, weight_dtype, accelerator):
|
|
# currently offload to cpu for some models
|
|
name = self.get_flux_model_name(args)
|
|
# if we load to cpu, flux.to(fp8) takes a long time
|
|
model = flux_utils.load_flow_model(name, args.pretrained_model_name_or_path, weight_dtype, "cpu")
|
|
|
|
if args.split_mode:
|
|
model = self.prepare_split_model(model, weight_dtype, accelerator, args)
|
|
|
|
clip_l = flux_utils.load_clip_l(args.clip_l, weight_dtype, "cpu")
|
|
clip_l.eval()
|
|
|
|
# loading t5xxl to cpu takes a long time, so we should load to gpu in future
|
|
t5xxl = flux_utils.load_t5xxl(args.t5xxl, weight_dtype, "cpu")
|
|
t5xxl.eval()
|
|
|
|
ae = flux_utils.load_ae(name, args.ae, weight_dtype, "cpu")
|
|
|
|
return flux_utils.MODEL_VERSION_FLUX_V1, [clip_l, t5xxl], ae, model
|
|
|
|
def prepare_split_model(self, model, weight_dtype, accelerator, args):
|
|
from accelerate import init_empty_weights
|
|
|
|
logger.info("prepare split model")
|
|
with init_empty_weights():
|
|
flux_upper = flux_models.FluxUpper(model.params)
|
|
flux_lower = flux_models.FluxLower(model.params)
|
|
sd = model.state_dict()
|
|
|
|
# lower (trainable)
|
|
logger.info("load state dict for lower")
|
|
flux_lower.load_state_dict(sd, strict=False, assign=True)
|
|
flux_lower.to(dtype=weight_dtype)
|
|
|
|
# upper (frozen)
|
|
logger.info("load state dict for upper")
|
|
flux_upper.load_state_dict(sd, strict=False, assign=True)
|
|
|
|
logger.info("prepare upper model")
|
|
target_dtype = torch.float8_e4m3fn if args.fp8_base else weight_dtype
|
|
flux_upper.to(accelerator.device, dtype=target_dtype)
|
|
flux_upper.eval()
|
|
|
|
if args.fp8_base:
|
|
# this is required to run on fp8
|
|
flux_upper = accelerator.prepare(flux_upper)
|
|
|
|
flux_upper.to("cpu")
|
|
|
|
self.flux_upper = flux_upper
|
|
del model # we don't need model anymore
|
|
clean_memory_on_device(accelerator.device)
|
|
|
|
logger.info("split model prepared")
|
|
|
|
return flux_lower
|
|
|
|
def get_tokenize_strategy(self, args):
|
|
name = self.get_flux_model_name(args)
|
|
|
|
if args.t5xxl_max_token_length is None:
|
|
if name == "schnell":
|
|
t5xxl_max_token_length = 256
|
|
else:
|
|
t5xxl_max_token_length = 512
|
|
else:
|
|
t5xxl_max_token_length = args.t5xxl_max_token_length
|
|
|
|
logger.info(f"t5xxl_max_token_length: {t5xxl_max_token_length}")
|
|
return strategy_flux.FluxTokenizeStrategy(t5xxl_max_token_length, args.tokenizer_cache_dir)
|
|
|
|
def get_tokenizers(self, tokenize_strategy: strategy_flux.FluxTokenizeStrategy):
|
|
return [tokenize_strategy.clip_l, tokenize_strategy.t5xxl]
|
|
|
|
def get_latents_caching_strategy(self, args):
|
|
latents_caching_strategy = strategy_flux.FluxLatentsCachingStrategy(args.cache_latents_to_disk, args.vae_batch_size, False)
|
|
return latents_caching_strategy
|
|
|
|
def get_text_encoding_strategy(self, args):
|
|
return strategy_flux.FluxTextEncodingStrategy(apply_t5_attn_mask=args.apply_t5_attn_mask)
|
|
|
|
def get_models_for_text_encoding(self, args, accelerator, text_encoders):
|
|
return text_encoders # + [accelerator.unwrap_model(text_encoders[-1])]
|
|
|
|
def get_text_encoder_outputs_caching_strategy(self, args):
|
|
if args.cache_text_encoder_outputs:
|
|
return strategy_flux.FluxTextEncoderOutputsCachingStrategy(
|
|
args.cache_text_encoder_outputs_to_disk, None, False, apply_t5_attn_mask=args.apply_t5_attn_mask
|
|
)
|
|
else:
|
|
return None
|
|
|
|
def cache_text_encoder_outputs_if_needed(
|
|
self, args, accelerator: Accelerator, unet, vae, text_encoders, dataset: train_util.DatasetGroup, weight_dtype
|
|
):
|
|
if args.cache_text_encoder_outputs:
|
|
if not args.lowram:
|
|
# reduce memory consumption
|
|
logger.info("move vae and unet to cpu to save memory")
|
|
org_vae_device = vae.device
|
|
org_unet_device = unet.device
|
|
vae.to("cpu")
|
|
unet.to("cpu")
|
|
clean_memory_on_device(accelerator.device)
|
|
|
|
# When TE is not be trained, it will not be prepared so we need to use explicit autocast
|
|
logger.info("move text encoders to gpu")
|
|
text_encoders[0].to(accelerator.device, dtype=weight_dtype)
|
|
text_encoders[1].to(accelerator.device, dtype=weight_dtype)
|
|
with accelerator.autocast():
|
|
dataset.new_cache_text_encoder_outputs(text_encoders, accelerator.is_main_process)
|
|
|
|
# cache sample prompts
|
|
|
|
if args.sample_prompts is not None:
|
|
logger.info(f"cache Text Encoder outputs for sample prompt: {args.sample_prompts}")
|
|
|
|
tokenize_strategy: strategy_flux.FluxTokenizeStrategy = strategy_base.TokenizeStrategy.get_strategy()
|
|
text_encoding_strategy: strategy_flux.FluxTextEncodingStrategy = strategy_base.TextEncodingStrategy.get_strategy()
|
|
|
|
prompts = []
|
|
for line in args.sample_prompts:
|
|
line = line.strip()
|
|
if len(line) > 0 and line[0] != "#":
|
|
prompts.append(line)
|
|
|
|
# preprocess prompts
|
|
for i in range(len(prompts)):
|
|
prompt_dict = prompts[i]
|
|
if isinstance(prompt_dict, str):
|
|
from .library.train_util import line_to_prompt_dict
|
|
|
|
prompt_dict = line_to_prompt_dict(prompt_dict)
|
|
prompts[i] = prompt_dict
|
|
assert isinstance(prompt_dict, dict)
|
|
|
|
# Adds an enumerator to the dict based on prompt position. Used later to name image files. Also cleanup of extra data in original prompt dict.
|
|
prompt_dict["enum"] = i
|
|
prompt_dict.pop("subset", None)
|
|
|
|
sample_prompts_te_outputs = {} # key: prompt, value: text encoder outputs
|
|
with accelerator.autocast(), torch.no_grad():
|
|
for prompt_dict in prompts:
|
|
for p in [prompt_dict.get("prompt", ""), prompt_dict.get("negative_prompt", "")]:
|
|
if p not in sample_prompts_te_outputs:
|
|
logger.info(f"cache Text Encoder outputs for prompt: {p}")
|
|
tokens_and_masks = tokenize_strategy.tokenize(p)
|
|
sample_prompts_te_outputs[p] = text_encoding_strategy.encode_tokens(
|
|
tokenize_strategy, text_encoders, tokens_and_masks, args.apply_t5_attn_mask
|
|
)
|
|
self.sample_prompts_te_outputs = sample_prompts_te_outputs
|
|
accelerator.wait_for_everyone()
|
|
|
|
logger.info("move text encoders back to cpu")
|
|
text_encoders[0].to("cpu") # , dtype=torch.float32) # Text Encoder doesn't work with fp16 on CPU
|
|
text_encoders[1].to("cpu") # , dtype=torch.float32)
|
|
clean_memory_on_device(accelerator.device)
|
|
|
|
if not args.lowram:
|
|
logger.info("move vae and unet back to original device")
|
|
vae.to(org_vae_device)
|
|
unet.to(org_unet_device)
|
|
else:
|
|
# Text Encoder
|
|
text_encoders[0].to(accelerator.device, dtype=weight_dtype)
|
|
text_encoders[1].to(accelerator.device, dtype=weight_dtype)
|
|
|
|
def sample_images_split_mode(self, accelerator, args, epoch, global_step, flux, ae, text_encoder, sample_prompts_te_outputs, validation_settings):
|
|
|
|
class FluxUpperLowerWrapper(torch.nn.Module):
|
|
def __init__(self, flux_upper: flux_models.FluxUpper, flux_lower: flux_models.FluxLower, device: torch.device):
|
|
super().__init__()
|
|
self.flux_upper = flux_upper
|
|
self.flux_lower = flux_lower
|
|
self.target_device = device
|
|
|
|
def forward(self, img, img_ids, txt, txt_ids, timesteps, y, guidance=None):
|
|
self.flux_lower.to("cpu")
|
|
clean_memory_on_device(self.target_device)
|
|
self.flux_upper.to(self.target_device)
|
|
img, txt, vec, pe = self.flux_upper(img, img_ids, txt, txt_ids, timesteps, y, guidance)
|
|
self.flux_upper.to("cpu")
|
|
clean_memory_on_device(self.target_device)
|
|
self.flux_lower.to(self.target_device)
|
|
return self.flux_lower(img, txt, vec, pe)
|
|
|
|
wrapper = FluxUpperLowerWrapper(self.flux_upper, flux, accelerator.device)
|
|
clean_memory_on_device(accelerator.device)
|
|
flux_train_utils.sample_images(
|
|
accelerator, args, epoch, global_step, wrapper, ae, text_encoder, sample_prompts_te_outputs, validation_settings
|
|
)
|
|
clean_memory_on_device(accelerator.device)
|
|
|
|
def get_noise_scheduler(self, args: argparse.Namespace, device: torch.device) -> Any:
|
|
noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.discrete_flow_shift)
|
|
self.noise_scheduler_copy = copy.deepcopy(noise_scheduler)
|
|
return noise_scheduler
|
|
|
|
def is_text_encoder_not_needed_for_training(self, args):
|
|
return args.cache_text_encoder_outputs
|
|
|
|
def encode_images_to_latents(self, args, accelerator, vae, images):
|
|
return vae.encode(images)
|
|
|
|
def shift_scale_latents(self, args, latents):
|
|
return latents
|
|
|
|
def get_noise_pred_and_target(
|
|
self,
|
|
args,
|
|
accelerator,
|
|
noise_scheduler,
|
|
latents,
|
|
batch,
|
|
text_encoder_conds,
|
|
unet: flux_models.Flux,
|
|
network,
|
|
weight_dtype,
|
|
train_unet,
|
|
):
|
|
# copy from sd3_train.py and modified
|
|
|
|
def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
|
|
sigmas = self.noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
|
|
schedule_timesteps = self.noise_scheduler_copy.timesteps.to(accelerator.device)
|
|
timesteps = timesteps.to(accelerator.device)
|
|
step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
|
|
|
|
sigma = sigmas[step_indices].flatten()
|
|
while len(sigma.shape) < n_dim:
|
|
sigma = sigma.unsqueeze(-1)
|
|
return sigma
|
|
|
|
def compute_density_for_timestep_sampling(
|
|
weighting_scheme: str, batch_size: int, logit_mean: float = None, logit_std: float = None, mode_scale: float = None
|
|
):
|
|
"""Compute the density for sampling the timesteps when doing SD3 training.
|
|
|
|
Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
|
|
|
|
SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
|
|
"""
|
|
if weighting_scheme == "logit_normal":
|
|
# See 3.1 in the SD3 paper ($rf/lognorm(0.00,1.00)$).
|
|
u = torch.normal(mean=logit_mean, std=logit_std, size=(batch_size,), device="cpu")
|
|
u = torch.nn.functional.sigmoid(u)
|
|
elif weighting_scheme == "mode":
|
|
u = torch.rand(size=(batch_size,), device="cpu")
|
|
u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u)
|
|
else:
|
|
u = torch.rand(size=(batch_size,), device="cpu")
|
|
return u
|
|
|
|
def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None):
|
|
"""Computes loss weighting scheme for SD3 training.
|
|
|
|
Courtesy: This was contributed by Rafie Walker in https://github.com/huggingface/diffusers/pull/8528.
|
|
|
|
SD3 paper reference: https://arxiv.org/abs/2403.03206v1.
|
|
"""
|
|
if weighting_scheme == "sigma_sqrt":
|
|
weighting = (sigmas**-2.0).float()
|
|
elif weighting_scheme == "cosmap":
|
|
bot = 1 - 2 * sigmas + 2 * sigmas**2
|
|
weighting = 2 / (math.pi * bot)
|
|
else:
|
|
weighting = torch.ones_like(sigmas)
|
|
return weighting
|
|
|
|
# Sample noise that we'll add to the latents
|
|
noise = torch.randn_like(latents)
|
|
bsz = latents.shape[0]
|
|
|
|
# get noisy model input and timesteps
|
|
noisy_model_input, timesteps, sigmas = flux_train_utils.get_noisy_model_input_and_timesteps(
|
|
args, noise_scheduler, latents, noise, accelerator.device, weight_dtype
|
|
)
|
|
|
|
# pack latents and get img_ids
|
|
packed_noisy_model_input = flux_utils.pack_latents(noisy_model_input) # b, c, h*2, w*2 -> b, h*w, c*4
|
|
packed_latent_height, packed_latent_width = noisy_model_input.shape[2] // 2, noisy_model_input.shape[3] // 2
|
|
img_ids = flux_utils.prepare_img_ids(bsz, packed_latent_height, packed_latent_width).to(device=accelerator.device)
|
|
|
|
# get guidance
|
|
guidance_vec = torch.full((bsz,), args.guidance_scale, device=accelerator.device)
|
|
|
|
# ensure the hidden state will require grad
|
|
if args.gradient_checkpointing:
|
|
noisy_model_input.requires_grad_(True)
|
|
for t in text_encoder_conds:
|
|
t.requires_grad_(True)
|
|
img_ids.requires_grad_(True)
|
|
guidance_vec.requires_grad_(True)
|
|
|
|
# Predict the noise residual
|
|
l_pooled, t5_out, txt_ids = text_encoder_conds
|
|
# print(
|
|
# f"model_input: {noisy_model_input.shape}, img_ids: {img_ids.shape}, t5_out: {t5_out.shape}, txt_ids: {txt_ids.shape}, l_pooled: {l_pooled.shape}, timesteps: {timesteps.shape}, guidance_vec: {guidance_vec.shape}"
|
|
# )
|
|
|
|
if not args.split_mode:
|
|
# normal forward
|
|
with accelerator.autocast():
|
|
# YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transformer model (we should not keep it but I want to keep the inputs same for the model for testing)
|
|
model_pred = unet(
|
|
img=packed_noisy_model_input,
|
|
img_ids=img_ids,
|
|
txt=t5_out,
|
|
txt_ids=txt_ids,
|
|
y=l_pooled,
|
|
timesteps=timesteps / 1000,
|
|
guidance=guidance_vec,
|
|
)
|
|
else:
|
|
# split forward to reduce memory usage
|
|
assert network.train_blocks == "single", "train_blocks must be single for split mode"
|
|
with accelerator.autocast():
|
|
# move flux lower to cpu, and then move flux upper to gpu
|
|
unet.to("cpu")
|
|
clean_memory_on_device(accelerator.device)
|
|
self.flux_upper.to(accelerator.device)
|
|
|
|
# upper model does not require grad
|
|
with torch.no_grad():
|
|
intermediate_img, intermediate_txt, vec, pe = self.flux_upper(
|
|
img=packed_noisy_model_input,
|
|
img_ids=img_ids,
|
|
txt=t5_out,
|
|
txt_ids=txt_ids,
|
|
y=l_pooled,
|
|
timesteps=timesteps / 1000,
|
|
guidance=guidance_vec,
|
|
)
|
|
|
|
# move flux upper back to cpu, and then move flux lower to gpu
|
|
self.flux_upper.to("cpu")
|
|
clean_memory_on_device(accelerator.device)
|
|
unet.to(accelerator.device)
|
|
|
|
# lower model requires grad
|
|
intermediate_img.requires_grad_(True)
|
|
intermediate_txt.requires_grad_(True)
|
|
vec.requires_grad_(True)
|
|
pe.requires_grad_(True)
|
|
model_pred = unet(img=intermediate_img, txt=intermediate_txt, vec=vec, pe=pe)
|
|
|
|
# unpack latents
|
|
model_pred = flux_utils.unpack_latents(model_pred, packed_latent_height, packed_latent_width)
|
|
|
|
# apply model prediction type
|
|
model_pred, weighting = flux_train_utils.apply_model_prediction_type(args, model_pred, noisy_model_input, sigmas)
|
|
|
|
# flow matching loss: this is different from SD3
|
|
target = noise - latents
|
|
|
|
return model_pred, target, timesteps, None, weighting
|
|
|
|
def post_process_loss(self, loss, args, timesteps, noise_scheduler):
|
|
return loss
|
|
|
|
def get_sai_model_spec(self, args):
|
|
return train_util.get_sai_model_spec(None, args, False, True, False, flux="dev")
|
|
|
|
def update_metadata(self, metadata, args):
|
|
metadata["ss_apply_t5_attn_mask"] = args.apply_t5_attn_mask
|
|
metadata["ss_weighting_scheme"] = args.weighting_scheme
|
|
metadata["ss_logit_mean"] = args.logit_mean
|
|
metadata["ss_logit_std"] = args.logit_std
|
|
metadata["ss_mode_scale"] = args.mode_scale
|
|
metadata["ss_guidance_scale"] = args.guidance_scale
|
|
metadata["ss_timestep_sampling"] = args.timestep_sampling
|
|
metadata["ss_sigmoid_scale"] = args.sigmoid_scale
|
|
metadata["ss_model_prediction_type"] = args.model_prediction_type
|
|
metadata["ss_discrete_flow_shift"] = args.discrete_flow_shift
|