| 34 | |
| 35 | |
| 36 | class Pipeline(LightningModule): |
| 37 | def __init__( |
| 38 | self, |
| 39 | learning_rate: float = 1e-4, |
| 40 | num_workers: int = 4, |
| 41 | train: bool = True, |
| 42 | T: int = 1000, |
| 43 | weight_decay: float = 1e-2, |
| 44 | every_plot_step: int = 2000, |
| 45 | shift: float = 3.0, |
| 46 | logit_mean: float = 0.0, |
| 47 | logit_std: float = 1.0, |
| 48 | timestep_densities_type: str = "logit_normal", |
| 49 | ssl_coeff: float = 1.0, |
| 50 | checkpoint_dir=None, |
| 51 | max_steps: int = 200000, |
| 52 | warmup_steps: int = 10, |
| 53 | dataset_path: str = "./data/your_dataset_path", |
| 54 | lora_config_path: str = None, |
| 55 | adapter_name: str = "lora_adapter", |
| 56 | ): |
| 57 | super().__init__() |
| 58 | |
| 59 | self.save_hyperparameters() |
| 60 | self.is_train = train |
| 61 | self.T = T |
| 62 | |
| 63 | # Initialize scheduler |
| 64 | self.scheduler = self.get_scheduler() |
| 65 | |
| 66 | # step 1: load model |
| 67 | acestep_pipeline = ACEStepPipeline(checkpoint_dir) |
| 68 | acestep_pipeline.load_checkpoint(acestep_pipeline.checkpoint_dir) |
| 69 | |
| 70 | transformers = acestep_pipeline.ace_step_transformer.float().cpu() |
| 71 | transformers.enable_gradient_checkpointing() |
| 72 | |
| 73 | assert lora_config_path is not None, "Please provide a LoRA config path" |
| 74 | if lora_config_path is not None: |
| 75 | try: |
| 76 | from peft import LoraConfig |
| 77 | except ImportError: |
| 78 | raise ImportError("Please install peft library to use LoRA training") |
| 79 | with open(lora_config_path, encoding="utf-8") as f: |
| 80 | import json |
| 81 | lora_config = json.load(f) |
| 82 | lora_config = LoraConfig(**lora_config) |
| 83 | transformers.add_adapter(adapter_config=lora_config, adapter_name=adapter_name) |
| 84 | self.adapter_name = adapter_name |
| 85 | |
| 86 | self.transformers = transformers |
| 87 | |
| 88 | self.dcae = acestep_pipeline.music_dcae.float().cpu() |
| 89 | self.dcae.requires_grad_(False) |
| 90 | |
| 91 | self.text_encoder_model = acestep_pipeline.text_encoder_model.float().cpu() |
| 92 | self.text_encoder_model.requires_grad_(False) |
| 93 | self.text_tokenizer = acestep_pipeline.text_tokenizer |