ComfyUI Node: YuE2 LoRA Trainer

YuE2 LoRA Trainer:

The YuE2LoRATrainer node is designed to facilitate the training of Low-Rank Adaptation (LoRA) models on the Non-Autoregressive (NAR) branch of YuE2's acoustic model. This node is particularly beneficial for users looking to customize and enhance the YuE2 model's performance by training it with their own audio datasets. By leveraging the flow-matching technique, the node allows for efficient training of LoRA adapters, which are lightweight and can be easily integrated into existing models. The primary goal of this node is to enable users to create personalized audio models that can generate unique styles or effects based on their input data. The node operates in a fully standalone manner, utilizing the official YuE2 model code and does not require any third-party custom nodes, making it accessible and straightforward for users to implement.

YuE2 LoRA Trainer Input Parameters:

dataset

The dataset parameter specifies the audio dataset to be used for training. It scans an audio folder and caches VAE latents, which are essential for the training process. This parameter is crucial as it determines the source of the audio data that the model will learn from.

checkpoint

The checkpoint parameter allows you to select the native YuE2 all-in-one checkpoint file from which the LoRA will be trained. This file is the same one used for generation, ensuring consistency between training and inference. The parameter provides a list of available checkpoints, and if none are found, it prompts you to refer to the README for guidance. This parameter is vital for ensuring that the training process is aligned with the model's current state.

trigger_word

The trigger_word parameter is a string that acts as a token or word placed at the start of the style prompt. It is used during the generation phase to invoke the specific style or effect that the LoRA model has been trained to produce. The default value is "mystyle," and it is recommended to use this word in your style prompt at generation time for optimal results.

steps

The steps parameter defines the total number of training steps, with each step involving one random clip from the dataset. The default value is 3000, with a minimum of 1 and a maximum of 100000. This parameter significantly impacts the strength and quality of the style learned by the model, with fewer steps resulting in a weaker style and more steps increasing the risk of overfitting.

learning_rate

The learning_rate parameter sets the learning rate for the AdamW optimizer used during training. The default value is 1e-4, with a range from 1e-7 to 1e-2. This parameter is crucial for controlling the speed and stability of the training process. Adjusting the learning rate can help prevent overfitting or underfitting, with higher values potentially leading to faster convergence and lower values providing more stable training.

rank

The rank parameter determines the LoRA rank, which refers to the adapter's capacity. The default value is 32, with a range from 1 to 256. This parameter affects the size and detail of the resulting LoRA model, with lower ranks producing smaller files with less detail and higher ranks offering more capacity but increasing the risk of overfitting.

YuE2 LoRA Trainer Output Parameters:

lora_path

The lora_path output parameter provides the file path to the trained LoRA model. This path is essential for accessing and utilizing the newly trained model in subsequent tasks or for further analysis. The output ensures that users can easily locate and apply their customized LoRA models.

YuE2 LoRA Trainer Usage Tips:

• Ensure that your audio dataset is well-organized and contains a diverse range of clips to improve the model's ability to generalize and produce high-quality outputs.
• Start with the default parameter values and gradually adjust them based on the results you observe. This approach helps in understanding the impact of each parameter on the training process.
• Use the trigger_word consistently in your style prompts during generation to achieve the desired effects and styles that the model has been trained to produce.

YuE2 LoRA Trainer Common Errors and Solutions:

Sequence exceeds context

• Explanation: This error occurs when the sequence length, including conditioning tokens and latents, exceeds the model's maximum position embeddings.
• Solution: Reduce the number of conditioning tokens or ensure that the sequence length is within the model's capacity.

No checkpoints found

• Explanation: This error indicates that no checkpoint files are available in the specified directory.
• Solution: Verify that the checkpoint files are correctly placed in the models/checkpoints directory and consult the README for additional setup instructions.

Overfitting observed

• Explanation: Overfitting can occur if the model is trained for too many steps or with a high learning rate, leading to repetitive or muffled outputs.
• Solution: Reduce the number of training steps or lower the learning rate to prevent overfitting and improve model generalization.