hugging-apps/echo-memory

1import lightning as pl2from peft import LoraConfig, inject_adapter_in_model3import torch, os4from ..data.simple_text_image import TextImageDataset5from modelscope.hub.api import HubApi6from ..models.utils import load_state_dict7 8 9 10class LightningModelForT2ILoRA(pl.LightningModule):11    def __init__(12        self,13        learning_rate=1e-4,14        use_gradient_checkpointing=True,15        state_dict_converter=None,16    ):17        super().__init__()18        # Set parameters19        self.learning_rate = learning_rate20        self.use_gradient_checkpointing = use_gradient_checkpointing21        self.state_dict_converter = state_dict_converter22        self.lora_alpha = None23 24 25    def load_models(self):26        # This function is implemented in other modules27        self.pipe = None28 29 30    def freeze_parameters(self):31        # Freeze parameters32        self.pipe.requires_grad_(False)33        self.pipe.eval()34        self.pipe.denoising_model().train()35 36    37    def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out", init_lora_weights="gaussian", pretrained_lora_path=None, state_dict_converter=None):38        # Add LoRA to UNet39        self.lora_alpha = lora_alpha40        if init_lora_weights == "kaiming":41            init_lora_weights = True42            43        lora_config = LoraConfig(44            r=lora_rank,45            lora_alpha=lora_alpha,46            init_lora_weights=init_lora_weights,47            target_modules=lora_target_modules.split(","),48        )49        model = inject_adapter_in_model(lora_config, model)50        for param in model.parameters():51            # Upcast LoRA parameters into fp3252            if param.requires_grad:53                param.data = param.to(torch.float32)54 55        # Lora pretrained lora weights56        if pretrained_lora_path is not None:57            state_dict = load_state_dict(pretrained_lora_path)58            if state_dict_converter is not None:59                state_dict = state_dict_converter(state_dict)60            missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)61            all_keys = [i for i, _ in model.named_parameters()]62            num_updated_keys = len(all_keys) - len(missing_keys)63            num_unexpected_keys = len(unexpected_keys)64            print(f"{num_updated_keys} parameters are loaded from {pretrained_lora_path}. {num_unexpected_keys} parameters are unexpected.")65 66 67    def training_step(self, batch, batch_idx):68        # Data69        text, image = batch["text"], batch["image"]70 71        # Prepare input parameters72        self.pipe.device = self.device73        prompt_emb = self.pipe.encode_prompt(text, positive=True)74        if "latents" in batch:75            latents = batch["latents"].to(dtype=self.pipe.torch_dtype, device=self.device)76        else:77            latents = self.pipe.vae_encoder(image.to(dtype=self.pipe.torch_dtype, device=self.device))78        noise = torch.randn_like(latents)79        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))80        timestep = self.pipe.scheduler.timesteps[timestep_id].to(self.device)81        extra_input = self.pipe.prepare_extra_input(latents)82        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)83        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)84 85        # Compute loss86        noise_pred = self.pipe.denoising_model()(87            noisy_latents, timestep=timestep, **prompt_emb, **extra_input,88            use_gradient_checkpointing=self.use_gradient_checkpointing89        )90        loss = torch.nn.functional.mse_loss(noise_pred.float(), training_target.float())91        loss = loss * self.pipe.scheduler.training_weight(timestep)92 93        # Record log94        self.log("train_loss", loss, prog_bar=True)95        return loss96 97 98    def configure_optimizers(self):99        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.denoising_model().parameters())100        optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate)101        return optimizer102    103 104    def on_save_checkpoint(self, checkpoint):105        checkpoint.clear()106        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.pipe.denoising_model().named_parameters()))107        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])108        state_dict = self.pipe.denoising_model().state_dict()109        lora_state_dict = {}110        for name, param in state_dict.items():111            if name in trainable_param_names:112                lora_state_dict[name] = param113        if self.state_dict_converter is not None:114            lora_state_dict = self.state_dict_converter(lora_state_dict, alpha=self.lora_alpha)115        checkpoint.update(lora_state_dict)116 117 118 119def add_general_parsers(parser):120    parser.add_argument(121        "--dataset_path",122        type=str,123        default=None,124        required=True,125        help="The path of the Dataset.",126    )127    parser.add_argument(128        "--output_path",129        type=str,130        default="./",131        help="Path to save the model.",132    )133    parser.add_argument(134        "--steps_per_epoch",135        type=int,136        default=500,137        help="Number of steps per epoch.",138    )139    parser.add_argument(140        "--height",141        type=int,142        default=1024,143        help="Image height.",144    )145    parser.add_argument(146        "--width",147        type=int,148        default=1024,149        help="Image width.",150    )151    parser.add_argument(152        "--center_crop",153        default=False,154        action="store_true",155        help=(156            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"157            " cropped. The images will be resized to the resolution first before cropping."158        ),159    )160    parser.add_argument(161        "--random_flip",162        default=False,163        action="store_true",164        help="Whether to randomly flip images horizontally",165    )166    parser.add_argument(167        "--batch_size",168        type=int,169        default=1,170        help="Batch size (per device) for the training dataloader.",171    )172    parser.add_argument(173        "--dataloader_num_workers",174        type=int,175        default=0,176        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",177    )178    parser.add_argument(179        "--precision",180        type=str,181        default="16-mixed",182        choices=["32", "16", "16-mixed", "bf16"],183        help="Training precision",184    )185    parser.add_argument(186        "--learning_rate",187        type=float,188        default=1e-4,189        help="Learning rate.",190    )191    parser.add_argument(192        "--lora_rank",193        type=int,194        default=4,195        help="The dimension of the LoRA update matrices.",196    )197    parser.add_argument(198        "--lora_alpha",199        type=float,200        default=4.0,201        help="The weight of the LoRA update matrices.",202    )203    parser.add_argument(204        "--init_lora_weights",205        type=str,206        default="kaiming",207        choices=["gaussian", "kaiming"],208        help="The initializing method of LoRA weight.",209    )210    parser.add_argument(211        "--use_gradient_checkpointing",212        default=False,213        action="store_true",214        help="Whether to use gradient checkpointing.",215    )216    parser.add_argument(217        "--accumulate_grad_batches",218        type=int,219        default=1,220        help="The number of batches in gradient accumulation.",221    )222    parser.add_argument(223        "--training_strategy",224        type=str,225        default="auto",226        choices=["auto", "deepspeed_stage_1", "deepspeed_stage_2", "deepspeed_stage_3"],227        help="Training strategy",228    )229    parser.add_argument(230        "--max_epochs",231        type=int,232        default=1,233        help="Number of epochs.",234    )235    parser.add_argument(236        "--modelscope_model_id",237        type=str,238        default=None,239        help="Model ID on ModelScope (https://www.modelscope.cn/). The model will be uploaded to ModelScope automatically if you provide a Model ID.",240    )241    parser.add_argument(242        "--modelscope_access_token",243        type=str,244        default=None,245        help="Access key on ModelScope (https://www.modelscope.cn/). Required if you want to upload the model to ModelScope.",246    )247    parser.add_argument(248        "--pretrained_lora_path",249        type=str,250        default=None,251        help="Pretrained LoRA path. Required if the training is resumed.",252    )253    parser.add_argument(254        "--use_swanlab",255        default=False,256        action="store_true",257        help="Whether to use SwanLab logger.",258    )259    parser.add_argument(260        "--swanlab_mode",261        default=None,262        help="SwanLab mode (cloud or local).",263    )264    return parser265 266 267def launch_training_task(model, args):268    # dataset and data loader269    dataset = TextImageDataset(270        args.dataset_path,271        steps_per_epoch=args.steps_per_epoch * args.batch_size,272        height=args.height,273        width=args.width,274        center_crop=args.center_crop,275        random_flip=args.random_flip276    )277    train_loader = torch.utils.data.DataLoader(278        dataset,279        shuffle=True,280        batch_size=args.batch_size,281        num_workers=args.dataloader_num_workers282    )283    # train284    if args.use_swanlab:285        from swanlab.integration.pytorch_lightning import SwanLabLogger286        swanlab_config = {"UPPERFRAMEWORK": "DiffSynth-Studio"}287        swanlab_config.update(vars(args))288        swanlab_logger = SwanLabLogger(289            project="diffsynth_studio", 290            name="diffsynth_studio",291            config=swanlab_config,292            mode=args.swanlab_mode,293            logdir=os.path.join(args.output_path, "swanlog"),294        )295        logger = [swanlab_logger]296    else:297        logger = None298    trainer = pl.Trainer(299        max_epochs=args.max_epochs,300        accelerator="gpu",301        devices="auto",302        precision=args.precision,303        strategy=args.training_strategy,304        default_root_dir=args.output_path,305        accumulate_grad_batches=args.accumulate_grad_batches,306        callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)],307        logger=logger,308    )309    trainer.fit(model=model, train_dataloaders=train_loader)310 311    # Upload models312    if args.modelscope_model_id is not None and args.modelscope_access_token is not None:313        print(f"Uploading models to modelscope. model_id: {args.modelscope_model_id} local_path: {trainer.log_dir}")314        with open(os.path.join(trainer.log_dir, "configuration.json"), "w", encoding="utf-8") as f:315            f.write('{"framework":"Pytorch","task":"text-to-image-synthesis"}\n')316        api = HubApi()317        api.login(args.modelscope_access_token)318        api.push_model(model_id=args.modelscope_model_id, model_dir=trainer.log_dir)319