client_example.py

"""
CusVoiceModel API 클라이언트 예제
FastAPI 서버 사용법을 보여주는 완전한 예제
"""

import requests
import time
import json
from pathlib import Path
import argparse

class CusVoiceModelClient:
    """CusVoiceModel FastAPI 서버용 클라이언트"""

    def __init__(self, base_url: str = "http://localhost:8090"):
        self.base_url = base_url.rstrip('/')

    def health_check(self):
        """서버 상태 확인"""
        try:
            response = requests.get(f"{self.base_url}/health")
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"상태 확인 실패: {e}")
            return None

    def generate_speech(self, text: str, **kwargs):
        """텍스트에서 음성 생성"""
        payload = {
            "text": text,
            **kwargs
        }

        try:
            response = requests.post(f"{self.base_url}/generate", json=payload)
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"음성 생성 요청 실패: {e}")
            return None

    def generate_with_voice(self, text: str, voice_file_path: str, **kwargs):
        """음성 클로닝으로 음성 생성"""
        files = {"voice_file": open(voice_file_path, "rb")}
        data = {
            "text": text,
            **kwargs
        }

        try:
            response = requests.post(f"{self.base_url}/generate-with-voice",
                                   files=files, data=data)
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"음성 클로닝 요청 실패: {e}")
            return None
        finally:
            files["voice_file"].close()

    def get_task_status(self, task_id: str):
        """작업 상태 확인"""
        try:
            response = requests.get(f"{self.base_url}/status/{task_id}")
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            print(f"상태 확인 실패: {e}")
            return None

    def download_audio(self, task_id: str, output_path: str):
        """생성된 오디오 다운로드"""
        try:
            response = requests.get(f"{self.base_url}/download/{task_id}")
            response.raise_for_status()

            with open(output_path, "wb") as f:
                f.write(response.content)
            return True
        except requests.RequestException as e:
            print(f"다운로드 실패: {e}")
            return False

    def wait_for_completion(self, task_id: str, timeout: int = 300):
        """작업 완료 대기"""
        start_time = time.time()

        while time.time() - start_time < timeout:
            status = self.get_task_status(task_id)

            if not status:
                return None

            print(f"상태: {status['status']} - {status['message']}")

            if status["status"] == "completed":
                return status
            elif status["status"] == "failed":
                print(f"작업 실패: {status['message']}")
                return status

            time.sleep(2)

        print("완료 대기 시간 초과")
        return None

def main():
    parser = argparse.ArgumentParser(description="CusVoiceModel API 클라이언트 예제")
    parser.add_argument("--url", default="http://localhost:8090",
                       help="API 서버 URL")
    parser.add_argument("--text", default="안녕하세요, CusVoiceModel API 테스트입니다.",
                       help="음성으로 변환할 텍스트")
    parser.add_argument("--model", default="CusVoiceModel-1.5B",
                       choices=["CusVoiceModel-1.5B", "CusVoiceModel-Large", "CusVoiceModel-Large-Quant-4Bit"],
                       help="사용할 모델")
    parser.add_argument("--voice", help="클로닝용 참조 음성 파일 경로")
    parser.add_argument("--output", default="generated_speech.wav",
                       help="출력 오디오 파일 경로")
    parser.add_argument("--seed", type=int, default=42, help="랜덤 시드")
    parser.add_argument("--steps", type=int, default=20, help="디퓨전 스텝")
    parser.add_argument("--cfg-scale", type=float, default=1.3, help="CFG 스케일")

    args = parser.parse_args()

    # 클라이언트 초기화
    client = CusVoiceModelClient(args.url)

    # 서버 상태 확인
    print("서버 상태 확인 중...")
    health = client.health_check()
    if not health:
        print("서버가 응답하지 않습니다!")
        return

    print(f"서버 상태: {health['status']}")
    print(f"모델 로드됨: {health['model_loaded']}")
    print(f"GPU 사용 가능: {health['gpu_available']}")

    # 음성 생성
    print(f"\n음성 생성 중...")
    print(f"텍스트: {args.text}")
    print(f"모델: {args.model}")

    if args.voice:
        print(f"음성 클로닝 사용: {args.voice}")
        result = client.generate_with_voice(
            text=args.text,
            voice_file_path=args.voice,
            model=args.model,
            seed=args.seed,
            diffusion_steps=args.steps,
            cfg_scale=args.cfg_scale
        )
    else:
        print("합성 음성 사용")
        result = client.generate_speech(
            text=args.text,
            model=args.model,
            seed=args.seed,
            diffusion_steps=args.steps,
            cfg_scale=args.cfg_scale
        )

    if not result:
        print("생성 시작 실패!")
        return

    task_id = result["task_id"]
    print(f"작업 ID: {task_id}")

    # 완료 대기
    print("\n완료 대기 중...")
    final_status = client.wait_for_completion(task_id)

    if final_status and final_status["status"] == "completed":
        print("\n생성 완료!")

        # 오디오 다운로드
        print(f"{args.output}로 오디오 다운로드 중...")
        if client.download_audio(task_id, args.output):
            print(f"✅ 오디오 저장 성공: {args.output}")

            # 결과 정보 표시
            if final_status.get("result"):
                duration = final_status["result"].get("duration_seconds")
                if duration:
                    print(f"오디오 길이: {duration:.2f}초")
        else:
            print("❌ 오디오 다운로드 실패")
    else:
        print("❌ 생성 실패 또는 시간 초과")

