| 5 | |
| 6 | |
| 7 | class LightningModel(LightningModelForT2ILoRA): |
| 8 | def __init__( |
| 9 | self, |
| 10 | torch_dtype=torch.float16, pretrained_weights=[], |
| 11 | learning_rate=1e-4, use_gradient_checkpointing=True, |
| 12 | 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, |
| 13 | ): |
| 14 | super().__init__(learning_rate=learning_rate, use_gradient_checkpointing=use_gradient_checkpointing) |
| 15 | # Load models |
| 16 | model_manager = ModelManager(torch_dtype=torch_dtype, device=self.device) |
| 17 | model_manager.load_models(pretrained_weights) |
| 18 | self.pipe = SDImagePipeline.from_model_manager(model_manager) |
| 19 | self.pipe.scheduler.set_timesteps(1000) |
| 20 | |
| 21 | self.freeze_parameters() |
| 22 | self.add_lora_to_model( |
| 23 | self.pipe.denoising_model(), |
| 24 | lora_rank=lora_rank, |
| 25 | lora_alpha=lora_alpha, |
| 26 | lora_target_modules=lora_target_modules, |
| 27 | init_lora_weights=init_lora_weights, |
| 28 | pretrained_lora_path=pretrained_lora_path, |
| 29 | ) |
| 30 | |
| 31 | |
| 32 | def parse_args(): |