document_to_podcast_kfp.py

pipeline.py

from kfp import dsl, compiler, kubernetes


@dsl.component(packages_to_install=['requests'])
def download_document(document_url: str, document_path: str):
    """Download a document from the given URL.

    Args:
        document_url (str): The URL of the document to download.
        document_path (str): The local path where the document will be saved.
    """
    import requests
    import logging

    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

    logger.info(f"Downloading document from {document_url} to {document_path}...")

    response = requests.get(document_url)
    response.raise_for_status()  # Raise an error for bad responses
 
    with open(document_path, 'wb') as f:
        f.write(response.content)

    logger.info(f"Document downloaded successfully to {document_path}")
    
    
@dsl.container_component
def process_data(input_path: str, output_path: str):
    """Preprocess the raw input document.

    Args:
        input_path: Path to the input document.
        output_path: Path to the output document.
    """
    return dsl.ContainerSpec(
        image='dpoulopoulos/kfp-discovery-transformer:v0.0.2',
        args=["python", "transformer.py", "--input", input_path, "--output", output_path]
    )
  

@dsl.container_component
def text_to_script(
    input_path: str,
    output_path: str,
    host_voice_profile: str,
    cohost_voice_profile: str):
    """Convert the preprocessed document to a script.

    Args:
        input_path: Path to the input document.
        output_path: Path to the output document.
        voice_profiles: Tuple of voice profiles to use for the conversion.
    """
    return dsl.ContainerSpec(
        image='dpoulopoulos/kfp-discovery-ttt-amd64:v0.0.2',
        args=[
            "python", "text_to_script.py",
            "--input", input_path,
            "--output", output_path,
            "--voice-profiles", host_voice_profile, cohost_voice_profile
        ]
    )
  
  
@dsl.container_component
def text_to_speech(
    input_path: str,
    output_path: str,
    host_voice_profile: str,
    cohost_voice_profile: str):
    """Convert the preprocessed document to a script.

    Args:
        input_path: Path to the input script.
        output_path: Path to the output audio file.
        voice_profiles: Tuple of voice profiles to use for the conversion.
    """
    return dsl.ContainerSpec(
        image='dpoulopoulos/kfp-discovery-sts-amd64:v0.0.2',
        args=[
            "python", "script_to_speech.py",
            "--input", input_path,
            "--output", output_path,
            "--voice-profiles", host_voice_profile, cohost_voice_profile
        ]
    )
  
  
@dsl.pipeline
def document_to_podcast(
    document_url: str,
    input_path: str,
    proccesed_text_path: str,
    script_path: str,
    audio_path: str,
    host_voice_profile: str = 'af_sarah',
    cohost_voice_profile: str = 'am_michael'
) -> None:
    """Convert a document to a podcast.

    This pipeline downloads a document, processes it, converts it to a script,
    and finally converts the script to speech.

    Args:
        input_path: Path to the input document.
        output_path: Path to the output document.
        script_path: Path to the output script.
        host_voice_profile: Voice profile for the host.
        cohost_voice_profile: Voice profile for the co-host.
    """
    pvc = kubernetes.CreatePVC(
        pvc_name='doc-to-podcast-pvc',
        access_modes=['ReadWriteMany'],
        size='1.0Gi',
        storage_class_name='shared-vast'
    )

    download_document_step = download_document(
        document_url=document_url,
        document_path=input_path
    ).after(pvc)
    download_document_step.set_caching_options(False)

    process_data_step = process_data(
        input_path=input_path,
        output_path=proccesed_text_path
    ).after(download_document_step)
    process_data_step.set_caching_options(False)

    text_to_script_step = text_to_script(
        input_path=proccesed_text_path,
        output_path=script_path,
        host_voice_profile=host_voice_profile,
        cohost_voice_profile=cohost_voice_profile
    ).after(process_data_step)
    text_to_script_step.set_accelerator_type("nvidia.com/gpu")
    text_to_script_step.set_accelerator_limit(1)
    text_to_script_step.set_caching_options(False)

    text_to_speech_step = text_to_speech(
        input_path=script_path,
        output_path=audio_path,
        host_voice_profile=host_voice_profile,
        cohost_voice_profile=cohost_voice_profile
    ).after(text_to_script_step)
    text_to_speech_step.set_accelerator_type("nvidia.com/gpu")
    text_to_speech_step.set_accelerator_limit(1)
    text_to_speech_step.set_caching_options(False)

    kubernetes.mount_pvc(
        download_document_step,
        pvc_name=pvc.outputs['name'],
        mount_path='/mnt/data'
    )
    kubernetes.mount_pvc(
        process_data_step,
        pvc_name=pvc.outputs['name'],
        mount_path='/mnt/data'
    )
    kubernetes.mount_pvc(
        text_to_script_step,
        pvc_name=pvc.outputs['name'],
        mount_path='/mnt/data'
    )
    kubernetes.mount_pvc(
        text_to_speech_step,
        pvc_name=pvc.outputs['name'],
        mount_path='/mnt/data'
    )