import torch, os, argparse, accelerate import copy import math from diffsynth.core import UnifiedDataset, gradient_checkpoint_forward from diffsynth.diffusion import * from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig, model_fn_flux2 os.environ["TOKENIZERS_PARALLELISM"] = "false" def _parse_int_list(value): if value is None or value == "": return None return [int(i) for i in value.split(",") if i != ""] def _parse_float_list(value): if value is None or value == "": return None return [float(i) for i in value.split(",") if i != ""] def _get_optimal_groups(num_channels): if num_channels <= 32: groups = max(1, num_channels // 4) else: groups = 32 while groups > 1 and num_channels % groups != 0: groups -= 1 assert num_channels % groups == 0, f"{num_channels} not divisible by {groups}" return groups class FluxDMD2Discriminator(torch.nn.Module): """DMD2 GAN discriminator scoring teacher hidden features at `feature_indices`. Only instantiated when `gan_loss_weight > 0`. """ def __init__(self, feature_indices=None, num_blocks=40, inner_dim=3072): super().__init__() if feature_indices is None: feature_indices = [int(num_blocks // 2)] self.feature_indices = sorted({int(i) for i in feature_indices if 0 <= int(i) < num_blocks}) if len(self.feature_indices) == 0: raise ValueError("DMD2 discriminator requires at least one valid feature index.") self.num_features = len(self.feature_indices) self.inner_dim = inner_dim hidden_channels = inner_dim // 2 self.heads = torch.nn.ModuleList([ torch.nn.Sequential( torch.nn.Conv2d(inner_dim, hidden_channels, kernel_size=4, stride=2, padding=1), torch.nn.GroupNorm(_get_optimal_groups(hidden_channels), hidden_channels), torch.nn.LeakyReLU(0.2), torch.nn.Conv2d(hidden_channels, 1, kernel_size=1, stride=1, padding=0), torch.nn.AdaptiveAvgPool2d((1, 1)), torch.nn.Flatten(), ) for _ in self.feature_indices ]) def forward(self, feats): if not isinstance(feats, list) or len(feats) != self.num_features: raise ValueError( f"Expected list of {self.num_features} feature tensors, " f"got {type(feats)} with length {len(feats) if isinstance(feats, list) else 'N/A'}." ) logits = [] for head, feat in zip(self.heads, feats): param = next(head.parameters()) feat = feat.to(device=param.device, dtype=param.dtype) logits.append(head(feat)) return torch.cat(logits, dim=1) def model_fn_flux2_features( dit, latents=None, timestep=None, embedded_guidance=None, prompt_embeds=None, text_ids=None, image_ids=None, edit_latents=None, edit_image_ids=None, kv_cache=None, extra_text_embedding=None, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, feature_indices=None, **kwargs, ): """Flux.2 DiT forward exposing hidden features at `feature_indices` for the discriminator.""" feature_indices = set() if feature_indices is None else set(feature_indices) image_seq_len = latents.shape[1] if edit_latents is not None: image_seq_len = latents.shape[1] latents = torch.concat([latents, edit_latents], dim=1) image_ids = torch.concat([image_ids, edit_image_ids], dim=1) if embedded_guidance is None: embedded_guidance = None elif isinstance(embedded_guidance, torch.Tensor): embedded_guidance = embedded_guidance.to(device=latents.device, dtype=latents.dtype).flatten() if embedded_guidance.numel() == 1: embedded_guidance = embedded_guidance.expand(latents.shape[0]) elif embedded_guidance.numel() != latents.shape[0]: raise ValueError("`embedded_guidance` must be a scalar or match the latent batch size.") else: embedded_guidance = torch.full((latents.shape[0],), float(embedded_guidance), device=latents.device, dtype=latents.dtype) if extra_text_embedding is not None: extra_text_ids = torch.zeros((1, extra_text_embedding.shape[1], 4), dtype=text_ids.dtype, device=text_ids.device) extra_text_ids[:, :, -1] = torch.arange(prompt_embeds.shape[1], prompt_embeds.shape[1] + extra_text_embedding.shape[1]) prompt_embeds = torch.concat([prompt_embeds, extra_text_embedding], dim=1) text_ids = torch.concat([text_ids, extra_text_ids], dim=1) height, width = kwargs.get("height"), kwargs.get("width") if height is not None and width is not None: feature_height, feature_width = int(height) // 16, int(width) // 16 else: feature_height = int(math.sqrt(image_seq_len)) feature_width = image_seq_len // feature_height if feature_height > 0 else 0 if feature_height * feature_width != image_seq_len: raise ValueError("Flux2 feature extraction requires height/width or square latent tokens.") features = [] def append_feature(feat): feat = feat[:, :image_seq_len] batch_size, _, channels = feat.shape feat = feat.permute(0, 2, 1).reshape(batch_size, channels, feature_height, feature_width) features.append(feat) if len(features) == len(feature_indices): return features return None num_txt_tokens = prompt_embeds.shape[1] timestep = timestep.to(latents.dtype) guidance = None if embedded_guidance is None else embedded_guidance.to(latents.dtype) * 1000 temb = dit.time_guidance_embed(timestep, guidance) double_stream_mod_img = dit.double_stream_modulation_img(temb) double_stream_mod_txt = dit.double_stream_modulation_txt(temb) single_stream_mod = dit.single_stream_modulation(temb)[0] hidden_states = dit.x_embedder(latents) encoder_hidden_states = dit.context_embedder(prompt_embeds) if image_ids.ndim == 3: image_ids = image_ids[0] if text_ids.ndim == 3: text_ids = text_ids[0] image_rotary_emb = dit.pos_embed(image_ids) text_rotary_emb = dit.pos_embed(text_ids) concat_rotary_emb = ( torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), ) for block_id, block in enumerate(dit.transformer_blocks): encoder_hidden_states, hidden_states = gradient_checkpoint_forward( block, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb_mod_params_img=double_stream_mod_img, temb_mod_params_txt=double_stream_mod_txt, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=None, kv_cache=None if kv_cache is None else kv_cache.get(f"double_{block_id}"), ) if block_id in feature_indices: selected_features = append_feature(hidden_states) if selected_features is not None: return selected_features hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) num_double_blocks = len(dit.transformer_blocks) for block_id, block in enumerate(dit.single_transformer_blocks): hidden_states = gradient_checkpoint_forward( block, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, hidden_states=hidden_states, encoder_hidden_states=None, temb_mod_params=single_stream_mod, image_rotary_emb=concat_rotary_emb, joint_attention_kwargs=None, kv_cache=None if kv_cache is None else kv_cache.get(f"single_{block_id}"), ) feature_id = block_id + num_double_blocks if feature_id in feature_indices: selected_features = append_feature(hidden_states[:, num_txt_tokens:num_txt_tokens + image_seq_len]) if selected_features is not None: return selected_features if len(features) != len(feature_indices): raise ValueError(f"Only collected {len(features)} feature maps for {len(feature_indices)} requested feature indices.") return features def model_fn_flux2_dmd2( pipe, dit, timestep, progress_id, num_inference_steps, inputs_shared, inputs_posi, feature_indices=None, return_features=False, ): """Dispatcher used by `DMD2Loss`: returns flow prediction or hidden features.""" if not return_features: return model_fn_flux2( dit=dit, **inputs_shared, **inputs_posi, timestep=timestep, progress_id=progress_id, num_inference_steps=num_inference_steps, ) return model_fn_flux2_features( dit=dit, **inputs_shared, **inputs_posi, timestep=timestep, progress_id=progress_id, num_inference_steps=num_inference_steps, feature_indices=feature_indices, ) class Flux2DMD2TrainingModule(DiffusionTrainingModule): def __init__(self, args, device="cpu"): config = DMD2Config( student_update_freq=args.dmd2_student_update_freq, student_sample_steps=args.dmd2_student_sample_steps, student_sample_type=args.dmd2_student_sample_type, student_schedule=args.dmd2_student_schedule, student_t_list=_parse_float_list(args.dmd2_student_t_list), matching_t_min=args.dmd2_matching_t_min, matching_t_max=args.dmd2_matching_t_max, matching_t_sampling=args.dmd2_matching_t_sampling, matching_t_mean=args.dmd2_matching_t_mean, matching_t_std=args.dmd2_matching_t_std, gan_loss_weight=args.dmd2_gan_loss_weight, gan_r1_reg_weight=args.dmd2_gan_r1_reg_weight, gan_r1_reg_alpha=args.dmd2_gan_r1_reg_alpha, fake_score_learning_rate=args.dmd2_fake_score_learning_rate, discriminator_learning_rate=args.dmd2_discriminator_learning_rate, feature_indices=_parse_int_list(args.dmd2_feature_indices), teacher_cfg_scale=args.dmd2_teacher_cfg_scale, student_grad_clip_norm=args.dmd2_student_grad_clip_norm, ) super().__init__() # Load models model_configs = self.parse_model_configs(args.model_paths, args.model_id_with_origin_paths, fp8_models=args.fp8_models, offload_models=args.offload_models, device=device) tokenizer_config = self.parse_path_or_model_id(args.tokenizer_path,default_value=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/")) self.pipe = Flux2ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device,model_configs=model_configs, tokenizer_config=tokenizer_config) # Training mode self.switch_pipe_to_training_mode( self.pipe, args.trainable_models, args.lora_base_model, args.lora_target_modules, args.lora_rank, args.lora_checkpoint, args.preset_lora_path, args.preset_lora_model, ) # DMD2 auxiliary models (frozen teacher, trainable fake_score, optional discriminator) self.pipe.dit_teacher = copy.deepcopy(self.pipe.dit) self.pipe.dit_fake_score = copy.deepcopy(self.pipe.dit) self.pipe.dmd2_discriminator = None if config.gan_loss_weight > 0: self.pipe.dmd2_discriminator = FluxDMD2Discriminator(feature_indices=config.feature_indices) self.pipe.dit_teacher.eval().requires_grad_(False) self.pipe.dit_fake_score.train().requires_grad_(True) if self.pipe.dmd2_discriminator is not None: self.pipe.dmd2_discriminator.train().requires_grad_(True) self.resume_from_checkpoint(args.resume_from_checkpoint, args.remove_prefix_in_ckpt) # Other configs self.use_gradient_checkpointing = args.use_gradient_checkpointing self.use_gradient_checkpointing_offload = args.use_gradient_checkpointing_offload self.extra_inputs = args.extra_inputs.split(",") if args.extra_inputs is not None else [] self.fp8_models = args.fp8_models self.embedded_guidance = args.embedded_guidance self.dmd2_student_model_name = "dit" self.dmd2_teacher_model_name = "dit_teacher" self.dmd2_fake_score_model_name = "dit_fake_score" self.dmd2_discriminator_model_name = "dmd2_discriminator" self.dmd2_model_fn_student = model_fn_flux2_dmd2 self.dmd2_model_fn_teacher = model_fn_flux2_dmd2 self.dmd2_model_fn_fake_score = model_fn_flux2_dmd2 self.dmd2_config = config self.loss = DMD2Loss(config) self._dmd2_student_param_names = {name for name, param in self.pipe.dit.named_parameters() if param.requires_grad} self._dmd2_fake_score_param_names = {name for name, _ in self.pipe.dit_fake_score.named_parameters()} def get_pipeline_inputs(self, data): inputs_posi = {"prompt": data["prompt"]} inputs_nega = {"negative_prompt": ""} inputs_shared = { "input_image": data["image"], "height": data["image"].size[1], "width": data["image"].size[0], "embedded_guidance": self.embedded_guidance, "cfg_scale": self.dmd2_config.teacher_cfg_scale, "rand_device": self.pipe.device, "use_gradient_checkpointing": self.use_gradient_checkpointing, "use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload, } inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared) return inputs_shared, inputs_posi, inputs_nega def forward(self, data, inputs=None, iteration=None): if inputs is None: inputs = self.get_pipeline_inputs(data) inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype) for unit in self.pipe.units: inputs = self.pipe.unit_runner(unit, self.pipe, *inputs) return self.loss(self, inputs, 0 if iteration is None else iteration) def export_trainable_state_dict(self, state_dict, remove_prefix=None): return export_dmd2_trainable_state_dict(self, state_dict, remove_prefix=remove_prefix) def flux2_dmd2_parser(): parser = argparse.ArgumentParser(description="Flux.2 DMD2 training script.") parser = add_general_config(parser) parser = add_image_size_config(parser) parser = add_dmd2_config(parser) parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.") parser.add_argument("--embedded_guidance", type=float, default=1.0, help="Flux.2 embedded guidance value.") parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.") return parser if __name__ == "__main__": parser = flux2_dmd2_parser() args = parser.parse_args() args.find_unused_parameters = True accelerator = accelerate.Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)], ) dataset = UnifiedDataset( base_path=args.dataset_base_path, metadata_path=args.dataset_metadata_path, repeat=args.dataset_repeat, data_file_keys=args.data_file_keys.split(","), main_data_operator=UnifiedDataset.default_image_operator( base_path=args.dataset_base_path, max_pixels=args.max_pixels, height=args.height, width=args.width, height_division_factor=16, width_division_factor=16, ), ) model = Flux2DMD2TrainingModule( args, device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device, ) model_logger = ModelLogger( args.output_path, remove_prefix_in_ckpt=args.remove_prefix_in_ckpt, enable_tensorboard_log=args.enable_tensorboard_log, enable_swanlab_log=args.enable_swanlab_log, swanlab_project=args.swanlab_project, enable_wandb_log=args.enable_wandb_log, wandb_project=args.wandb_project, ) launch_dmd2_training_task(accelerator, dataset, model, model_logger, args=args)