From a10635818ab5633021e6ba1a71f16c7d5a048e86 Mon Sep 17 00:00:00 2001 From: "xuyixuan.xyx" Date: Tue, 22 Apr 2025 10:40:40 +0800 Subject: [PATCH] multi-node --- test.py | 26 --- train_flux_reference.py | 15 +- train_flux_reference_multi_node.py | 248 +++++++++++++++++++++++++++++ 3 files changed, 258 insertions(+), 31 deletions(-) delete mode 100644 test.py create mode 100644 train_flux_reference_multi_node.py diff --git a/test.py b/test.py deleted file mode 100644 index da929c5..0000000 --- a/test.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch -from diffsynth import ModelManager, FluxImagePipeline, download_models, load_state_dict -from diffsynth.models.flux_reference_embedder import FluxReferenceEmbedder -from PIL import Image - - -model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cuda") -model_manager.load_models([ - "models/FLUX/FLUX.1-dev/text_encoder/model.safetensors", - "models/FLUX/FLUX.1-dev/text_encoder_2", - "models/FLUX/FLUX.1-dev/ae.safetensors", - "models/FLUX/FLUX.1-dev/flux1-dev.safetensors" -]) -pipe = FluxImagePipeline.from_model_manager(model_manager) - -pipe.reference_embedder = FluxReferenceEmbedder().to(dtype=torch.bfloat16, device="cuda") -pipe.reference_embedder.init() - -for i in range(4): - image = pipe( - prompt="a girl.", - num_inference_steps=30, embedded_guidance=3.5, - height=512, width=512, - reference_images=[Image.open("data/example4.jpg").resize((512, 512))] - ) - image.save(f"image_{i}.jpg") \ No newline at end of file diff --git a/train_flux_reference.py b/train_flux_reference.py index 3d626af..bc08564 100644 --- a/train_flux_reference.py +++ b/train_flux_reference.py @@ -42,6 +42,8 @@ class LightningModel(LightningModelForT2ILoRA): self.freeze_parameters() self.pipe.reference_embedder.requires_grad_(True) self.pipe.reference_embedder.train() + self.pipe.dit.requires_grad_(True) + self.pipe.dit.train() # self.add_lora_to_model( # self.pipe.denoising_model(), # lora_rank=lora_rank, @@ -192,27 +194,30 @@ if __name__ == '__main__': dataset_list=[ SingleTaskDataset( "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_change_add_remove", - metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250411_dataset_change_add_remove.json", + keys=(("image_1", "image_2", "editing_instruction"), ("image_2", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_change_add_remove.json", height=512, width=512, ), SingleTaskDataset( "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_zoomin_zoomout", - metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250411_dataset_zoomin_zoomout.json", + keys=(("image_1", "image_2", "editing_instruction"), ("image_2", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_zoomin_zoomout.json", height=512, width=512, ), SingleTaskDataset( "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_style_transfer", keys=(("image_1", "image_4", "editing_instruction"), ("image_4", "image_1", "reverse_editing_instruction")), - metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250411_dataset_style_transfer.json", + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_style_transfer.json", height=512, width=512, ), SingleTaskDataset( "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_faceid", - metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250411_dataset_faceid.json", + keys=(("image_1", "image_2", "editing_instruction"), ("image_2", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_faceid.json", height=512, width=512, ), ], - dataset_weight=(4, 2, 2, 1), + dataset_weight=(4, 1, 4, 1), steps_per_epoch=args.steps_per_epoch, ) train_loader = torch.utils.data.DataLoader( diff --git a/train_flux_reference_multi_node.py b/train_flux_reference_multi_node.py new file mode 100644 index 0000000..cfe62d8 --- /dev/null +++ b/train_flux_reference_multi_node.py @@ -0,0 +1,248 @@ +from diffsynth import ModelManager, FluxImagePipeline +from diffsynth.trainers.text_to_image import LightningModelForT2ILoRA, add_general_parsers, launch_training_task +from diffsynth.models.lora import FluxLoRAConverter +import torch, os, argparse +import lightning as pl +from diffsynth.data.image_pulse import SingleTaskDataset, MultiTaskDataset +from diffsynth.pipelines.flux_image import lets_dance_flux +from diffsynth.models.flux_reference_embedder import FluxReferenceEmbedder +os.environ["TOKENIZERS_PARALLELISM"] = "True" + + +class LightningModel(LightningModelForT2ILoRA): + def __init__( + self, + torch_dtype=torch.float16, pretrained_weights=[], preset_lora_path=None, + learning_rate=1e-4, use_gradient_checkpointing=True, + lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out", init_lora_weights="kaiming", pretrained_lora_path=None, + state_dict_converter=None, quantize = None + ): + super().__init__(learning_rate=learning_rate, use_gradient_checkpointing=use_gradient_checkpointing, state_dict_converter=state_dict_converter) + # Load models + model_manager = ModelManager(torch_dtype=torch_dtype, device=self.device) + if quantize is None: + model_manager.load_models(pretrained_weights) + else: + model_manager.load_models(pretrained_weights[1:]) + model_manager.load_model(pretrained_weights[0], torch_dtype=quantize) + if preset_lora_path is not None: + preset_lora_path = preset_lora_path.split(",") + for path in preset_lora_path: + model_manager.load_lora(path) + + self.pipe = FluxImagePipeline.from_model_manager(model_manager) + self.pipe.reference_embedder = FluxReferenceEmbedder() + self.pipe.reference_embedder.init() + + if quantize is not None: + self.pipe.dit.quantize() + + self.pipe.scheduler.set_timesteps(1000, training=True) + + self.freeze_parameters() + self.pipe.reference_embedder.requires_grad_(True) + self.pipe.reference_embedder.train() + self.pipe.dit.requires_grad_(True) + self.pipe.dit.train() + # self.add_lora_to_model( + # self.pipe.denoising_model(), + # lora_rank=lora_rank, + # lora_alpha=lora_alpha, + # lora_target_modules=lora_target_modules, + # init_lora_weights=init_lora_weights, + # pretrained_lora_path=pretrained_lora_path, + # state_dict_converter=FluxLoRAConverter.align_to_diffsynth_format + # ) + + + def training_step(self, batch, batch_idx): + # Data + text, image = batch["instruction"], batch["image_2"] + image_ref = batch["image_1"] + + # Prepare input parameters + self.pipe.device = self.device + prompt_emb = self.pipe.encode_prompt(text, positive=True) + if "latents" in batch: + latents = batch["latents"].to(dtype=self.pipe.torch_dtype, device=self.device) + else: + latents = self.pipe.vae_encoder(image.to(dtype=self.pipe.torch_dtype, device=self.device)) + noise = torch.randn_like(latents) + timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,)) + timestep = self.pipe.scheduler.timesteps[timestep_id].to(self.device) + extra_input = self.pipe.prepare_extra_input(latents) + noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep) + training_target = self.pipe.scheduler.training_target(latents, noise, timestep) + + # Reference image + hidden_states_ref = self.pipe.vae_encoder(image_ref.to(dtype=self.pipe.torch_dtype, device=self.device)) + + # Compute loss + noise_pred = lets_dance_flux( + self.pipe.denoising_model(), + reference_embedder=self.pipe.reference_embedder, + hidden_states_ref=hidden_states_ref, + hidden_states=noisy_latents, timestep=timestep, **prompt_emb, **extra_input, + use_gradient_checkpointing=self.use_gradient_checkpointing + ) + loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float()) + loss = loss * self.pipe.scheduler.training_weight(timestep) + + # Record log + self.log("train_loss", loss, prog_bar=True) + return loss + + + def configure_optimizers(self): + trainable_modules = filter(lambda p: p.requires_grad, self.pipe.parameters()) + optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate) + return optimizer + + + def on_save_checkpoint(self, checkpoint): + checkpoint.clear() + trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.pipe.named_parameters())) + trainable_param_names = set([named_param[0] for named_param in trainable_param_names]) + state_dict = self.pipe.state_dict() + lora_state_dict = {} + for name, param in state_dict.items(): + if name in trainable_param_names: + lora_state_dict[name] = param + if self.state_dict_converter is not None: + lora_state_dict = self.state_dict_converter(lora_state_dict, alpha=self.lora_alpha) + checkpoint.update(lora_state_dict) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Simple example of a training script.") + parser.add_argument( + "--pretrained_text_encoder_path", + type=str, + default=None, + required=True, + help="Path to pretrained text encoder model. For example, `models/FLUX/FLUX.1-dev/text_encoder/model.safetensors`.", + ) + parser.add_argument( + "--pretrained_text_encoder_2_path", + type=str, + default=None, + required=True, + help="Path to pretrained t5 text encoder model. For example, `models/FLUX/FLUX.1-dev/text_encoder_2`.", + ) + parser.add_argument( + "--pretrained_dit_path", + type=str, + default=None, + required=True, + help="Path to pretrained dit model. For example, `models/FLUX/FLUX.1-dev/flux1-dev.safetensors`.", + ) + parser.add_argument( + "--pretrained_vae_path", + type=str, + default=None, + required=True, + help="Path to pretrained vae model. For example, `models/FLUX/FLUX.1-dev/ae.safetensors`.", + ) + parser.add_argument( + "--lora_target_modules", + type=str, + default="a_to_qkv,b_to_qkv,ff_a.0,ff_a.2,ff_b.0,ff_b.2,a_to_out,b_to_out,proj_out,norm.linear,norm1_a.linear,norm1_b.linear,to_qkv_mlp", + help="Layers with LoRA modules.", + ) + parser.add_argument( + "--align_to_opensource_format", + default=False, + action="store_true", + help="Whether to export lora files aligned with other opensource format.", + ) + parser.add_argument( + "--quantize", + type=str, + default=None, + choices=["float8_e4m3fn"], + help="Whether to use quantization when training the model, and in which format.", + ) + parser.add_argument( + "--preset_lora_path", + type=str, + default=None, + help="Preset LoRA path.", + ) + parser.add_argument( + "--num_nodes", + type=int, + default=1, + help="Num nodes.", + ) + parser = add_general_parsers(parser) + args = parser.parse_args() + return args + + +if __name__ == '__main__': + args = parse_args() + model = LightningModel( + torch_dtype={"32": torch.float32, "bf16": torch.bfloat16}.get(args.precision, torch.float16), + pretrained_weights=[args.pretrained_dit_path, args.pretrained_text_encoder_path, args.pretrained_text_encoder_2_path, args.pretrained_vae_path], + preset_lora_path=args.preset_lora_path, + learning_rate=args.learning_rate, + use_gradient_checkpointing=args.use_gradient_checkpointing, + lora_rank=args.lora_rank, + lora_alpha=args.lora_alpha, + lora_target_modules=args.lora_target_modules, + init_lora_weights=args.init_lora_weights, + pretrained_lora_path=args.pretrained_lora_path, + state_dict_converter=FluxLoRAConverter.align_to_opensource_format if args.align_to_opensource_format else None, + quantize={"float8_e4m3fn": torch.float8_e4m3fn}.get(args.quantize, None), + ) + # dataset and data loader + dataset = MultiTaskDataset( + dataset_list=[ + SingleTaskDataset( + "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_change_add_remove", + keys=(("image_1", "image_2", "editing_instruction"), ("image_2", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_change_add_remove.json", + height=512, width=512, + ), + SingleTaskDataset( + "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_zoomin_zoomout", + keys=(("image_1", "image_2", "editing_instruction"), ("image_2", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_zoomin_zoomout.json", + height=512, width=512, + ), + SingleTaskDataset( + "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_style_transfer", + keys=(("image_1", "image_4", "editing_instruction"), ("image_4", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_style_transfer.json", + height=512, width=512, + ), + SingleTaskDataset( + "/shark/zhongjie/data/image_pulse_datasets/task1/data/dataset_faceid", + keys=(("image_1", "image_2", "editing_instruction"), ("image_2", "image_1", "reverse_editing_instruction")), + metadata_path="/shark/zhongjie/data/image_pulse_datasets/task1/data/metadata/20250418_dataset_faceid.json", + height=512, width=512, + ), + ], + dataset_weight=(4, 1, 4, 1), + steps_per_epoch=args.steps_per_epoch, + ) + train_loader = torch.utils.data.DataLoader( + dataset, + shuffle=True, + batch_size=args.batch_size, + num_workers=args.dataloader_num_workers + ) + # train + trainer = pl.Trainer( + max_epochs=args.max_epochs, + accelerator="gpu", + devices="auto", + num_nodes=args.num_nodes, + precision=args.precision, + strategy="ddp", + default_root_dir=args.output_path, + accumulate_grad_batches=args.accumulate_grad_batches, + callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)], + logger=None, + ) + trainer.fit(model=model, train_dataloaders=train_loader)