!git clone https://github.com/SesameAILabs/csm
%cd csm
!pip install -r requirements.txt

%cd csm

from generator import load_csm_1b

generator = load_csm_1b(device="cuda")

import torchaudio

audio = generator.generate(
    text="Hi Bob, long time no see! did you catch the horse races this weekend?",
    speaker=0,
    context=[],
    max_audio_length_ms=10_000,
)
torchaudio.save("audio.wav", audio.unsqueeze(0).cpu(), generator.sample_rate)

from IPython.display import display, Audio

display(Audio("audio.wav"))

import torchaudio
from IPython.display import display, Audio

audio = generator.generate(
    text="Hey! it's been so long since we last caught up!",
    speaker=0,
    context=[],
    max_audio_length_ms=10_000,
)

torchaudio.save("speaker_0_1.wav", audio.unsqueeze(0).cpu(), generator.sample_rate)
display(Audio("speaker_0_1.wav"))

from generator import Segment

def load_audio(audio_path):
    audio_tensor, sample_rate = torchaudio.load(audio_path)
    audio_tensor = torchaudio.functional.resample(
        audio_tensor.squeeze(0), orig_freq=sample_rate, new_freq=generator.sample_rate
    )
    return audio_tensor

contexts = [
    Segment(
        text="Hi Bob, long time no see! did you catch the horse races this weekend?",
        speaker=0,
        audio=load_audio("speaker_0_1.wav")
    )
]

audio = generator.generate(
    text="Yea! I've really missed our chats. How have you been?",
    speaker=1,
    context=contexts,
    max_audio_length_ms=10_000,
)

torchaudio.save("speaker_1_1.wav", audio.unsqueeze(0).cpu(), generator.sample_rate)
display(Audio("speaker_1_1.wav"))

contexts.append(
    Segment(
        text="Yea! I've really missed our chats. How have you been?",
        speaker=1,
        audio=load_audio("speaker_1_1.wav")
    )
)

audio = generator.generate(
    text="I’ve been doing well, thanks. I’ve also kept up with our shared passion for horse racing.",
    speaker=0,
    context=contexts,
    max_audio_length_ms=10_000,
)

torchaudio.save("speaker_0_2.wav", audio.unsqueeze(0).cpu(), generator.sample_rate)
display(Audio("speaker_0_2.wav"))

contexts.append(
    Segment(
        text="I've been doing well, thanks. I've also kept up with our shared passion for horse racing.",
        speaker=0,
        audio=load_audio("speaker_0_2.wav")
    )
)

audio = generator.generate(
    text="That's great to hear! I watched a few races recently and the competition was fierce.",
    speaker=1,
    context=contexts,
    max_audio_length_ms=10_000,
)

torchaudio.save("speaker_1_2.wav", audio.unsqueeze(0).cpu(), generator.sample_rate)
display(Audio("speaker_1_2.wav"))

import torch
import torchaudio

device = "cuda" if torch.cuda.is_available() else "cpu"
generator = load_csm_1b(device=device)

# 2者の発話を自動生成し、初期プロンプトとして利用
default_text_a = "Hello, How's it going?"
default_text_b = "Hello, How's it going?"

audio_a = generator.generate(
    text=default_text_a,
    speaker=0,
    context=[],  # 初回なのでコンテキストなし
    max_audio_length_ms=10_000,
)
audio_b = generator.generate(
    text=default_text_b,
    speaker=1,
    context=[],  # 初回なのでコンテキストなし
    max_audio_length_ms=10_000,
)

# 初期プロンプトとしての Segment を作成
initial_prompt_a = Segment(text=default_text_a, speaker=0, audio=audio_a)
initial_prompt_b = Segment(text=default_text_b, speaker=1, audio=audio_b)

# 初期コンテキストに初回発話を格納
conversation_segments = [initial_prompt_a, initial_prompt_b]

# 会話履歴を定義、各行が交互の発話として扱う
conversation_text = """\
Hi, it's been so long since we last caught up!
Yea, I know! I've really missed our chats. How have you been?
I've been doing well, thanks. I've also kept up with our shared passion for horse racing.
That's great to hear! I watched a few races recently and the competition was fierce.
Absolutely, the races were exhilarating, especially that unexpected win by an underdog.
I couldn't agree more—the strategy displayed by the jockey was simply brilliant.
Have you read about the new training techniques being introduced? They seem to be changing the game.
Yes, I saw an article on that! It looks like they could really enhance performance in future races.
I'm excited to see how these developments will shape upcoming events. It really adds another layer of thrill.
Definitely. We should plan to catch a race together soon and discuss all these new trends in person!
"""

utterances = [line.strip() for line in conversation_text.split("\n") if line.strip()]

# 各発話の生成
for i, utterance in enumerate(utterances):
    # 発話は交互に Speaker A (0) と Speaker B (1) に割り当て
    speaker = i % 2
    print(f"Generating audio for Speaker {'A' if speaker == 0 else 'B'}: {utterance}")

    audio_tensor = generator.generate(
        text=utterance,
        speaker=speaker,
        context=conversation_segments,
        max_audio_length_ms=10_000,
    )

    segment = Segment(text=utterance, speaker=speaker, audio=audio_tensor)
    conversation_segments.append(segment)

# 全ての発話を連結して1つの wav ファイルに出力
# ※最初のプロンプトの元になった2つの音声は除いている
all_audio_tensors = [segment.audio for segment in conversation_segments[2:]]
final_audio = torch.cat(all_audio_tensors, dim=0)

output_filename = "final_conversation.wav"
torchaudio.save(output_filename, final_audio.unsqueeze(0).cpu(), generator.sample_rate)

display(Audio("final_conversation.wav"))