def example_batch_processing():
    """여러 텍스트 배치 처리 예제"""
    client = CusVoiceModelClient()

    texts = [
        "첫 번째 예제 텍스트입니다.",
        "두 번째 텍스트에는 무음이 있습니다. [pause:1000] 들리셨나요?",
        "세 번째 텍스트는 긴 내용 처리 기능을 보여줍니다.",
        "마지막으로, 네 번째 텍스트로 배치 예제를 마칩니다. [pause] 감사합니다!"
    ]

    print("배치 처리 예제...")
    tasks = []

    # 모든 작업 제출
    for i, text in enumerate(texts):
        print(f"작업 {i+1}/{len(texts)} 제출 중...")
        result = client.generate_speech(
            text=text,
            seed=42 + i,  # 각각 다른 시드
            model="CusVoiceModel-1.5B"
        )

        if result:
            tasks.append((result["task_id"], f"batch_output_{i+1}.wav"))
        else:
            print(f"작업 {i+1} 제출 실패")

    # 모든 완료 대기
    print(f"\n{len(tasks)}개 작업 완료 대기 중...")
    for task_id, output_file in tasks:
        print(f"작업 {task_id} 대기 중...")
        status = client.wait_for_completion(task_id)

        if status and status["status"] == "completed":
            if client.download_audio(task_id, output_file):
                print(f"✅ 다운로드 완료: {output_file}")
            else:
                print(f"❌ 다운로드 실패: {output_file}")
        else:
            print(f"❌ 작업 실패: {task_id}")

def example_voice_cloning():
    """음성 클로닝 기능 예제"""
    client = CusVoiceModelClient()

    # 참조 음성 파일을 제공해야 합니다
    reference_voice = "reference_voice.wav"

    if not Path(reference_voice).exists():
        print(f"참조 음성 파일을 찾을 수 없음: {reference_voice}")
        print("클로닝 예제를 위해 참조 음성 파일을 제공해주세요")
        return

    text = "음성 클로닝 예제입니다. AI가 참조 음성을 모방하려고 시도합니다."

    print("음성 클로닝 예제...")
    result = client.generate_with_voice(
        text=text,
        voice_file_path=reference_voice,
        model="CusVoiceModel-1.5B",
        seed=42
    )

    if result:
        task_id = result["task_id"]
        status = client.wait_for_completion(task_id)

        if status and status["status"] == "completed":
            if client.download_audio(task_id, "cloned_voice_output.wav"):
                print("✅ 음성 클로닝 완료: cloned_voice_output.wav")
            else:
                print("❌ 클론된 음성 다운로드 실패")
        else:
            print("❌ 음성 클로닝 실패")
    else:
        print("❌ 음성 클로닝 시작 실패")

def example_pause_functionality():
    """무음 기능 시연 예제"""
    client = CusVoiceModelClient()

    text_with_pauses = """
    무음 기능 시연에 오신 것을 환영합니다. [pause]
    이것은 1초 무음입니다. [pause:1500]
    방금 것은 1.5초 무음이었습니다. [pause:500]
    이것은 0.5초 무음이었습니다. [pause:2000]
    그리고 방금 것은 2초 무음이었습니다. 들어주셔서 감사합니다!
    """

    print("무음 기능 예제...")
    result = client.generate_speech(
        text=text_with_pauses,
        model="CusVoiceModel-1.5B",
        seed=42
    )

    if result:
        task_id = result["task_id"]
        status = client.wait_for_completion(task_id)

        if status and status["status"] == "completed":
            if client.download_audio(task_id, "pause_example.wav"):
                print("✅ 무음 예제 완료: pause_example.wav")
                if status.get("result"):
                    duration = status["result"].get("duration_seconds")
                    if duration:
                        print(f"무음 포함 총 길이: {duration:.2f}초")
            else:
                print("❌ 무음 예제 다운로드 실패")
        else:
            print("❌ 무음 예제 실패")
    else:
        print("❌ 무음 예제 시작 실패")

if __name__ == "__main__":
    # 명령줄 인자 확인
    import sys

    if len(sys.argv) > 1:
        # 메인 CLI 인터페이스 실행
        main()
    else:
        # 예제 실행
        print("CusVoiceModel API 예제 실행 중...")
        print("CLI 옵션은 --help 사용, 또는 아래 예제들을 실행:\n")

        # 기본 예제
        print("=== 기본 음성 생성 ===")
        main_args = argparse.Namespace(
            url="http://localhost:8090",
            text="안녕하세요, CusVoiceModel API 기본 예제입니다.",
            model="CusVoiceModel-1.5B",
            voice=None,
            output="basic_example.wav",
            seed=42,
            steps=20,
            cfg_scale=1.3
        )

        # 예제를 위해 sys.argv 임시 교체
        original_argv = sys.argv[:]
        sys.argv = ["client_example.py", "--text", main_args.text]
        try:
            main()
        except:
            pass
        finally:
            sys.argv = original_argv

        print("\n=== 무음 기능 ===")
        example_pause_functionality()

        print("\n=== 배치 처리 ===")
        example_batch_processing()

        print("\n=== 음성 클로닝 (참조 파일이 있는 경우) ===")
        example_voice_cloning()

        print("\n예제 완료! 생성된 오디오 파일들을 확인해보세요.")