From a9c5aa1f9336cedf1e294fd3c8c22bb649d51015 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 5 Jan 2025 22:28:51 +0900 Subject: [PATCH 01/40] add CFG to FLUX.1 sample image --- library/flux_train_utils.py | 152 ++++++++++++++++++++++++------------ 1 file changed, 104 insertions(+), 48 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index f7f06c5c..9f954f58 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -40,7 +40,7 @@ def sample_images( text_encoders, sample_prompts_te_outputs, prompt_replacement=None, - controlnet=None + controlnet=None, ): if steps == 0: if not args.sample_at_first: @@ -101,7 +101,7 @@ def sample_images( steps, sample_prompts_te_outputs, prompt_replacement, - controlnet + controlnet, ) else: # Creating list with N elements, where each element is a list of prompt_dicts, and N is the number of processes available (number of devices available) @@ -125,7 +125,7 @@ def sample_images( steps, sample_prompts_te_outputs, prompt_replacement, - controlnet + controlnet, ) torch.set_rng_state(rng_state) @@ -147,14 +147,14 @@ def sample_image_inference( steps, sample_prompts_te_outputs, prompt_replacement, - controlnet + controlnet, ): assert isinstance(prompt_dict, dict) - # negative_prompt = prompt_dict.get("negative_prompt") + negative_prompt = prompt_dict.get("negative_prompt") sample_steps = prompt_dict.get("sample_steps", 20) width = prompt_dict.get("width", 512) height = prompt_dict.get("height", 512) - scale = prompt_dict.get("scale", 3.5) + scale = prompt_dict.get("scale", 1.0) # 1.0 means no guidance seed = prompt_dict.get("seed") controlnet_image = prompt_dict.get("controlnet_image") prompt: str = prompt_dict.get("prompt", "") @@ -162,8 +162,8 @@ def sample_image_inference( if prompt_replacement is not None: prompt = prompt.replace(prompt_replacement[0], prompt_replacement[1]) - # if negative_prompt is not None: - # negative_prompt = negative_prompt.replace(prompt_replacement[0], prompt_replacement[1]) + if negative_prompt is not None: + negative_prompt = negative_prompt.replace(prompt_replacement[0], prompt_replacement[1]) if seed is not None: torch.manual_seed(seed) @@ -173,16 +173,18 @@ def sample_image_inference( torch.seed() torch.cuda.seed() - # if negative_prompt is None: - # negative_prompt = "" + if negative_prompt is None: + negative_prompt = "" height = max(64, height - height % 16) # round to divisible by 16 width = max(64, width - width % 16) # round to divisible by 16 logger.info(f"prompt: {prompt}") - # logger.info(f"negative_prompt: {negative_prompt}") + if scale != 1.0: + logger.info(f"negative_prompt: {negative_prompt}") logger.info(f"height: {height}") logger.info(f"width: {width}") logger.info(f"sample_steps: {sample_steps}") - logger.info(f"scale: {scale}") + if scale != 1.0: + logger.info(f"scale: {scale}") # logger.info(f"sample_sampler: {sampler_name}") if seed is not None: logger.info(f"seed: {seed}") @@ -191,26 +193,37 @@ def sample_image_inference( tokenize_strategy = strategy_base.TokenizeStrategy.get_strategy() encoding_strategy = strategy_base.TextEncodingStrategy.get_strategy() - text_encoder_conds = [] - if sample_prompts_te_outputs and prompt in sample_prompts_te_outputs: - text_encoder_conds = sample_prompts_te_outputs[prompt] - print(f"Using cached text encoder outputs for prompt: {prompt}") - if text_encoders is not None: - print(f"Encoding prompt: {prompt}") - tokens_and_masks = tokenize_strategy.tokenize(prompt) - # strategy has apply_t5_attn_mask option - encoded_text_encoder_conds = encoding_strategy.encode_tokens(tokenize_strategy, text_encoders, tokens_and_masks) + def encode_prompt(prpt): + text_encoder_conds = [] + if sample_prompts_te_outputs and prpt in sample_prompts_te_outputs: + text_encoder_conds = sample_prompts_te_outputs[prpt] + print(f"Using cached text encoder outputs for prompt: {prpt}") + if text_encoders is not None: + print(f"Encoding prompt: {prpt}") + tokens_and_masks = tokenize_strategy.tokenize(prpt) + # strategy has apply_t5_attn_mask option + encoded_text_encoder_conds = encoding_strategy.encode_tokens(tokenize_strategy, text_encoders, tokens_and_masks) - # if text_encoder_conds is not cached, use encoded_text_encoder_conds - if len(text_encoder_conds) == 0: - text_encoder_conds = encoded_text_encoder_conds - else: - # if encoded_text_encoder_conds is not None, update cached text_encoder_conds - for i in range(len(encoded_text_encoder_conds)): - if encoded_text_encoder_conds[i] is not None: - text_encoder_conds[i] = encoded_text_encoder_conds[i] + # if text_encoder_conds is not cached, use encoded_text_encoder_conds + if len(text_encoder_conds) == 0: + text_encoder_conds = encoded_text_encoder_conds + else: + # if encoded_text_encoder_conds is not None, update cached text_encoder_conds + for i in range(len(encoded_text_encoder_conds)): + if encoded_text_encoder_conds[i] is not None: + text_encoder_conds[i] = encoded_text_encoder_conds[i] + return text_encoder_conds - l_pooled, t5_out, txt_ids, t5_attn_mask = text_encoder_conds + l_pooled, t5_out, txt_ids, t5_attn_mask = encode_prompt(prompt) + # encode negative prompts + if scale != 1.0: + neg_l_pooled, neg_t5_out, _, neg_t5_attn_mask = encode_prompt(negative_prompt) + neg_t5_attn_mask = ( + neg_t5_attn_mask.to(accelerator.device) if args.apply_t5_attn_mask and neg_t5_attn_mask is not None else None + ) + neg_cond = (scale, neg_l_pooled, neg_t5_out, neg_t5_attn_mask) + else: + neg_cond = None # sample image weight_dtype = ae.dtype # TOFO give dtype as argument @@ -235,7 +248,20 @@ def sample_image_inference( controlnet_image = controlnet_image.permute(2, 0, 1).unsqueeze(0).to(weight_dtype).to(accelerator.device) with accelerator.autocast(), torch.no_grad(): - x = denoise(flux, noise, img_ids, t5_out, txt_ids, l_pooled, timesteps=timesteps, guidance=scale, t5_attn_mask=t5_attn_mask, controlnet=controlnet, controlnet_img=controlnet_image) + x = denoise( + flux, + noise, + img_ids, + t5_out, + txt_ids, + l_pooled, + timesteps=timesteps, + guidance=scale, + t5_attn_mask=t5_attn_mask, + controlnet=controlnet, + controlnet_img=controlnet_image, + neg_cond=neg_cond, + ) x = flux_utils.unpack_latents(x, packed_latent_height, packed_latent_width) @@ -305,22 +331,24 @@ def denoise( model: flux_models.Flux, img: torch.Tensor, img_ids: torch.Tensor, - txt: torch.Tensor, + txt: torch.Tensor, # t5_out txt_ids: torch.Tensor, - vec: torch.Tensor, + vec: torch.Tensor, # l_pooled timesteps: list[float], guidance: float = 4.0, t5_attn_mask: Optional[torch.Tensor] = None, controlnet: Optional[flux_models.ControlNetFlux] = None, controlnet_img: Optional[torch.Tensor] = None, + neg_cond: Optional[Tuple[float, torch.Tensor, torch.Tensor, torch.Tensor]] = None, ): # this is ignored for schnell guidance_vec = torch.full((img.shape[0],), guidance, device=img.device, dtype=img.dtype) - + do_cfg = neg_cond is not None for t_curr, t_prev in zip(tqdm(timesteps[:-1]), timesteps[1:]): t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device) model.prepare_block_swap_before_forward() + if controlnet is not None: block_samples, block_single_samples = controlnet( img=img, @@ -336,20 +364,48 @@ def denoise( else: block_samples = None block_single_samples = None - pred = model( - img=img, - img_ids=img_ids, - txt=txt, - txt_ids=txt_ids, - y=vec, - block_controlnet_hidden_states=block_samples, - block_controlnet_single_hidden_states=block_single_samples, - timesteps=t_vec, - guidance=guidance_vec, - txt_attention_mask=t5_attn_mask, - ) - img = img + (t_prev - t_curr) * pred + if not do_cfg: + pred = model( + img=img, + img_ids=img_ids, + txt=txt, + txt_ids=txt_ids, + y=vec, + block_controlnet_hidden_states=block_samples, + block_controlnet_single_hidden_states=block_single_samples, + timesteps=t_vec, + guidance=guidance_vec, + txt_attention_mask=t5_attn_mask, + ) + + img = img + (t_prev - t_curr) * pred + else: + cfg_scale, neg_l_pooled, neg_t5_out, neg_t5_attn_mask = neg_cond + nc_c_t5_attn_mask = None if t5_attn_mask is None else torch.cat([neg_t5_attn_mask, t5_attn_mask], dim=0) + + # TODO is it ok to use the same block samples for both cond and uncond? + block_samples = None if block_samples is None else torch.cat([block_samples, block_samples], dim=0) + block_single_samples = ( + None if block_single_samples is None else torch.cat([block_single_samples, block_single_samples], dim=0) + ) + + nc_c_pred = model( + img=torch.cat([img, img], dim=0), + img_ids=torch.cat([img_ids, img_ids], dim=0), + txt=torch.cat([neg_t5_out, txt], dim=0), + txt_ids=torch.cat([txt_ids, txt_ids], dim=0), + y=torch.cat([neg_l_pooled, vec], dim=0), + block_controlnet_hidden_states=block_samples, + block_controlnet_single_hidden_states=block_single_samples, + timesteps=t_vec, + guidance=guidance_vec, + txt_attention_mask=nc_c_t5_attn_mask, + ) + neg_pred, pred = torch.chunk(nc_c_pred, 2, dim=0) + pred = neg_pred + (pred - neg_pred) * cfg_scale + + img = img + (t_prev - t_curr) * pred model.prepare_block_swap_before_forward() return img @@ -567,7 +623,7 @@ def add_flux_train_arguments(parser: argparse.ArgumentParser): "--controlnet_model_name_or_path", type=str, default=None, - help="path to controlnet (*.sft or *.safetensors) / controlnetのパス（*.sftまたは*.safetensors）" + help="path to controlnet (*.sft or *.safetensors) / controlnetのパス（*.sftまたは*.safetensors）", ) parser.add_argument( "--t5xxl_max_token_length", From 0b25a05e3c0b983d7a4fa74f40798705a00992e3 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 18 Mar 2025 15:40:40 -0400 Subject: [PATCH 02/40] Add IP noise gamma for Flux --- library/flux_train_utils.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index f7f06c5c..f866fd4a 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -415,6 +415,16 @@ def get_noisy_model_input_and_timesteps( bsz, _, h, w = latents.shape sigmas = None + ip_noise_gamma = 0.0 + + # Add noise to the latents according to the noise magnitude at each timestep + # (this is the forward diffusion process) + if args.ip_noise_gamma: + if args.ip_noise_gamma_random_strength: + ip_noise_gamma = torch.rand(1, device=latents.device) * args.ip_noise_gamma + else: + ip_noise_gamma = args.ip_noise_gamma + if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling if args.timestep_sampling == "sigmoid": @@ -425,7 +435,7 @@ def get_noisy_model_input_and_timesteps( timesteps = t * 1000.0 t = t.view(-1, 1, 1, 1) - noisy_model_input = (1 - t) * latents + t * noise + noisy_model_input = (1 - t) * latents + t * noise + ip_noise_gamma elif args.timestep_sampling == "shift": shift = args.discrete_flow_shift logits_norm = torch.randn(bsz, device=device) @@ -435,7 +445,7 @@ def get_noisy_model_input_and_timesteps( t = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * noise + noisy_model_input = (1 - t) * latents + t * noise + ip_noise_gamma elif args.timestep_sampling == "flux_shift": logits_norm = torch.randn(bsz, device=device) logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling @@ -445,7 +455,7 @@ def get_noisy_model_input_and_timesteps( t = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * noise + noisy_model_input = (1 - t) * latents + t * noise + ip_noise_gamma else: # Sample a random timestep for each image # for weighting schemes where we sample timesteps non-uniformly @@ -461,7 +471,8 @@ def get_noisy_model_input_and_timesteps( # Add noise according to flow matching. sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) - noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents + noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents + ip_noise_gamma + return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From c8be141ae0576119ecd8ae329f00700098ee83a2 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 18 Mar 2025 15:42:18 -0400 Subject: [PATCH 03/40] Apply IP gamma to noise fix --- library/flux_train_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index f866fd4a..557f61e7 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -471,7 +471,7 @@ def get_noisy_model_input_and_timesteps( # Add noise according to flow matching. sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) - noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents + ip_noise_gamma + noisy_model_input = sigmas * noise + ip_noise_gamma + (1.0 - sigmas) * latents return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From b425466e7be64e12238b267862468dc9f0b0bb6e Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 18 Mar 2025 18:42:35 -0400 Subject: [PATCH 04/40] Fix IP noise gamma to use random values --- library/flux_train_utils.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 557f61e7..f0744747 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -415,15 +415,15 @@ def get_noisy_model_input_and_timesteps( bsz, _, h, w = latents.shape sigmas = None - ip_noise_gamma = 0.0 - # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) if args.ip_noise_gamma: if args.ip_noise_gamma_random_strength: - ip_noise_gamma = torch.rand(1, device=latents.device) * args.ip_noise_gamma + ip_noise = (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) else: - ip_noise_gamma = args.ip_noise_gamma + ip_noise = args.ip_noise_gamma * torch.randn_like(latents) + else: + ip_noise = torch.zeros_like(latents) if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling @@ -435,7 +435,7 @@ def get_noisy_model_input_and_timesteps( timesteps = t * 1000.0 t = t.view(-1, 1, 1, 1) - noisy_model_input = (1 - t) * latents + t * noise + ip_noise_gamma + noisy_model_input = (1 - t) * latents + t * (noise + ip_noise) elif args.timestep_sampling == "shift": shift = args.discrete_flow_shift logits_norm = torch.randn(bsz, device=device) @@ -445,7 +445,7 @@ def get_noisy_model_input_and_timesteps( t = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * noise + ip_noise_gamma + noisy_model_input = (1 - t) * latents + t * (noise + ip_noise) elif args.timestep_sampling == "flux_shift": logits_norm = torch.randn(bsz, device=device) logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling @@ -455,7 +455,7 @@ def get_noisy_model_input_and_timesteps( t = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * noise + ip_noise_gamma + noisy_model_input = (1 - t) * latents + t * (noise + ip_noise) else: # Sample a random timestep for each image # for weighting schemes where we sample timesteps non-uniformly @@ -471,7 +471,7 @@ def get_noisy_model_input_and_timesteps( # Add noise according to flow matching. sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) - noisy_model_input = sigmas * noise + ip_noise_gamma + (1.0 - sigmas) * latents + noisy_model_input = sigmas * (noise + ip_noise) + (1.0 - sigmas) * latents return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From a4f3a9fc1a4f4f964a6971bc4b0ae15c94f0d672 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 18 Mar 2025 18:44:21 -0400 Subject: [PATCH 05/40] Use ones_like --- library/flux_train_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index f0744747..8cf95858 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -423,7 +423,7 @@ def get_noisy_model_input_and_timesteps( else: ip_noise = args.ip_noise_gamma * torch.randn_like(latents) else: - ip_noise = torch.zeros_like(latents) + ip_noise = torch.ones_like(latents) if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling From 6f4d3657756a9d679dfa76f7c6c7bd1c957130ca Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 18 Mar 2025 18:53:34 -0400 Subject: [PATCH 06/40] zeros_like because we are adding --- library/flux_train_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 8cf95858..f0744747 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -423,7 +423,7 @@ def get_noisy_model_input_and_timesteps( else: ip_noise = args.ip_noise_gamma * torch.randn_like(latents) else: - ip_noise = torch.ones_like(latents) + ip_noise = torch.zeros_like(latents) if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling From b81bcd0b01aa81bf616b6125ca1da4d6d3c9dd82 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Tue, 18 Mar 2025 21:36:55 -0400 Subject: [PATCH 07/40] Move IP noise gamma to noise creation to remove complexity and align noise for target loss --- flux_train_network.py | 9 +++++++++ library/flux_train_utils.py | 19 ++++--------------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/flux_train_network.py b/flux_train_network.py index def44155..d85584f5 100644 --- a/flux_train_network.py +++ b/flux_train_network.py @@ -350,6 +350,15 @@ class FluxNetworkTrainer(train_network.NetworkTrainer): ): # Sample noise that we'll add to the latents noise = torch.randn_like(latents) + + # Add noise to the latents according to the noise magnitude at each timestep + # (this is the forward diffusion process) + if args.ip_noise_gamma: + if args.ip_noise_gamma_random_strength: + noise = noise + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) + else: + noise = noise + args.ip_noise_gamma * torch.randn_like(latents) + bsz = latents.shape[0] # get noisy model input and timesteps diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index f0744747..f7f06c5c 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -415,16 +415,6 @@ def get_noisy_model_input_and_timesteps( bsz, _, h, w = latents.shape sigmas = None - # Add noise to the latents according to the noise magnitude at each timestep - # (this is the forward diffusion process) - if args.ip_noise_gamma: - if args.ip_noise_gamma_random_strength: - ip_noise = (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) - else: - ip_noise = args.ip_noise_gamma * torch.randn_like(latents) - else: - ip_noise = torch.zeros_like(latents) - if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling if args.timestep_sampling == "sigmoid": @@ -435,7 +425,7 @@ def get_noisy_model_input_and_timesteps( timesteps = t * 1000.0 t = t.view(-1, 1, 1, 1) - noisy_model_input = (1 - t) * latents + t * (noise + ip_noise) + noisy_model_input = (1 - t) * latents + t * noise elif args.timestep_sampling == "shift": shift = args.discrete_flow_shift logits_norm = torch.randn(bsz, device=device) @@ -445,7 +435,7 @@ def get_noisy_model_input_and_timesteps( t = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * (noise + ip_noise) + noisy_model_input = (1 - t) * latents + t * noise elif args.timestep_sampling == "flux_shift": logits_norm = torch.randn(bsz, device=device) logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling @@ -455,7 +445,7 @@ def get_noisy_model_input_and_timesteps( t = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * (noise + ip_noise) + noisy_model_input = (1 - t) * latents + t * noise else: # Sample a random timestep for each image # for weighting schemes where we sample timesteps non-uniformly @@ -471,8 +461,7 @@ def get_noisy_model_input_and_timesteps( # Add noise according to flow matching. sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) - noisy_model_input = sigmas * (noise + ip_noise) + (1.0 - sigmas) * latents - + noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From 7197266703d8ac9219dda8b5a58bbd60d029d597 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Wed, 19 Mar 2025 00:25:51 -0400 Subject: [PATCH 08/40] Perturbed noise should be separate of input noise --- flux_train_network.py | 9 --------- library/flux_train_utils.py | 13 ++++++++++++- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/flux_train_network.py b/flux_train_network.py index d85584f5..def44155 100644 --- a/flux_train_network.py +++ b/flux_train_network.py @@ -350,15 +350,6 @@ class FluxNetworkTrainer(train_network.NetworkTrainer): ): # Sample noise that we'll add to the latents noise = torch.randn_like(latents) - - # Add noise to the latents according to the noise magnitude at each timestep - # (this is the forward diffusion process) - if args.ip_noise_gamma: - if args.ip_noise_gamma_random_strength: - noise = noise + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) - else: - noise = noise + args.ip_noise_gamma * torch.randn_like(latents) - bsz = latents.shape[0] # get noisy model input and timesteps diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index f7f06c5c..775e0c33 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -410,11 +410,22 @@ def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None): def get_noisy_model_input_and_timesteps( - args, noise_scheduler, latents, noise, device, dtype + args, noise_scheduler, latents: torch.Tensor, input_noise: torch.Tensor, device, dtype ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, _, h, w = latents.shape sigmas = None + # Add noise to the latents according to the noise magnitude at each timestep + # (this is the forward diffusion process) + if args.ip_noise_gamma: + if args.ip_noise_gamma_random_strength: + noise = input_noise.detach().clone() + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) + else: + noise = input_noise.detach().clone() + args.ip_noise_gamma * torch.randn_like(latents) + else: + noise = input_noise + + if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling if args.timestep_sampling == "sigmoid": From d93ad90a717beb2fd322d2fae73992e9ea5213ea Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Wed, 19 Mar 2025 00:37:27 -0400 Subject: [PATCH 09/40] Add perturbation on noisy_model_input if needed --- library/flux_train_utils.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 775e0c33..0fe81da7 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -410,20 +410,11 @@ def compute_loss_weighting_for_sd3(weighting_scheme: str, sigmas=None): def get_noisy_model_input_and_timesteps( - args, noise_scheduler, latents: torch.Tensor, input_noise: torch.Tensor, device, dtype + args, noise_scheduler, latents: torch.Tensor, noise: torch.Tensor, device, dtype ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, _, h, w = latents.shape sigmas = None - # Add noise to the latents according to the noise magnitude at each timestep - # (this is the forward diffusion process) - if args.ip_noise_gamma: - if args.ip_noise_gamma_random_strength: - noise = input_noise.detach().clone() + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) - else: - noise = input_noise.detach().clone() + args.ip_noise_gamma * torch.randn_like(latents) - else: - noise = input_noise if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": @@ -474,6 +465,15 @@ def get_noisy_model_input_and_timesteps( sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents + # Add noise to the latents according to the noise magnitude at each timestep + # (this is the forward diffusion process) + if args.ip_noise_gamma: + if args.ip_noise_gamma_random_strength: + xi = noise.detach().clone() + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) + else: + xi = noise.detach().clone() + args.ip_noise_gamma * torch.randn_like(latents) + noisy_model_input += xi + return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From 8e6817b0c2d6e312b8da0d84baa2ecc72c83767f Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Wed, 19 Mar 2025 00:45:13 -0400 Subject: [PATCH 10/40] Remove double noise --- library/flux_train_utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 0fe81da7..9808ad0a 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -415,8 +415,6 @@ def get_noisy_model_input_and_timesteps( bsz, _, h, w = latents.shape sigmas = None - - if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling if args.timestep_sampling == "sigmoid": @@ -469,10 +467,10 @@ def get_noisy_model_input_and_timesteps( # (this is the forward diffusion process) if args.ip_noise_gamma: if args.ip_noise_gamma_random_strength: - xi = noise.detach().clone() + (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(latents) + noise_perturbation = (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(noise) else: - xi = noise.detach().clone() + args.ip_noise_gamma * torch.randn_like(latents) - noisy_model_input += xi + noise_perturbation = args.ip_noise_gamma * torch.randn_like(noise) + noisy_model_input += noise_perturbation return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From 1eddac26b010d23ce5f0eb6a8ac12fbca66ee50b Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Wed, 19 Mar 2025 00:49:42 -0400 Subject: [PATCH 11/40] Separate random to a variable, and make sure on device --- library/flux_train_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 9808ad0a..107f351f 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -466,11 +466,12 @@ def get_noisy_model_input_and_timesteps( # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) if args.ip_noise_gamma: + xi = torch.randn_like(latents, device=latents.device, dtype=dtype) if args.ip_noise_gamma_random_strength: - noise_perturbation = (torch.rand(1, device=latents.device) * args.ip_noise_gamma) * torch.randn_like(noise) + ip_noise_gamma = (torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma) else: - noise_perturbation = args.ip_noise_gamma * torch.randn_like(noise) - noisy_model_input += noise_perturbation + ip_noise_gamma = args.ip_noise_gamma + noisy_model_input += ip_noise_gamma * xi return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From 5d5a7d2acf884077b6a24db269c8f4facb5b7487 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Wed, 19 Mar 2025 13:50:04 -0400 Subject: [PATCH 12/40] Fix IP noise calculation --- library/flux_train_utils.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 107f351f..0cb07e3d 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -423,29 +423,24 @@ def get_noisy_model_input_and_timesteps( else: t = torch.rand((bsz,), device=device) + sigmas = t.view(-1, 1, 1, 1) timesteps = t * 1000.0 - t = t.view(-1, 1, 1, 1) - noisy_model_input = (1 - t) * latents + t * noise elif args.timestep_sampling == "shift": shift = args.discrete_flow_shift logits_norm = torch.randn(bsz, device=device) logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling timesteps = logits_norm.sigmoid() timesteps = (timesteps * shift) / (1 + (shift - 1) * timesteps) - - t = timesteps.view(-1, 1, 1, 1) + sigmas = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * noise elif args.timestep_sampling == "flux_shift": logits_norm = torch.randn(bsz, device=device) logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling timesteps = logits_norm.sigmoid() mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) timesteps = time_shift(mu, 1.0, timesteps) - - t = timesteps.view(-1, 1, 1, 1) + sigmas = timesteps.view(-1, 1, 1, 1) timesteps = timesteps * 1000.0 - noisy_model_input = (1 - t) * latents + t * noise else: # Sample a random timestep for each image # for weighting schemes where we sample timesteps non-uniformly @@ -458,10 +453,7 @@ def get_noisy_model_input_and_timesteps( ) indices = (u * noise_scheduler.config.num_train_timesteps).long() timesteps = noise_scheduler.timesteps[indices].to(device=device) - - # Add noise according to flow matching. sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) - noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) @@ -471,7 +463,9 @@ def get_noisy_model_input_and_timesteps( ip_noise_gamma = (torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma) else: ip_noise_gamma = args.ip_noise_gamma - noisy_model_input += ip_noise_gamma * xi + noisy_model_input = sigmas * (noise + ip_noise_gamma * xi) + (1.0 - sigmas) * latents + else: + noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From f974c6b2577348acbe948bcc668dd7b061feb73e Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Wed, 19 Mar 2025 14:27:43 -0400 Subject: [PATCH 13/40] change order to match upstream --- library/flux_train_utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 0cb07e3d..7bf2faf0 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -413,8 +413,6 @@ def get_noisy_model_input_and_timesteps( args, noise_scheduler, latents: torch.Tensor, noise: torch.Tensor, device, dtype ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, _, h, w = latents.shape - sigmas = None - if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random t-based noise sampling if args.timestep_sampling == "sigmoid": @@ -463,9 +461,9 @@ def get_noisy_model_input_and_timesteps( ip_noise_gamma = (torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma) else: ip_noise_gamma = args.ip_noise_gamma - noisy_model_input = sigmas * (noise + ip_noise_gamma * xi) + (1.0 - sigmas) * latents + noisy_model_input = (1.0 - sigmas) * latents + sigmas * (noise + ip_noise_gamma * xi) else: - noisy_model_input = sigmas * noise + (1.0 - sigmas) * latents + noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise return noisy_model_input.to(dtype), timesteps.to(dtype), sigmas From 16cef81aeaec1ebc07de30c7a1448982a61167e1 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 20 Mar 2025 14:32:56 -0400 Subject: [PATCH 14/40] Refactor sigmas and timesteps --- library/flux_train_utils.py | 41 ++++++++++++++++++------------------- 1 file changed, 20 insertions(+), 21 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 7bf2faf0..9110da89 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -366,8 +366,6 @@ def get_sigmas(noise_scheduler, timesteps, device, n_dim=4, dtype=torch.float32) 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 @@ -413,32 +411,30 @@ def get_noisy_model_input_and_timesteps( args, noise_scheduler, latents: torch.Tensor, noise: torch.Tensor, device, dtype ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, _, h, w = latents.shape + num_timesteps = noise_scheduler.config.num_train_timesteps if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": - # Simple random t-based noise sampling + # Simple random sigma-based noise sampling if args.timestep_sampling == "sigmoid": # https://github.com/XLabs-AI/x-flux/tree/main - t = torch.sigmoid(args.sigmoid_scale * torch.randn((bsz,), device=device)) + sigmas = torch.sigmoid(args.sigmoid_scale * torch.randn((bsz,), device=device)) else: - t = torch.rand((bsz,), device=device) + sigmas = torch.rand((bsz,), device=device) - sigmas = t.view(-1, 1, 1, 1) - timesteps = t * 1000.0 + timesteps = sigmas * num_timesteps elif args.timestep_sampling == "shift": shift = args.discrete_flow_shift - logits_norm = torch.randn(bsz, device=device) - logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling - timesteps = logits_norm.sigmoid() - timesteps = (timesteps * shift) / (1 + (shift - 1) * timesteps) - sigmas = timesteps.view(-1, 1, 1, 1) - timesteps = timesteps * 1000.0 + sigmas = torch.randn(bsz, device=device) + sigmas = sigmas * args.sigmoid_scale # larger scale for more uniform sampling + sigmas = sigmas.sigmoid() + sigmas = (sigmas * shift) / (1 + (shift - 1) * sigmas) + timesteps = sigmas * num_timesteps elif args.timestep_sampling == "flux_shift": - logits_norm = torch.randn(bsz, device=device) - logits_norm = logits_norm * args.sigmoid_scale # larger scale for more uniform sampling - timesteps = logits_norm.sigmoid() - mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) - timesteps = time_shift(mu, 1.0, timesteps) - sigmas = timesteps.view(-1, 1, 1, 1) - timesteps = timesteps * 1000.0 + sigmas = torch.randn(bsz, device=device) + sigmas = sigmas * args.sigmoid_scale # larger scale for more uniform sampling + sigmas = sigmas.sigmoid() + mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) # we are pre-packed so must adjust for packed size + sigmas = time_shift(mu, 1.0, sigmas) + timesteps = sigmas * num_timesteps else: # Sample a random timestep for each image # for weighting schemes where we sample timesteps non-uniformly @@ -449,10 +445,13 @@ def get_noisy_model_input_and_timesteps( logit_std=args.logit_std, mode_scale=args.mode_scale, ) - indices = (u * noise_scheduler.config.num_train_timesteps).long() + indices = (u * num_timesteps).long() timesteps = noise_scheduler.timesteps[indices].to(device=device) sigmas = get_sigmas(noise_scheduler, timesteps, device, n_dim=latents.ndim, dtype=dtype) + # Broadcast sigmas to latent shape + sigmas = sigmas.view(-1, 1, 1, 1) + # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process) if args.ip_noise_gamma: From e8b32548580ebf0001cd457d7b6f796e2eb169ff Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 20 Mar 2025 15:01:15 -0400 Subject: [PATCH 15/40] Add flux_train_utils tests for get get_noisy_model_input_and_timesteps --- library/flux_train_utils.py | 1 + tests/library/test_flux_train_utils.py | 220 +++++++++++++++++++++++++ 2 files changed, 221 insertions(+) create mode 100644 tests/library/test_flux_train_utils.py diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index 9110da89..0e73a01d 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -411,6 +411,7 @@ def get_noisy_model_input_and_timesteps( args, noise_scheduler, latents: torch.Tensor, noise: torch.Tensor, device, dtype ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bsz, _, h, w = latents.shape + assert bsz > 0, "Batch size not large enough" num_timesteps = noise_scheduler.config.num_train_timesteps if args.timestep_sampling == "uniform" or args.timestep_sampling == "sigmoid": # Simple random sigma-based noise sampling diff --git a/tests/library/test_flux_train_utils.py b/tests/library/test_flux_train_utils.py new file mode 100644 index 00000000..a4c7ba3b --- /dev/null +++ b/tests/library/test_flux_train_utils.py @@ -0,0 +1,220 @@ +import pytest +import torch +from unittest.mock import MagicMock, patch +from library.flux_train_utils import ( + get_noisy_model_input_and_timesteps, +) + +# Mock classes and functions +class MockNoiseScheduler: + def __init__(self, num_train_timesteps=1000): + self.config = MagicMock() + self.config.num_train_timesteps = num_train_timesteps + self.timesteps = torch.arange(num_train_timesteps, dtype=torch.long) + + +# Create fixtures for commonly used objects +@pytest.fixture +def args(): + args = MagicMock() + args.timestep_sampling = "uniform" + args.weighting_scheme = "uniform" + args.logit_mean = 0.0 + args.logit_std = 1.0 + args.mode_scale = 1.0 + args.sigmoid_scale = 1.0 + args.discrete_flow_shift = 3.1582 + args.ip_noise_gamma = None + args.ip_noise_gamma_random_strength = False + return args + + +@pytest.fixture +def noise_scheduler(): + return MockNoiseScheduler(num_train_timesteps=1000) + + +@pytest.fixture +def latents(): + return torch.randn(2, 4, 8, 8) + + +@pytest.fixture +def noise(): + return torch.randn(2, 4, 8, 8) + + +@pytest.fixture +def device(): + # return "cuda" if torch.cuda.is_available() else "cpu" + return "cpu" + + +# Mock the required functions +@pytest.fixture(autouse=True) +def mock_functions(): + with ( + patch("torch.sigmoid", side_effect=torch.sigmoid), + patch("torch.rand", side_effect=torch.rand), + patch("torch.randn", side_effect=torch.randn), + ): + yield + + +# Test different timestep sampling methods +def test_uniform_sampling(args, noise_scheduler, latents, noise, device): + args.timestep_sampling = "uniform" + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + assert noisy_input.dtype == dtype + assert timesteps.dtype == dtype + + +def test_sigmoid_sampling(args, noise_scheduler, latents, noise, device): + args.timestep_sampling = "sigmoid" + args.sigmoid_scale = 10.0 + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + + +def test_shift_sampling(args, noise_scheduler, latents, noise, device): + args.timestep_sampling = "shift" + args.sigmoid_scale = 1.0 + args.discrete_flow_shift = 3.1582 + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + + +def test_flux_shift_sampling(args, noise_scheduler, latents, noise, device): + args.timestep_sampling = "flux_shift" + args.sigmoid_scale = 10.0 + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + + +def test_weighting_scheme(args, noise_scheduler, latents, noise, device): + # Mock the necessary functions for this specific test + with patch("library.flux_train_utils.compute_density_for_timestep_sampling", + return_value=torch.tensor([0.3, 0.7], device=device)), \ + patch("library.flux_train_utils.get_sigmas", + return_value=torch.tensor([[0.3], [0.7]], device=device).view(-1, 1, 1, 1)): + + args.timestep_sampling = "other" # Will trigger the weighting scheme path + args.weighting_scheme = "uniform" + args.logit_mean = 0.0 + args.logit_std = 1.0 + args.mode_scale = 1.0 + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps( + args, noise_scheduler, latents, noise, device, dtype + ) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + + +# Test IP noise options +def test_with_ip_noise(args, noise_scheduler, latents, noise, device): + args.ip_noise_gamma = 0.5 + args.ip_noise_gamma_random_strength = False + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + + +def test_with_random_ip_noise(args, noise_scheduler, latents, noise, device): + args.ip_noise_gamma = 0.1 + args.ip_noise_gamma_random_strength = True + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (latents.shape[0],) + assert sigmas.shape == (latents.shape[0], 1, 1, 1) + + +# Test different data types +def test_float16_dtype(args, noise_scheduler, latents, noise, device): + dtype = torch.float16 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.dtype == dtype + assert timesteps.dtype == dtype + + +# Test different batch sizes +def test_different_batch_size(args, noise_scheduler, device): + latents = torch.randn(5, 4, 8, 8) # batch size of 5 + noise = torch.randn(5, 4, 8, 8) + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (5,) + assert sigmas.shape == (5, 1, 1, 1) + + +# Test different image sizes +def test_different_image_size(args, noise_scheduler, device): + latents = torch.randn(2, 4, 16, 16) # larger image size + noise = torch.randn(2, 4, 16, 16) + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (2,) + assert sigmas.shape == (2, 1, 1, 1) + + +# Test edge cases +def test_zero_batch_size(args, noise_scheduler, device): + with pytest.raises(AssertionError): # expecting an error with zero batch size + latents = torch.randn(0, 4, 8, 8) + noise = torch.randn(0, 4, 8, 8) + dtype = torch.float32 + + get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + +def test_different_timestep_count(args, device): + noise_scheduler = MockNoiseScheduler(num_train_timesteps=500) # different timestep count + latents = torch.randn(2, 4, 8, 8) + noise = torch.randn(2, 4, 8, 8) + dtype = torch.float32 + + noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) + + assert noisy_input.shape == latents.shape + assert timesteps.shape == (2,) + # Check that timesteps are within the proper range + assert torch.all(timesteps < 500) From 8aa126582efbdf0472b0b8db800d50860870f3cd Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 20 Mar 2025 15:09:11 -0400 Subject: [PATCH 16/40] Scale sigmoid to default 1.0 --- pytest.ini | 1 + requirements.txt | 2 +- tests/library/test_flux_train_utils.py | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pytest.ini b/pytest.ini index 484d3aef..34b7e9c1 100644 --- a/pytest.ini +++ b/pytest.ini @@ -6,3 +6,4 @@ filterwarnings = ignore::DeprecationWarning ignore::UserWarning ignore::FutureWarning +pythonpath = . diff --git a/requirements.txt b/requirements.txt index de39f588..8fe8c762 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ opencv-python==4.8.1.78 einops==0.7.0 pytorch-lightning==1.9.0 bitsandbytes==0.44.0 -prodigyopt==1.0 +prodigyopt>=1.0 lion-pytorch==0.0.6 schedulefree==1.4 tensorboard diff --git a/tests/library/test_flux_train_utils.py b/tests/library/test_flux_train_utils.py index a4c7ba3b..2ad7ce4e 100644 --- a/tests/library/test_flux_train_utils.py +++ b/tests/library/test_flux_train_utils.py @@ -77,7 +77,7 @@ def test_uniform_sampling(args, noise_scheduler, latents, noise, device): def test_sigmoid_sampling(args, noise_scheduler, latents, noise, device): args.timestep_sampling = "sigmoid" - args.sigmoid_scale = 10.0 + args.sigmoid_scale = 1.0 dtype = torch.float32 noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) @@ -102,7 +102,7 @@ def test_shift_sampling(args, noise_scheduler, latents, noise, device): def test_flux_shift_sampling(args, noise_scheduler, latents, noise, device): args.timestep_sampling = "flux_shift" - args.sigmoid_scale = 10.0 + args.sigmoid_scale = 1.0 dtype = torch.float32 noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) From d40f5b1e4ef5e7e6b51df26914be3a661b006d34 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 20 Mar 2025 15:09:50 -0400 Subject: [PATCH 17/40] Revert "Scale sigmoid to default 1.0" This reverts commit 8aa126582efbdf0472b0b8db800d50860870f3cd. --- pytest.ini | 1 - requirements.txt | 2 +- tests/library/test_flux_train_utils.py | 4 ++-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/pytest.ini b/pytest.ini index 34b7e9c1..484d3aef 100644 --- a/pytest.ini +++ b/pytest.ini @@ -6,4 +6,3 @@ filterwarnings = ignore::DeprecationWarning ignore::UserWarning ignore::FutureWarning -pythonpath = . diff --git a/requirements.txt b/requirements.txt index 8fe8c762..de39f588 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ opencv-python==4.8.1.78 einops==0.7.0 pytorch-lightning==1.9.0 bitsandbytes==0.44.0 -prodigyopt>=1.0 +prodigyopt==1.0 lion-pytorch==0.0.6 schedulefree==1.4 tensorboard diff --git a/tests/library/test_flux_train_utils.py b/tests/library/test_flux_train_utils.py index 2ad7ce4e..a4c7ba3b 100644 --- a/tests/library/test_flux_train_utils.py +++ b/tests/library/test_flux_train_utils.py @@ -77,7 +77,7 @@ def test_uniform_sampling(args, noise_scheduler, latents, noise, device): def test_sigmoid_sampling(args, noise_scheduler, latents, noise, device): args.timestep_sampling = "sigmoid" - args.sigmoid_scale = 1.0 + args.sigmoid_scale = 10.0 dtype = torch.float32 noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) @@ -102,7 +102,7 @@ def test_shift_sampling(args, noise_scheduler, latents, noise, device): def test_flux_shift_sampling(args, noise_scheduler, latents, noise, device): args.timestep_sampling = "flux_shift" - args.sigmoid_scale = 1.0 + args.sigmoid_scale = 10.0 dtype = torch.float32 noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) From 89f0d27a5930ae0a355caacfedc546fb04a7345d Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 20 Mar 2025 15:10:33 -0400 Subject: [PATCH 18/40] Set sigmoid_scale to default 1.0 --- tests/library/test_flux_train_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/library/test_flux_train_utils.py b/tests/library/test_flux_train_utils.py index a4c7ba3b..2ad7ce4e 100644 --- a/tests/library/test_flux_train_utils.py +++ b/tests/library/test_flux_train_utils.py @@ -77,7 +77,7 @@ def test_uniform_sampling(args, noise_scheduler, latents, noise, device): def test_sigmoid_sampling(args, noise_scheduler, latents, noise, device): args.timestep_sampling = "sigmoid" - args.sigmoid_scale = 10.0 + args.sigmoid_scale = 1.0 dtype = torch.float32 noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) @@ -102,7 +102,7 @@ def test_shift_sampling(args, noise_scheduler, latents, noise, device): def test_flux_shift_sampling(args, noise_scheduler, latents, noise, device): args.timestep_sampling = "flux_shift" - args.sigmoid_scale = 10.0 + args.sigmoid_scale = 1.0 dtype = torch.float32 noisy_input, timesteps, sigmas = get_noisy_model_input_and_timesteps(args, noise_scheduler, latents, noise, device, dtype) From 381303d64fa87c61da145b85b737c410d42555fa Mon Sep 17 00:00:00 2001 From: DKnight54 <126916963+DKnight54@users.noreply.github.com> Date: Sat, 29 Mar 2025 02:26:18 +0800 Subject: [PATCH 19/40] Update train_network.py --- train_network.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/train_network.py b/train_network.py index 7bf125dc..6953bb17 100644 --- a/train_network.py +++ b/train_network.py @@ -912,14 +912,22 @@ class NetworkTrainer: if "latents" in batch and batch["latents"] is not None: latents = batch["latents"].to(accelerator.device).to(dtype=weight_dtype) else: - with torch.no_grad(): - # latentに変換 - latents = vae.encode(batch["images"].to(dtype=vae_dtype)).latent_dist.sample().to(dtype=weight_dtype) - + if args.vae_batch_size is None or len(batch["images"]) <= args.vae_batch_size: + with torch.no_grad(): + # latentに変換 + latents = vae.encode(batch["images"].to(dtype=vae_dtype)).latent_dist.sample().to(dtype=weight_dtype) + else: + chunks = [batch["images"][i:i + args.vae_batch_size] for i in range(0, len(batch["images"]), args.vae_batch_size)] + list_latents = [] + for chunk in chunks: + with torch.no_grad(): + # latentに変換 + list_latents.append(vae.encode(chunk.to(dtype=vae_dtype)).latent_dist.sample().to(dtype=weight_dtype)) + latents = torch.cat(list_latents, dim=0) # NaNが含まれていれば警告を表示し0に置き換える - if torch.any(torch.isnan(latents)): - accelerator.print("NaN found in latents, replacing with zeros") - latents = torch.nan_to_num(latents, 0, out=latents) + if torch.any(torch.isnan(latents)): + accelerator.print("NaN found in latents, replacing with zeros") + latents = torch.nan_to_num(latents, 0, out=latents) latents = latents * self.vae_scale_factor # get multiplier for each sample From 4589262f8f35914b592992f4144bde5e746a6e36 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 6 Apr 2025 21:34:27 +0900 Subject: [PATCH 20/40] README.md: Update recent updates section to include IP noise gamma feature for FLUX.1 --- README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/README.md b/README.md index ae417d05..2e80a697 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,10 @@ The command to install PyTorch is as follows: ### Recent Updates +Apr 6, 2025: +- IP noise gamma has been enabled in FLUX.1. Thanks to rockerBOO for PR [#1992](https://github.com/kohya-ss/sd-scripts/pull/1992). See the PR for details. + - `--ip_noise_gamma` and `--ip_noise_gamma_random_strength` are available. + Mar 30, 2025: - LoRA-GGPO is added for FLUX.1 LoRA training. Thank you to rockerBOO for PR [#1974](https://github.com/kohya-ss/sd-scripts/pull/1974). - Specify `--network_args ggpo_sigma=0.03 ggpo_beta=0.01` in the command line or `network_args = ["ggpo_sigma=0.03", "ggpo_beta=0.01"]` in .toml file. See PR for details. From 629073cd9dd21296ca8aa97a5267d4dc7f6e5fdb Mon Sep 17 00:00:00 2001 From: Kohya S Date: Wed, 16 Apr 2025 21:50:36 +0900 Subject: [PATCH 21/40] Add guidance scale for prompt param and flux sampling --- library/flux_train_utils.py | 10 +++++++--- library/train_util.py | 5 +++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index ce381829..d2ff347d 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -154,6 +154,7 @@ def sample_image_inference( sample_steps = prompt_dict.get("sample_steps", 20) width = prompt_dict.get("width", 512) height = prompt_dict.get("height", 512) + guidance_scale = prompt_dict.get("guidance_scale", args.guidance_scale) scale = prompt_dict.get("scale", 1.0) # 1.0 means no guidance seed = prompt_dict.get("seed") controlnet_image = prompt_dict.get("controlnet_image") @@ -180,9 +181,12 @@ def sample_image_inference( logger.info(f"prompt: {prompt}") if scale != 1.0: logger.info(f"negative_prompt: {negative_prompt}") + elif negative_prompt != "": + logger.info(f"negative prompt is ignored because scale is 1.0") logger.info(f"height: {height}") logger.info(f"width: {width}") logger.info(f"sample_steps: {sample_steps}") + logger.info(f"guidance_scale: {guidance_scale}") if scale != 1.0: logger.info(f"scale: {scale}") # logger.info(f"sample_sampler: {sampler_name}") @@ -256,7 +260,7 @@ def sample_image_inference( txt_ids, l_pooled, timesteps=timesteps, - guidance=scale, + guidance=guidance_scale, t5_attn_mask=t5_attn_mask, controlnet=controlnet, controlnet_img=controlnet_image, @@ -489,7 +493,7 @@ def get_noisy_model_input_and_timesteps( sigmas = torch.randn(bsz, device=device) sigmas = sigmas * args.sigmoid_scale # larger scale for more uniform sampling sigmas = sigmas.sigmoid() - mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) # we are pre-packed so must adjust for packed size + mu = get_lin_function(y1=0.5, y2=1.15)((h // 2) * (w // 2)) # we are pre-packed so must adjust for packed size sigmas = time_shift(mu, 1.0, sigmas) timesteps = sigmas * num_timesteps else: @@ -514,7 +518,7 @@ def get_noisy_model_input_and_timesteps( if args.ip_noise_gamma: xi = torch.randn_like(latents, device=latents.device, dtype=dtype) if args.ip_noise_gamma_random_strength: - ip_noise_gamma = (torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma) + ip_noise_gamma = torch.rand(1, device=latents.device, dtype=dtype) * args.ip_noise_gamma else: ip_noise_gamma = args.ip_noise_gamma noisy_model_input = (1.0 - sigmas) * latents + sigmas * (noise + ip_noise_gamma * xi) diff --git a/library/train_util.py b/library/train_util.py index 6c39f8d9..e152f30f 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -6178,6 +6178,11 @@ def line_to_prompt_dict(line: str) -> dict: prompt_dict["scale"] = float(m.group(1)) continue + m = re.match(r"g ([\d\.]+)", parg, re.IGNORECASE) + if m: # guidance scale + prompt_dict["guidance_scale"] = float(m.group(1)) + continue + m = re.match(r"n (.+)", parg, re.IGNORECASE) if m: # negative prompt prompt_dict["negative_prompt"] = m.group(1) From 26db64be17835de5bff22bfc6d671ae1a2ffb4a4 Mon Sep 17 00:00:00 2001 From: Glen Date: Sat, 19 Apr 2025 11:54:12 -0600 Subject: [PATCH 22/40] fix: update hf_hub_download parameters to fix wd14 tagger regression --- finetune/tag_images_by_wd14_tagger.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/finetune/tag_images_by_wd14_tagger.py b/finetune/tag_images_by_wd14_tagger.py index cbc3d2d6..f8f6ddd9 100644 --- a/finetune/tag_images_by_wd14_tagger.py +++ b/finetune/tag_images_by_wd14_tagger.py @@ -100,15 +100,19 @@ def main(args): else: for file in SUB_DIR_FILES: hf_hub_download( - args.repo_id, - file, + repo_id=args.repo_id, + filename=file, subfolder=SUB_DIR, - cache_dir=os.path.join(model_location, SUB_DIR), + local_dir=os.path.join(model_location, SUB_DIR), force_download=True, - force_filename=file, ) for file in files: - hf_hub_download(args.repo_id, file, cache_dir=model_location, force_download=True, force_filename=file) + hf_hub_download( + repo_id=args.repo_id, + filename=file, + local_dir=model_location, + force_download=True, + ) else: logger.info("using existing wd14 tagger model") From 7c61c0dfe0e879fd6b66ccb70273e4b99deaf1c5 Mon Sep 17 00:00:00 2001 From: saibit Date: Tue, 22 Apr 2025 16:06:55 +0800 Subject: [PATCH 23/40] Add autocast warpper for forward functions in deepspeed_utils.py to try aligning precision when using mixed precision in training process --- library/deepspeed_utils.py | 32 ++++++++++++++++++++++++++++++++ library/flux_models.py | 2 +- library/train_util.py | 5 +++++ requirements.txt | 3 ++- 4 files changed, 40 insertions(+), 2 deletions(-) diff --git a/library/deepspeed_utils.py b/library/deepspeed_utils.py index 99a7b2b3..3018def7 100644 --- a/library/deepspeed_utils.py +++ b/library/deepspeed_utils.py @@ -94,6 +94,7 @@ def prepare_deepspeed_plugin(args: argparse.Namespace): deepspeed_plugin.deepspeed_config["train_batch_size"] = ( args.train_batch_size * args.gradient_accumulation_steps * int(os.environ["WORLD_SIZE"]) ) + deepspeed_plugin.set_mixed_precision(args.mixed_precision) if args.mixed_precision.lower() == "fp16": deepspeed_plugin.deepspeed_config["fp16"]["initial_scale_power"] = 0 # preventing overflow. @@ -122,18 +123,49 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models): class DeepSpeedWrapper(torch.nn.Module): def __init__(self, **kw_models) -> None: super().__init__() + self.models = torch.nn.ModuleDict() + + warp_model_forward_with_torch_autocast = args.mixed_precision is not "no" for key, model in kw_models.items(): if isinstance(model, list): model = torch.nn.ModuleList(model) + assert isinstance( model, torch.nn.Module ), f"model must be an instance of torch.nn.Module, but got {key} is {type(model)}" + + if warp_model_forward_with_torch_autocast: + model = self.__warp_with_torch_autocast(model) + self.models.update(torch.nn.ModuleDict({key: model})) + def __warp_with_torch_autocast(self, model): + if isinstance(model, torch.nn.ModuleList): + for i in range(len(model)): + model[i] = self.__warp_model_forward_with_torch_autocast(model[i]) + else: + model = self.__warp_model_forward_with_torch_autocast(model) + return model + + def __warp_model_forward_with_torch_autocast(self, model): + + assert hasattr(model, "forward"), f"model must have a forward method." + + forward_fn = model.forward + + def forward(*args, **kwargs): + device_type= "cuda" if torch.cuda.is_available() else "cpu" + with torch.autocast(device_type=device_type): + return forward_fn(*args, **kwargs) + model.forward = forward + + return model + def get_models(self): return self.models + ds_model = DeepSpeedWrapper(**models) return ds_model diff --git a/library/flux_models.py b/library/flux_models.py index 328ad481..12151ee8 100644 --- a/library/flux_models.py +++ b/library/flux_models.py @@ -1005,7 +1005,7 @@ class Flux(nn.Module): return self.offloader_double.prepare_block_devices_before_forward(self.double_blocks) self.offloader_single.prepare_block_devices_before_forward(self.single_blocks) - + def forward( self, img: Tensor, diff --git a/library/train_util.py b/library/train_util.py index 6c39f8d9..dbbfda3e 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -5495,6 +5495,11 @@ def load_target_model(args, weight_dtype, accelerator, unet_use_linear_projectio def patch_accelerator_for_fp16_training(accelerator): + + from accelerate import DistributedType + if accelerator.distributed_type == DistributedType.DEEPSPEED: + return + org_unscale_grads = accelerator.scaler._unscale_grads_ def _unscale_grads_replacer(optimizer, inv_scale, found_inf, allow_fp16): diff --git a/requirements.txt b/requirements.txt index 767d9e8e..bead3f90 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ accelerate==0.33.0 transformers==4.44.0 -diffusers[torch]==0.25.0 +diffusers==0.25.0 +deepspeed==0.16.7 ftfy==6.1.1 # albumentations==1.3.0 opencv-python==4.8.1.78 From d33d5eccd16970e489359ee02b89a6259559e4b9 Mon Sep 17 00:00:00 2001 From: saibit Date: Tue, 22 Apr 2025 16:12:06 +0800 Subject: [PATCH 24/40] # --- library/flux_models.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/library/flux_models.py b/library/flux_models.py index 12151ee8..d7840d51 100644 --- a/library/flux_models.py +++ b/library/flux_models.py @@ -1004,8 +1004,7 @@ class Flux(nn.Module): if self.blocks_to_swap is None or self.blocks_to_swap == 0: return self.offloader_double.prepare_block_devices_before_forward(self.double_blocks) - self.offloader_single.prepare_block_devices_before_forward(self.single_blocks) - + self.offloader_single.prepare_block_devices_before_forward(self.single_blocks) def forward( self, img: Tensor, From 7f984f47758f9e17f4a82b92cb9dbc97b3ba982f Mon Sep 17 00:00:00 2001 From: saibit Date: Tue, 22 Apr 2025 16:15:12 +0800 Subject: [PATCH 25/40] # --- library/flux_models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/library/flux_models.py b/library/flux_models.py index d7840d51..328ad481 100644 --- a/library/flux_models.py +++ b/library/flux_models.py @@ -1004,7 +1004,8 @@ class Flux(nn.Module): if self.blocks_to_swap is None or self.blocks_to_swap == 0: return self.offloader_double.prepare_block_devices_before_forward(self.double_blocks) - self.offloader_single.prepare_block_devices_before_forward(self.single_blocks) + self.offloader_single.prepare_block_devices_before_forward(self.single_blocks) + def forward( self, img: Tensor, From c8af252a44a7dbc54a0c1622946faedef4e7c52b Mon Sep 17 00:00:00 2001 From: Robert Date: Tue, 22 Apr 2025 16:19:14 +0800 Subject: [PATCH 26/40] refactor --- library/deepspeed_utils.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/library/deepspeed_utils.py b/library/deepspeed_utils.py index 3018def7..f6eac367 100644 --- a/library/deepspeed_utils.py +++ b/library/deepspeed_utils.py @@ -126,7 +126,7 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models): self.models = torch.nn.ModuleDict() - warp_model_forward_with_torch_autocast = args.mixed_precision is not "no" + wrap_model_forward_with_torch_autocast = args.mixed_precision is not "no" for key, model in kw_models.items(): if isinstance(model, list): @@ -136,31 +136,30 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models): model, torch.nn.Module ), f"model must be an instance of torch.nn.Module, but got {key} is {type(model)}" - if warp_model_forward_with_torch_autocast: - model = self.__warp_with_torch_autocast(model) + if wrap_model_forward_with_torch_autocast: + model = self.__wrap_model_with_torch_autocast(model) self.models.update(torch.nn.ModuleDict({key: model})) - def __warp_with_torch_autocast(self, model): + def __wrap_model_with_torch_autocast(self, model): if isinstance(model, torch.nn.ModuleList): - for i in range(len(model)): - model[i] = self.__warp_model_forward_with_torch_autocast(model[i]) + model = [self.__wrap_model_forward_with_torch_autocast(m) for m in model] else: - model = self.__warp_model_forward_with_torch_autocast(model) + model = self.__wrap_model_forward_with_torch_autocast(model) return model - def __warp_model_forward_with_torch_autocast(self, model): + def __wrap_model_forward_with_torch_autocast(self, model): assert hasattr(model, "forward"), f"model must have a forward method." forward_fn = model.forward def forward(*args, **kwargs): - device_type= "cuda" if torch.cuda.is_available() else "cpu" + device_type = "cuda" if torch.cuda.is_available() else "cpu" with torch.autocast(device_type=device_type): return forward_fn(*args, **kwargs) + model.forward = forward - return model def get_models(self): From adb775c6165d93a856e33d0d9058efd629cf2a2d Mon Sep 17 00:00:00 2001 From: saibit Date: Wed, 23 Apr 2025 17:05:20 +0800 Subject: [PATCH 27/40] Update: requirement diffusers[torch]==0.25.0 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index bead3f90..9e97eed3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ accelerate==0.33.0 transformers==4.44.0 -diffusers==0.25.0 +diffusers[torch]==0.25.0 deepspeed==0.16.7 ftfy==6.1.1 # albumentations==1.3.0 From abf2c44bc5650afef8bebbb1ef278c66f44c4dda Mon Sep 17 00:00:00 2001 From: sharlynxy Date: Wed, 23 Apr 2025 18:57:19 +0800 Subject: [PATCH 28/40] Dynamically set device in deepspeed wrapper (#2) * get device type from model * add logger warning * format * format * format --- library/deepspeed_utils.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/library/deepspeed_utils.py b/library/deepspeed_utils.py index f6eac367..09c6f7b9 100644 --- a/library/deepspeed_utils.py +++ b/library/deepspeed_utils.py @@ -5,6 +5,8 @@ from accelerate import DeepSpeedPlugin, Accelerator from .utils import setup_logging +from .device_utils import get_preferred_device + setup_logging() import logging @@ -153,13 +155,21 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models): assert hasattr(model, "forward"), f"model must have a forward method." forward_fn = model.forward - + def forward(*args, **kwargs): - device_type = "cuda" if torch.cuda.is_available() else "cpu" - with torch.autocast(device_type=device_type): + try: + device_type = model.device.type + except AttributeError: + logger.warning( + "[DeepSpeed] model.device is not available. Using get_preferred_device() " + "to determine the device_type for torch.autocast()." + ) + device_type = get_preferred_device().type + + with torch.autocast(device_type = device_type): return forward_fn(*args, **kwargs) - model.forward = forward + model.forward = forward return model def get_models(self): From 46ad3be0593df1df9d485c3ac2efb5aebd87730c Mon Sep 17 00:00:00 2001 From: saibit Date: Thu, 24 Apr 2025 11:26:36 +0800 Subject: [PATCH 29/40] update deepspeed wrapper --- library/deepspeed_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/library/deepspeed_utils.py b/library/deepspeed_utils.py index 09c6f7b9..a8a05c3a 100644 --- a/library/deepspeed_utils.py +++ b/library/deepspeed_utils.py @@ -134,18 +134,18 @@ def prepare_deepspeed_model(args: argparse.Namespace, **models): if isinstance(model, list): model = torch.nn.ModuleList(model) + if wrap_model_forward_with_torch_autocast: + model = self.__wrap_model_with_torch_autocast(model) + assert isinstance( model, torch.nn.Module ), f"model must be an instance of torch.nn.Module, but got {key} is {type(model)}" - if wrap_model_forward_with_torch_autocast: - model = self.__wrap_model_with_torch_autocast(model) - self.models.update(torch.nn.ModuleDict({key: model})) def __wrap_model_with_torch_autocast(self, model): if isinstance(model, torch.nn.ModuleList): - model = [self.__wrap_model_forward_with_torch_autocast(m) for m in model] + model = torch.nn.ModuleList([self.__wrap_model_forward_with_torch_autocast(m) for m in model]) else: model = self.__wrap_model_forward_with_torch_autocast(model) return model From 8387e0b95c1067e919f91a2abec11ddcd5ed15cb Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 27 Apr 2025 18:25:59 +0900 Subject: [PATCH 30/40] docs: update README to include CFG scale support in FLUX.1 training --- README.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 2e80a697..f9831aee 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,10 @@ The command to install PyTorch is as follows: ### Recent Updates +Apr 27, 2025: +- FLUX.1 training now supports CFG scale in the sample generation during training. Please use `--g` option, to specify the CFG scale (note that `--l` is used as the embedded guidance scale.) PR [#2064](https://github.com/kohya-ss/sd-scripts/pull/2064). + - See [here](#sample-image-generation-during-training) for details. + Apr 6, 2025: - IP noise gamma has been enabled in FLUX.1. Thanks to rockerBOO for PR [#1992](https://github.com/kohya-ss/sd-scripts/pull/1992). See the PR for details. - `--ip_noise_gamma` and `--ip_noise_gamma_random_strength` are available. @@ -1344,11 +1348,13 @@ masterpiece, best quality, 1boy, in business suit, standing at street, looking b Lines beginning with `#` are comments. You can specify options for the generated image with options like `--n` after the prompt. The following can be used. - * `--n` Negative prompt up to the next option. + * `--n` Negative prompt up to the next option. Ignored when CFG scale is `1.0`. * `--w` Specifies the width of the generated image. * `--h` Specifies the height of the generated image. * `--d` Specifies the seed of the generated image. * `--l` Specifies the CFG scale of the generated image. + * In guidance distillation models like FLUX.1, this value is used as the embedded guidance scale for backward compatibility. + * `--g` Specifies the CFG scale for the models with embedded guidance scale. The default is `1.0`, `1.0` means no CFG. In general, should not be changed unless you train the un-distilled FLUX.1 models. * `--s` Specifies the number of steps in the generation. The prompt weighting such as `( )` and `[ ]` are working. From f0b07c52abaf4ab33d619b427afabe17b69b7d05 Mon Sep 17 00:00:00 2001 From: "Kohya S." <52813779+kohya-ss@users.noreply.github.com> Date: Sun, 27 Apr 2025 21:28:38 +0900 Subject: [PATCH 31/40] Create FUNDING.yml --- .github/FUNDING.yml | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 .github/FUNDING.yml diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 00000000..3b8943c3 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1,3 @@ +# These are supported funding model platforms + +github: kohya-ss From fd3a445769910ddc0c8c02d13e535cac37b85d2e Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 27 Apr 2025 22:50:27 +0900 Subject: [PATCH 32/40] fix: revert default emb guidance scale and CFG scale for FLUX.1 sampling --- library/flux_train_utils.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py index d2ff347d..5f6867a8 100644 --- a/library/flux_train_utils.py +++ b/library/flux_train_utils.py @@ -154,8 +154,9 @@ def sample_image_inference( sample_steps = prompt_dict.get("sample_steps", 20) width = prompt_dict.get("width", 512) height = prompt_dict.get("height", 512) - guidance_scale = prompt_dict.get("guidance_scale", args.guidance_scale) - scale = prompt_dict.get("scale", 1.0) # 1.0 means no guidance + # TODO refactor variable names + cfg_scale = prompt_dict.get("guidance_scale", 1.0) + emb_guidance_scale = prompt_dict.get("scale", 3.5) seed = prompt_dict.get("seed") controlnet_image = prompt_dict.get("controlnet_image") prompt: str = prompt_dict.get("prompt", "") @@ -179,16 +180,16 @@ def sample_image_inference( height = max(64, height - height % 16) # round to divisible by 16 width = max(64, width - width % 16) # round to divisible by 16 logger.info(f"prompt: {prompt}") - if scale != 1.0: + if cfg_scale != 1.0: logger.info(f"negative_prompt: {negative_prompt}") elif negative_prompt != "": logger.info(f"negative prompt is ignored because scale is 1.0") logger.info(f"height: {height}") logger.info(f"width: {width}") logger.info(f"sample_steps: {sample_steps}") - logger.info(f"guidance_scale: {guidance_scale}") - if scale != 1.0: - logger.info(f"scale: {scale}") + logger.info(f"embedded guidance scale: {emb_guidance_scale}") + if cfg_scale != 1.0: + logger.info(f"CFG scale: {cfg_scale}") # logger.info(f"sample_sampler: {sampler_name}") if seed is not None: logger.info(f"seed: {seed}") @@ -220,12 +221,12 @@ def sample_image_inference( l_pooled, t5_out, txt_ids, t5_attn_mask = encode_prompt(prompt) # encode negative prompts - if scale != 1.0: + if cfg_scale != 1.0: neg_l_pooled, neg_t5_out, _, neg_t5_attn_mask = encode_prompt(negative_prompt) neg_t5_attn_mask = ( neg_t5_attn_mask.to(accelerator.device) if args.apply_t5_attn_mask and neg_t5_attn_mask is not None else None ) - neg_cond = (scale, neg_l_pooled, neg_t5_out, neg_t5_attn_mask) + neg_cond = (cfg_scale, neg_l_pooled, neg_t5_out, neg_t5_attn_mask) else: neg_cond = None @@ -260,7 +261,7 @@ def sample_image_inference( txt_ids, l_pooled, timesteps=timesteps, - guidance=guidance_scale, + guidance=emb_guidance_scale, t5_attn_mask=t5_attn_mask, controlnet=controlnet, controlnet_img=controlnet_image, From 29523c9b68bd56cdb1cce3f4985f2e45cefb1f2b Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 27 Apr 2025 23:34:37 +0900 Subject: [PATCH 33/40] docs: add note for user feedback on CFG scale in FLUX.1 training --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index f9831aee..18e8e659 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ The command to install PyTorch is as follows: Apr 27, 2025: - FLUX.1 training now supports CFG scale in the sample generation during training. Please use `--g` option, to specify the CFG scale (note that `--l` is used as the embedded guidance scale.) PR [#2064](https://github.com/kohya-ss/sd-scripts/pull/2064). - See [here](#sample-image-generation-during-training) for details. + - If you have any issues with this, please let us know. Apr 6, 2025: - IP noise gamma has been enabled in FLUX.1. Thanks to rockerBOO for PR [#1992](https://github.com/kohya-ss/sd-scripts/pull/1992). See the PR for details. From 4625b34f4ebcd30ffb24f1b03ece8a8362c7bfb2 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Tue, 29 Apr 2025 21:27:04 +0900 Subject: [PATCH 34/40] Fix mean image aspect ratio error calculation to avoid NaN values --- library/train_util.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 100ef475..b9d08f25 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -957,8 +957,11 @@ class BaseDataset(torch.utils.data.Dataset): self.bucket_info["buckets"][i] = {"resolution": reso, "count": len(bucket)} logger.info(f"bucket {i}: resolution {reso}, count: {len(bucket)}") - img_ar_errors = np.array(img_ar_errors) - mean_img_ar_error = np.mean(np.abs(img_ar_errors)) + if len(img_ar_errors) == 0: + mean_img_ar_error = 0 # avoid NaN + else: + img_ar_errors = np.array(img_ar_errors) + mean_img_ar_error = np.mean(np.abs(img_ar_errors)) self.bucket_info["mean_img_ar_error"] = mean_img_ar_error logger.info(f"mean ar error (without repeats): {mean_img_ar_error}") From 1684ababcd7fc4259c77f1471ef41d10e612a721 Mon Sep 17 00:00:00 2001 From: sharlynxy Date: Wed, 30 Apr 2025 19:51:09 +0800 Subject: [PATCH 35/40] remove deepspeed from requirements.txt --- requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 9e97eed3..767d9e8e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,6 @@ accelerate==0.33.0 transformers==4.44.0 diffusers[torch]==0.25.0 -deepspeed==0.16.7 ftfy==6.1.1 # albumentations==1.3.0 opencv-python==4.8.1.78 From a4fae93dce5b78a0c92ee328d6b2dd96be944a7d Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 1 May 2025 00:55:10 -0400 Subject: [PATCH 36/40] Add pythonpath to pytest.ini --- pytest.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/pytest.ini b/pytest.ini index 484d3aef..34b7e9c1 100644 --- a/pytest.ini +++ b/pytest.ini @@ -6,3 +6,4 @@ filterwarnings = ignore::DeprecationWarning ignore::UserWarning ignore::FutureWarning +pythonpath = . From f62c68df3c96639e83dcb7f5062d48c3067055e1 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 1 May 2025 01:37:57 -0400 Subject: [PATCH 37/40] Make grad_norm and combined_grad_norm None is not recording --- networks/lora_flux.py | 12 ++++++------ train_network.py | 6 ++++-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/networks/lora_flux.py b/networks/lora_flux.py index 92b3979a..0b30f1b8 100644 --- a/networks/lora_flux.py +++ b/networks/lora_flux.py @@ -955,26 +955,26 @@ class LoRANetwork(torch.nn.Module): for lora in self.text_encoder_loras + self.unet_loras: lora.update_grad_norms() - def grad_norms(self) -> Tensor: + def grad_norms(self) -> Tensor | None: grad_norms = [] for lora in self.text_encoder_loras + self.unet_loras: if hasattr(lora, "grad_norms") and lora.grad_norms is not None: grad_norms.append(lora.grad_norms.mean(dim=0)) - return torch.stack(grad_norms) if len(grad_norms) > 0 else torch.tensor([]) + return torch.stack(grad_norms) if len(grad_norms) > 0 else None - def weight_norms(self) -> Tensor: + def weight_norms(self) -> Tensor | None: weight_norms = [] for lora in self.text_encoder_loras + self.unet_loras: if hasattr(lora, "weight_norms") and lora.weight_norms is not None: weight_norms.append(lora.weight_norms.mean(dim=0)) - return torch.stack(weight_norms) if len(weight_norms) > 0 else torch.tensor([]) + return torch.stack(weight_norms) if len(weight_norms) > 0 else None - def combined_weight_norms(self) -> Tensor: + def combined_weight_norms(self) -> Tensor | None: combined_weight_norms = [] for lora in self.text_encoder_loras + self.unet_loras: if hasattr(lora, "combined_weight_norms") and lora.combined_weight_norms is not None: combined_weight_norms.append(lora.combined_weight_norms.mean(dim=0)) - return torch.stack(combined_weight_norms) if len(combined_weight_norms) > 0 else torch.tensor([]) + return torch.stack(combined_weight_norms) if len(combined_weight_norms) > 0 else None def load_weights(self, file): diff --git a/train_network.py b/train_network.py index d6bc66ed..2b4e6d3f 100644 --- a/train_network.py +++ b/train_network.py @@ -1444,8 +1444,10 @@ class NetworkTrainer: else: if hasattr(network, "weight_norms"): mean_norm = network.weight_norms().mean().item() - mean_grad_norm = network.grad_norms().mean().item() - mean_combined_norm = network.combined_weight_norms().mean().item() + grad_norms = network.grad_norms() + mean_grad_norm = grad_norms.mean().item() if grad_norms is not None else None + combined_weight_norms = network.combined_weight_norms() + mean_combined_norm = combined_weight_norms.mean().item() if combined_weight_norms is not None else None weight_norms = network.weight_norms() maximum_norm = weight_norms.max().item() if weight_norms.numel() > 0 else None keys_scaled = None From b4a89c3cdf7319b6840f1e4a28a5a1001643bc22 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Thu, 1 May 2025 02:03:22 -0400 Subject: [PATCH 38/40] Fix None --- train_network.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train_network.py b/train_network.py index 2b4e6d3f..1336a0b1 100644 --- a/train_network.py +++ b/train_network.py @@ -1443,13 +1443,13 @@ class NetworkTrainer: max_mean_logs = {"Keys Scaled": keys_scaled, "Average key norm": mean_norm} else: if hasattr(network, "weight_norms"): - mean_norm = network.weight_norms().mean().item() + weight_norms = network.weight_norms() + mean_norm = weight_norms.mean().item() if weight_norms is not None else None grad_norms = network.grad_norms() mean_grad_norm = grad_norms.mean().item() if grad_norms is not None else None combined_weight_norms = network.combined_weight_norms() mean_combined_norm = combined_weight_norms.mean().item() if combined_weight_norms is not None else None - weight_norms = network.weight_norms() - maximum_norm = weight_norms.max().item() if weight_norms.numel() > 0 else None + maximum_norm = weight_norms.max().item() if weight_norms is not None else None keys_scaled = None max_mean_logs = {} else: From 865c8d55e2b8cd9f0b6008a6d4ee4a07949d9acc Mon Sep 17 00:00:00 2001 From: Kohya S Date: Thu, 1 May 2025 23:29:19 +0900 Subject: [PATCH 39/40] README.md: Update recent updates and add DeepSpeed installation instructions --- README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/README.md b/README.md index 18e8e659..13c2320c 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,10 @@ The command to install PyTorch is as follows: ### Recent Updates +May 1, 2025: +- The error when training FLUX.1 with mixed precision in flux_train.py with DeepSpeed enabled has been resolved. Thanks to sharlynxy for PR [#2060](https://github.com/kohya-ss/sd-scripts/pull/2060). Please refer to the PR for details. + - If you enable DeepSpeed, please install deepseed with `pip install deepspeed==0.16.7`. + Apr 27, 2025: - FLUX.1 training now supports CFG scale in the sample generation during training. Please use `--g` option, to specify the CFG scale (note that `--l` is used as the embedded guidance scale.) PR [#2064](https://github.com/kohya-ss/sd-scripts/pull/2064). - See [here](#sample-image-generation-during-training) for details. @@ -875,6 +879,14 @@ Note: Some user reports ``ValueError: fp16 mixed precision requires a GPU`` is o (Single GPU with id `0` will be used.) +## DeepSpeed installation (experimental, Linux or WSL2 only) + +To install DeepSpeed, run the following command in your activated virtual environment: + +```bash +pip install deepspeed==0.16.7 +``` + ## Upgrade When a new release comes out you can upgrade your repo with the following command: From a27ace74d96d9519629283f4ff3d207c1ad8d98e Mon Sep 17 00:00:00 2001 From: Kohya S Date: Thu, 1 May 2025 23:31:23 +0900 Subject: [PATCH 40/40] doc: add DeepSpeed installation in header section --- README.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 13c2320c..497969ab 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,8 @@ __Please update PyTorch to 2.4.0. We have tested with `torch==2.4.0` and `torchv The command to install PyTorch is as follows: `pip3 install torch==2.4.0 torchvision==0.19.0 --index-url https://download.pytorch.org/whl/cu124` +If you are using DeepSpeed, please install DeepSpeed with `pip install deepspeed==0.16.7`. + - [FLUX.1 training](#flux1-training) - [SD3 training](#sd3-training) @@ -16,7 +18,7 @@ The command to install PyTorch is as follows: May 1, 2025: - The error when training FLUX.1 with mixed precision in flux_train.py with DeepSpeed enabled has been resolved. Thanks to sharlynxy for PR [#2060](https://github.com/kohya-ss/sd-scripts/pull/2060). Please refer to the PR for details. - - If you enable DeepSpeed, please install deepseed with `pip install deepspeed==0.16.7`. + - If you enable DeepSpeed, please install DeepSpeed with `pip install deepspeed==0.16.7`. Apr 27, 2025: - FLUX.1 training now supports CFG scale in the sample generation during training. Please use `--g` option, to specify the CFG scale (note that `--l` is used as the embedded guidance scale.) PR [#2064](https://github.com/kohya-ss/sd-scripts/pull/2064).