# /// script # requires-python = ">=3.10,<3.14" # dependencies = [ # "accelerate", # "bitsandbytes>=0.43.0", # "datasets==4.3.0", # "einx", # "hf_transfer", # "huggingface_hub>=0.34.0", # "marimo", # "omegaconf", # "peft", # "protobuf", # "sentencepiece", # "torchao>=0.16.0", # "torchcodec", # "transformers==4.56.2", # "triton>=3.2.0", # "trl==0.22.2", # "unsloth @ git+https://github.com/unslothai/unsloth", # "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo", # ] # # [tool.uv] # no-build-package = [ # "bitsandbytes", # "triton", # "vllm", # "xformers", # ] # /// import marimo __generated_with = "0.23.8" app = marimo.App() @app.cell def _(): import marimo as mo return (mo,) @app.cell def _(): import subprocess return @app.cell(hide_code=True) def _(mo): mo.md(r""" To run this notebook, hit the **▶ Run all** button in the bottom-right corner - or use `Ctrl/Cmd + Shift + R`.

Join Discord if you need help + ⭐ Star us on Github ⭐

To install Unsloth on your local device, follow [our guide](https://unsloth.ai/docs/get-started/install). This notebook is licensed [LGPL-3.0](https://github.com/unslothai/notebooks?tab=LGPL-3.0-1-ov-file#readme). You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & how to save it """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### News """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" Introducing **Unsloth Studio** - a new open source, no-code web UI to train and run LLMs. [Blog](https://unsloth.ai/docs/new/studio) • [Notebook](https://github.com/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb)

_{Train models — no code needed}

_{Run GGUF models on Mac, Windows & Linux}

Train MoEs - DeepSeek, GLM, Qwen and gpt-oss 12x faster with 35% less VRAM. [Blog](https://unsloth.ai/docs/new/faster-moe) Ultra Long-Context Reinforcement Learning is here with 7x more context windows! [Blog](https://unsloth.ai/docs/new/grpo-long-context) New in Reinforcement Learning: [FP8 RL](https://unsloth.ai/docs/new/fp8-reinforcement-learning) • [Vision RL](https://unsloth.ai/docs/new/vision-reinforcement-learning-vlm-rl) • [Standby](https://unsloth.ai/docs/basics/memory-efficient-rl) • [gpt-oss RL](https://unsloth.ai/docs/new/gpt-oss-reinforcement-learning) Visit our docs for all our [model uploads](https://unsloth.ai/docs/get-started/unsloth-model-catalog) and [notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks). """) return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Unsloth `FastModel` supports loading nearly any model now! This includes Vision and Text models! Thank you to [Etherl](https://huggingface.co/Etherll) for creating this notebook! """) return @app.cell def _(): from unsloth import FastModel import torch from huggingface_hub import snapshot_download max_seq_length = 2048 # Choose any for long context! fourbit_models = [ # 4bit dynamic quants for superior accuracy and low memory use "unsloth/gemma-3-4b-it-unsloth-bnb-4bit", "unsloth/gemma-3-12b-it-unsloth-bnb-4bit", "unsloth/gemma-3-27b-it-unsloth-bnb-4bit", # Qwen3 new models "unsloth/Qwen3-4B-unsloth-bnb-4bit", "unsloth/Qwen3-8B-unsloth-bnb-4bit", # Other very popular models! "unsloth/Llama-3.1-8B", "unsloth/Llama-3.2-3B", "unsloth/Llama-3.3-70B", "unsloth/mistral-7b-instruct-v0.3", "unsloth/Phi-4", ] # More models at https://huggingface.co/unsloth # Download model and code snapshot_download("unsloth/Spark-TTS-0.5B", local_dir="Spark-TTS-0.5B") model, tokenizer = FastModel.from_pretrained( model_name=f"Spark-TTS-0.5B/LLM", max_seq_length=max_seq_length, # Choose any for long context! dtype=torch.float32, # Spark seems to only work on float32 for now full_finetuning=True, # We support full finetuning now! load_in_4bit=False, # token = "YOUR_HF_TOKEN", # HF Token for gated models ) return FastModel, max_seq_length, model, tokenizer, torch @app.cell(hide_code=True) def _(mo): mo.md(r""" We now add LoRA adapters so we only need to update 1 to 10% of all parameters! """) return @app.cell def _(FastModel, model): # LoRA does not work with float32 only works with bfloat16 ! model_1 = FastModel.get_peft_model( model, r=128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], lora_alpha=128, lora_dropout=0, # Supports any, but = 0 is optimized bias="none", # Supports any, but = "none" is optimized use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context random_state=3407, use_rslora=False, # We support rank stabilized LoRA loftq_config=None, # And LoftQ ) return (model_1,) @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Data Prep We will use the `MrDragonFox/Elise`, which is designed for training TTS models. Ensure that your dataset follows the required format: **text, audio** for single-speaker models or **source, text, audio** for multi-speaker models. You can modify this section to accommodate your own dataset, but maintaining the correct structure is essential for optimal training. """) return @app.cell def _(): from datasets import load_dataset dataset = load_dataset("MrDragonFox/Elise", split="train") return (dataset,) @app.cell def _(dataset, torch): # @title Tokenization Function import locale import torchaudio.transforms as T import os import sys import numpy as np sys.path.append("Spark-TTS") from sparktts.models.audio_tokenizer import BiCodecTokenizer from sparktts.utils.audio import audio_volume_normalize audio_tokenizer = BiCodecTokenizer("Spark-TTS-0.5B", "cuda") def extract_wav2vec2_features(wavs: torch.Tensor) -> torch.Tensor: """extract wav2vec2 features""" if wavs.shape[0] != 1: raise ValueError(f"Expected batch size 1, but got shape {wavs.shape}") wav_np = wavs.squeeze(0).cpu().numpy() processed = audio_tokenizer.processor( wav_np, sampling_rate=16000, return_tensors="pt", padding=True ) input_values = processed.input_values input_values = input_values.to(audio_tokenizer.feature_extractor.device) model_output = audio_tokenizer.feature_extractor(input_values) if model_output.hidden_states is None: raise ValueError( "Wav2Vec2Model did not return hidden states. Ensure config `output_hidden_states=True`." ) num_layers = len(model_output.hidden_states) required_layers = [11, 14, 16] if any((l >= num_layers for l in required_layers)): raise IndexError( f"Requested hidden state indices {required_layers} out of range for model with {num_layers} layers." ) feats_mix = ( model_output.hidden_states[11] + model_output.hidden_states[14] + model_output.hidden_states[16] ) / 3 return feats_mix def formatting_audio_func(example): text = ( f"{example['source']}: {example['text']}" if "source" in example else example["text"] ) audio_array = example["audio"]["array"] sampling_rate = example["audio"]["sampling_rate"] target_sr = audio_tokenizer.config["sample_rate"] if sampling_rate != target_sr: resampler = T.Resample(orig_freq=sampling_rate, new_freq=target_sr) audio_tensor_temp = torch.from_numpy(audio_array).float() audio_array = resampler(audio_tensor_temp).numpy() if audio_tokenizer.config["volume_normalize"]: audio_array = audio_volume_normalize(audio_array) ref_wav_np = audio_tokenizer.get_ref_clip(audio_array) audio_tensor = ( torch.from_numpy(audio_array) .unsqueeze(0) .float() .to(audio_tokenizer.device) ) ref_wav_tensor = ( torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(audio_tokenizer.device) ) feat = extract_wav2vec2_features(audio_tensor) batch = { "wav": audio_tensor, "ref_wav": ref_wav_tensor, "feat": feat.to(audio_tokenizer.device), } semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(batch) global_tokens = "".join( [ f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze().cpu().numpy() ] ) semantic_tokens = "".join( [ f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze().cpu().numpy() ] ) inputs = [ "<|task_tts|>", "<|start_content|>", text, "<|end_content|>", "<|start_global_token|>", global_tokens, "<|end_global_token|>", "<|start_semantic_token|>", semantic_tokens, "<|end_semantic_token|>", "<|im_end|>", ] inputs = "".join(inputs) return {"text": inputs} dataset_1 = dataset.map(formatting_audio_func, remove_columns=["audio"]) print("Moving Bicodec model and Wav2Vec2Model to cpu.") audio_tokenizer.model.cpu() audio_tokenizer.feature_extractor.cpu() torch.cuda.empty_cache() # Squeeze batch dim return audio_tokenizer, dataset_1, np @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Train the model Now let's train our model. We do 60 steps to speed things up, but you can set `num_train_epochs=1` for a full run, and turn off `max_steps=None`. We also support `DPOTrainer` and `GRPOTrainer` for reinforcement learning!! """) return @app.cell def _(dataset_1, max_seq_length, model_1, tokenizer): from trl import SFTConfig, SFTTrainer trainer = SFTTrainer( model=model_1, tokenizer=tokenizer, train_dataset=dataset_1, dataset_text_field="text", max_seq_length=max_seq_length, # Choose any for long context! packing=False, # Can make training 5x faster for short sequences. args=SFTConfig( per_device_train_batch_size=2, gradient_accumulation_steps=4, warmup_steps=5, max_steps=60, learning_rate=0.0002, fp16=False, # We're doing full float32 so disable mixed precision bf16=False, # We're doing full float32 so disable mixed precision logging_steps=1, optim="adamw_8bit", weight_decay=0.001, lr_scheduler_type="linear", seed=3407, output_dir="outputs", report_to="none", # Use TrackIO/WandB etc ), ) return (trainer,) @app.cell def _(torch): # @title Show current memory stats gpu_stats = torch.cuda.get_device_properties(0) start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.") print(f"{start_gpu_memory} GB of memory reserved.") return max_memory, start_gpu_memory @app.cell def _(trainer): trainer_stats = trainer.train() return (trainer_stats,) @app.cell def _(max_memory, start_gpu_memory, torch, trainer_stats): # @title Show final memory and time stats used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) used_memory_for_lora = round(used_memory - start_gpu_memory, 3) used_percentage = round(used_memory / max_memory * 100, 3) lora_percentage = round(used_memory_for_lora / max_memory * 100, 3) print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.") print( f"{round(trainer_stats.metrics['train_runtime'] / 60, 2)} minutes used for training." ) print(f"Peak reserved memory = {used_memory} GB.") print(f"Peak reserved memory for training = {used_memory_for_lora} GB.") print(f"Peak reserved memory % of max memory = {used_percentage} %.") print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.") return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Inference Let's run the model! You can change the prompts """) return @app.cell def _(): input_text = "Hey there my name is Elise, and I'm a speech generation model that can sound like a person." chosen_voice = None # None for single-speaker return chosen_voice, input_text @app.cell def _( FastModel, audio_tokenizer, chosen_voice, input_text, model_1, np, tokenizer, torch, ): # @title Run Inference import re from typing import Dict, Any FastModel.for_inference(model_1) @torch.inference_mode() def generate_speech_from_text( text: str, temperature: float = 0.8, top_k: int = 50, top_p: float = 1, max_new_audio_tokens: int = 2048, device: torch.device = torch.device( "cuda" if torch.cuda.is_available() else "cpu" ), ) -> np.ndarray: """ Generates speech audio from text using default voice control parameters. # Enable native 2x faster inference Args: text (str): The text input to be converted to speech. temperature (float): Sampling temperature for generation. top_k (int): Top-k sampling parameter. # Generation temperature top_p (float): Top-p (nucleus) sampling parameter. # Generation top_k max_new_audio_tokens (int): Max number of new tokens to generate (limits audio length). # Generation top_p device (torch.device): Device to run inference on. # Max tokens for audio part Returns: np.ndarray: Generated waveform as a NumPy array. """ torch.compiler.reset() prompt = "".join( [ "<|task_tts|>", "<|start_content|>", text, "<|end_content|>", "<|start_global_token|>", ] ) model_inputs = tokenizer([prompt], return_tensors="pt").to(device) print("Generating token sequence...") generated_ids = model_1.generate( **model_inputs, max_new_tokens=max_new_audio_tokens, # Limit generation length do_sample=True, temperature=temperature, top_k=top_k, top_p=top_p, eos_token_id=tokenizer.eos_token_id, # Stop token pad_token_id=tokenizer.pad_token_id, # Use models pad token id ) print("Token sequence generated.") generated_ids_trimmed = generated_ids[:, model_inputs.input_ids.shape[1] :] predicts_text = tokenizer.batch_decode( generated_ids_trimmed, skip_special_tokens=False )[0] semantic_matches = re.findall("<\\|bicodec_semantic_(\\d+)\\|>", predicts_text) if not semantic_matches: print("Warning: No semantic tokens found in the generated output.") return np.array([], dtype=np.float32) pred_semantic_ids = ( # Add batch dim torch.tensor([int(token) for token in semantic_matches]).long().unsqueeze(0) ) global_matches = re.findall("<\\|bicodec_global_(\\d+)\\|>", predicts_text) if not global_matches: print( "Warning: No global tokens found in the generated output (controllable mode). Might use defaults or fail." ) pred_global_ids = torch.zeros((1, 1), dtype=torch.long) # Add batch dim else: pred_global_ids = ( # Add batch dim torch.tensor([int(token) for token in global_matches]) .long() .unsqueeze(0) ) pred_global_ids = pred_global_ids.unsqueeze(0) # Add batch dim print(f"Found {pred_semantic_ids.shape[1]} semantic tokens.") print(f"Found {pred_global_ids.shape[2]} global tokens.") print("Detokenizing audio tokens...") audio_tokenizer.device = device audio_tokenizer.model.to(device) wav_np = audio_tokenizer.detokenize( pred_global_ids.to(device).squeeze(0), pred_semantic_ids.to(device) ) print("Detokenization complete.") return wav_np if __name__ == "__main__": # Limit generation length print(f"Generating speech for: '{input_text}'") text = f"{chosen_voice}: " + input_text if chosen_voice else input_text generated_waveform = generate_speech_from_text(input_text) if generated_waveform.size > 0: import soundfile as sf # Stop token output_filename = ( "generated_speech_controllable.wav" # Use models pad token id ) sample_rate = audio_tokenizer.config.get("sample_rate", 16000) sf.write(output_filename, generated_waveform, sample_rate) print(f"Audio saved to {output_filename}") from IPython.display import Audio, display display(Audio(generated_waveform, rate=sample_rate)) else: print( "Audio generation failed (no tokens found?)." ) return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Saving, loading finetuned models To save the final model as LoRA adapters, either use Hugging Face's `push_to_hub` for an online save or `save_pretrained` for a local save. **[NOTE]** This ONLY saves the LoRA adapters, and not the full model. To save to 16bit or GGUF, scroll down! """) return @app.cell def _(model_1, tokenizer): model_1.save_pretrained("spark_tts_lora") tokenizer.save_pretrained("spark_tts_lora") return @app.cell(hide_code=True) def _(mo): mo.md(r""" ### Saving to float16 We also support saving to `float16` directly. Select `merged_16bit` for float16 or `merged_4bit` for int4. We also allow `lora` adapters as a fallback. Use `push_to_hub_merged` to upload to your Hugging Face account! You can go to https://huggingface.co/settings/tokens for your personal tokens. See [our docs](https://unsloth.ai/docs/basics/inference-and-deployment) for more deployment options. """) return @app.cell def _(model_1, tokenizer): # Merge to 16bit if False: model_1.save_pretrained_merged( "spark_tts_finetune_16bit", tokenizer, save_method="merged_16bit" ) if False: # Merge to 4bit model_1.push_to_hub_merged( "HF_USERNAME/spark_tts_finetune_16bit", tokenizer, save_method="merged_16bit", token="YOUR_HF_TOKEN", ) if False: model_1.save_pretrained_merged( "spark_tts_finetune_4bit", tokenizer, save_method="merged_4bit" ) # Just LoRA adapters if False: model_1.push_to_hub_merged( "HF_USERNAME/spark_tts_finetune_4bit", tokenizer, save_method="merged_4bit", token="YOUR_HF_TOKEN", ) if False: model_1.save_pretrained("spark_tts_lora") tokenizer.save_pretrained("spark_tts_lora") if False: model_1.push_to_hub("HF_USERNAME/spark_tts_lora", token="YOUR_HF_TOKEN") tokenizer.push_to_hub("HF_USERNAME/spark_tts_lora", token="YOUR_HF_TOKEN") return @app.cell(hide_code=True) def _(mo): mo.md(r""" And we're done! If you have any questions on Unsloth, we have a [Discord](https://discord.gg/unsloth) channel! If you find any bugs or want to keep updated with the latest LLM stuff, or need help, join projects etc, feel free to join our Discord! Some other resources: 1. Looking to use Unsloth locally? Read our [Installation Guide](https://unsloth.ai/docs/get-started/install) for details on installing Unsloth on Windows, Docker, AMD, Intel GPUs. 2. Learn how to do Reinforcement Learning with our [RL Guide and notebooks](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide). 3. Read our guides and notebooks for [Text-to-speech (TTS)](https://unsloth.ai/docs/basics/text-to-speech-tts-fine-tuning) and [vision](https://unsloth.ai/docs/basics/vision-fine-tuning) model support. 4. Explore our [LLM Tutorials Directory](https://unsloth.ai/docs/models/tutorials-how-to-fine-tune-and-run-llms) to find dedicated guides for each model. 5. Need help with Inference? Read our [Inference & Deployment page](https://unsloth.ai/docs/basics/inference-and-deployment) for details on using vLLM, llama.cpp, Ollama etc.

Join Discord if you need help + ⭐️ Star us on Github ⭐️ This notebook and all Unsloth notebooks are licensed [LGPL-3.0](https://github.com/unslothai/notebooks?tab=LGPL-3.0-1-ov-file#readme)

""") return if __name__ == "__main__": app.run()