modelscope/DiffSynth-Painter

1import lightning as pl2from peft import LoraConfig, inject_adapter_in_model3import torch, os4from ..data.simple_text_image import TextImageDataset5from modelscope.hub.api import HubApi6 7 8 9class LightningModelForT2ILoRA(pl.LightningModule):10    def __init__(11        self,12        learning_rate=1e-4,13        use_gradient_checkpointing=True,14    ):15        super().__init__()16        # Set parameters17        self.learning_rate = learning_rate18        self.use_gradient_checkpointing = use_gradient_checkpointing19 20 21    def load_models(self):22        # This function is implemented in other modules23        self.pipe = None24 25 26    def freeze_parameters(self):27        # Freeze parameters28        self.pipe.requires_grad_(False)29        self.pipe.eval()30        self.pipe.denoising_model().train()31 32    33    def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out"):34        # Add LoRA to UNet35        lora_config = LoraConfig(36            r=lora_rank,37            lora_alpha=lora_alpha,38            init_lora_weights="gaussian",39            target_modules=lora_target_modules.split(","),40        )41        model = inject_adapter_in_model(lora_config, model)42        for param in model.parameters():43            # Upcast LoRA parameters into fp3244            if param.requires_grad:45                param.data = param.to(torch.float32)46 47 48    def training_step(self, batch, batch_idx):49        # Data50        text, image = batch["text"], batch["image"]51 52        # Prepare input parameters53        self.pipe.device = self.device54        prompt_emb = self.pipe.encode_prompt(text, positive=True)55        latents = self.pipe.vae_encoder(image.to(dtype=self.pipe.torch_dtype, device=self.device))56        noise = torch.randn_like(latents)57        timestep_id = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,))58        timestep = self.pipe.scheduler.timesteps[timestep_id].to(self.device)59        extra_input = self.pipe.prepare_extra_input(latents)60        noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep)61        training_target = self.pipe.scheduler.training_target(latents, noise, timestep)62 63        # Compute loss64        noise_pred = self.pipe.denoising_model()(65            noisy_latents, timestep=timestep, **prompt_emb, **extra_input,66            use_gradient_checkpointing=self.use_gradient_checkpointing67        )68        loss = torch.nn.functional.mse_loss(noise_pred, training_target)69 70        # Record log71        self.log("train_loss", loss, prog_bar=True)72        return loss73 74 75    def configure_optimizers(self):76        trainable_modules = filter(lambda p: p.requires_grad, self.pipe.denoising_model().parameters())77        optimizer = torch.optim.AdamW(trainable_modules, lr=self.learning_rate)78        return optimizer79    80 81    def on_save_checkpoint(self, checkpoint):82        checkpoint.clear()83        trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.pipe.denoising_model().named_parameters()))84        trainable_param_names = set([named_param[0] for named_param in trainable_param_names])85        state_dict = self.pipe.denoising_model().state_dict()86        for name, param in state_dict.items():87            if name in trainable_param_names:88                checkpoint[name] = param89 90 91 92def add_general_parsers(parser):93    parser.add_argument(94        "--dataset_path",95        type=str,96        default=None,97        required=True,98        help="The path of the Dataset.",99    )100    parser.add_argument(101        "--output_path",102        type=str,103        default="./",104        help="Path to save the model.",105    )106    parser.add_argument(107        "--steps_per_epoch",108        type=int,109        default=500,110        help="Number of steps per epoch.",111    )112    parser.add_argument(113        "--height",114        type=int,115        default=1024,116        help="Image height.",117    )118    parser.add_argument(119        "--width",120        type=int,121        default=1024,122        help="Image width.",123    )124    parser.add_argument(125        "--center_crop",126        default=False,127        action="store_true",128        help=(129            "Whether to center crop the input images to the resolution. If not set, the images will be randomly"130            " cropped. The images will be resized to the resolution first before cropping."131        ),132    )133    parser.add_argument(134        "--random_flip",135        default=False,136        action="store_true",137        help="Whether to randomly flip images horizontally",138    )139    parser.add_argument(140        "--batch_size",141        type=int,142        default=1,143        help="Batch size (per device) for the training dataloader.",144    )145    parser.add_argument(146        "--dataloader_num_workers",147        type=int,148        default=0,149        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",150    )151    parser.add_argument(152        "--precision",153        type=str,154        default="16-mixed",155        choices=["32", "16", "16-mixed"],156        help="Training precision",157    )158    parser.add_argument(159        "--learning_rate",160        type=float,161        default=1e-4,162        help="Learning rate.",163    )164    parser.add_argument(165        "--lora_rank",166        type=int,167        default=4,168        help="The dimension of the LoRA update matrices.",169    )170    parser.add_argument(171        "--lora_alpha",172        type=float,173        default=4.0,174        help="The weight of the LoRA update matrices.",175    )176    parser.add_argument(177        "--use_gradient_checkpointing",178        default=False,179        action="store_true",180        help="Whether to use gradient checkpointing.",181    )182    parser.add_argument(183        "--accumulate_grad_batches",184        type=int,185        default=1,186        help="The number of batches in gradient accumulation.",187    )188    parser.add_argument(189        "--training_strategy",190        type=str,191        default="auto",192        choices=["auto", "deepspeed_stage_1", "deepspeed_stage_2", "deepspeed_stage_3"],193        help="Training strategy",194    )195    parser.add_argument(196        "--max_epochs",197        type=int,198        default=1,199        help="Number of epochs.",200    )201    parser.add_argument(202        "--modelscope_model_id",203        type=str,204        default=None,205        help="Model ID on ModelScope (https://www.modelscope.cn/). The model will be uploaded to ModelScope automatically if you provide a Model ID.",206    )207    parser.add_argument(208        "--modelscope_access_token",209        type=str,210        default=None,211        help="Access key on ModelScope (https://www.modelscope.cn/). Required if you want to upload the model to ModelScope.",212    )213    return parser214 215 216def launch_training_task(model, args):217    # dataset and data loader218    dataset = TextImageDataset(219        args.dataset_path,220        steps_per_epoch=args.steps_per_epoch * args.batch_size,221        height=args.height,222        width=args.width,223        center_crop=args.center_crop,224        random_flip=args.random_flip225    )226    train_loader = torch.utils.data.DataLoader(227        dataset,228        shuffle=True,229        batch_size=args.batch_size,230        num_workers=args.dataloader_num_workers231    )232 233    # train234    trainer = pl.Trainer(235        max_epochs=args.max_epochs,236        accelerator="gpu",237        devices="auto",238        precision=args.precision,239        strategy=args.training_strategy,240        default_root_dir=args.output_path,241        accumulate_grad_batches=args.accumulate_grad_batches,242        callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)]243    )244    trainer.fit(model=model, train_dataloaders=train_loader)245 246    # Upload models247    if args.modelscope_model_id is not None and args.modelscope_access_token is not None:248        print(f"Uploading models to modelscope. model_id: {args.modelscope_model_id} local_path: {trainer.log_dir}")249        with open(os.path.join(trainer.log_dir, "configuration.json"), "w", encoding="utf-8") as f:250            f.write('{"framework":"Pytorch","task":"text-to-image-synthesis"}\n')251        api = HubApi()252        api.login(args.modelscope_access_token)253        api.push_model(model_id=args.modelscope_model_id, model_dir=trainer.log_dir)254