ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter8/sesame/inference.py

@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Sesame CSM (1B) TTS - Inference Script
+
+This script loads a trained LoRA model and generates speech from text.
+"""
+
+import argparse
+import torch
+import soundfile as sf
+from pathlib import Path
+from datasets import load_dataset, Audio
+from unsloth import FastModel
+from transformers import CsmForConditionalGeneration
+from peft import PeftModel
+
+
+def load_model(base_model_name: str, lora_path: str = None, load_in_4bit: bool = False):
+    """
+    Load the base model and optionally apply LoRA adapters.
+    
+    Args:
+        base_model_name: Name or path of the base model
+        lora_path: Path to saved LoRA adapters (optional)
+        load_in_4bit: Whether to load model in 4-bit quantization
+    
+    Returns:
+        model, processor
+    """
+    print(f"Loading base model: {base_model_name}")
+    model, processor = FastModel.from_pretrained(
+        model_name=base_model_name,
+        max_seq_length=2048,
+        dtype=None,  # Auto-detection
+        auto_model=CsmForConditionalGeneration,
+        load_in_4bit=load_in_4bit,
+    )
+    
+    if lora_path:
+        print(f"Loading LoRA adapters from: {lora_path}")
+        model = PeftModel.from_pretrained(model, lora_path)
+    
+    return model, processor
+
+
+def load_dataset_for_context(dataset_name: str = "maxbsoft/mrdragonfox-elise", split: str = "train"):
+    """
+    Load the dataset for voice context examples.
+    
+    Args:
+        dataset_name: Name of the dataset
+        split: Dataset split to use
+    
+    Returns:
+        Dataset
+    """
+    print(f"Loading dataset: {dataset_name}")
+    raw_ds = load_dataset(dataset_name, split=split)
+    target_sampling_rate = 24000
+    raw_ds = raw_ds.cast_column("audio", Audio(sampling_rate=target_sampling_rate))
+    print(f"Loaded {len(raw_ds)} examples from dataset")
+    return raw_ds
+
+
+def generate_speech(
+    model,
+    processor,
+    text: str,
+    speaker_id: int = 0,
+    max_new_tokens: int = 125,
+    output_path: str = "output.wav",
+    dataset_context_idx: int = None,
+    dataset_name: str = "maxbsoft/mrdragonfox-elise",
+):
+    """
+    Generate speech from text.
+    
+    Args:
+        model: The loaded model
+        processor: The processor
+        text: Text to convert to speech
+        speaker_id: Speaker ID (for multi-speaker models)
+        max_new_tokens: Maximum number of tokens to generate (125 tokens ≈ 10 seconds)
+        output_path: Path to save the output audio file
+        dataset_context_idx: Optional dataset index to use for voice consistency
+        dataset_name: Name of the dataset to load context from
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Prepare inputs based on whether context is provided
+    if dataset_context_idx is not None:
+        print(f"Generating speech with voice context from dataset index: {dataset_context_idx}")
+        print(f"Target text: '{text}'")
+        
+        # Load dataset and get context example
+        raw_ds = load_dataset_for_context(dataset_name)
+        context_example = raw_ds[dataset_context_idx]
+        context_audio = context_example["audio"]["array"]
+        context_text = context_example["text"]
+        
+        print(f"Context text: '{context_text}'")
+        
+        # Use conversation format with audio context for voice consistency
+        conversation = [
+            {
+                "role": str(speaker_id),
+                "content": [
+                    {"type": "text", "text": context_text},
+                    {"type": "audio", "path": context_audio}
+                ]
+            },
+            {
+                "role": str(speaker_id),
+                "content": [{"type": "text", "text": text}]
+            },
+        ]
+        inputs = processor.apply_chat_template(
+            conversation,
+            tokenize=True,
+            return_dict=True,
+        ).to(device)
+    else:
+        print(f"Generating speech without context for: '{text}'")
+        
+        # Simple text-only input
+        inputs = processor(
+            f"[{speaker_id}]{text}",
+            add_special_tokens=True,
+            return_tensors="pt"
+        ).to(device)
+    
+    # Generate audio
+    print("Generating audio...")
+    with torch.no_grad():
+        audio_values = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs.get("attention_mask"),
+            max_new_tokens=max_new_tokens,
+            output_audio=True,
+        )
+    
+    # Save audio
+    audio = audio_values[0].to(torch.float32).cpu().numpy()
+    sf.write(output_path, audio, 24000)
+    print(f"Audio saved to: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate speech using Sesame CSM TTS model"
+    )
+    parser.add_argument(
+        "--base-model",
+        type=str,
+        default="unsloth/csm-1b",
+        help="Base model name or path (default: unsloth/csm-1b)"
+    )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        default=None,
+        help="Path to saved LoRA adapters (optional)"
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to convert to speech"
+    )
+    parser.add_argument(
+        "--speaker-id",
+        type=int,
+        default=0,
+        help="Speaker ID (default: 0)"
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output.wav",
+        help="Output audio file path (default: output.wav)"
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=125,
+        help="Maximum tokens to generate (125 ≈ 10 seconds) (default: 125)"
+    )
+    parser.add_argument(
+        "--load-in-4bit",
+        action="store_true",
+        help="Load model in 4-bit quantization to reduce memory usage"
+    )
+    parser.add_argument(
+        "--dataset-context-idx",
+        type=int,
+        default=None,
+        help="Dataset index to use for voice consistency (e.g., 3 or 4 from training examples)"
+    )
+    parser.add_argument(
+        "--dataset-name",
+        type=str,
+        default="maxbsoft/mrdragonfox-elise",
+        help="Dataset name to load context from (default: public Elise mirror)"
+    )
+    
+    args = parser.parse_args()
+    
+    # Load model
+    model, processor = load_model(
+        base_model_name=args.base_model,
+        lora_path=args.lora_path,
+        load_in_4bit=args.load_in_4bit
+    )
+    
+    # Generate speech
+    generate_speech(
+        model=model,
+        processor=processor,
+        text=args.text,
+        speaker_id=args.speaker_id,
+        max_new_tokens=args.max_tokens,
+        output_path=args.output,
+        dataset_context_idx=args.dataset_context_idx,
+        dataset_name=args.dataset_name,
+    )
+
+
+if __name__ == "__main__":
+    main()