lattifai 0.2.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

lattifai/bin/align.py

@@ -1,17 +1,40 @@
+import asyncio
+import os
+from pathlib import Path
+
 import click
 import colorful
 from lhotse.utils import Pathlike
 
 from lattifai.bin.cli_base import cli
+from lattifai.client import AsyncLattifAI, LattifAI
+from lattifai.io import INPUT_SUBTITLE_FORMATS, OUTPUT_SUBTITLE_FORMATS
 
 
 @cli.command()
 @click.option(
     '-F',
     '--input_format',
-    type=click.Choice(['srt', 'vtt', 'ass', 'txt', 'auto'], case_sensitive=False),
+    '--input-format',
+    type=click.Choice(INPUT_SUBTITLE_FORMATS, case_sensitive=False),
     default='auto',
-    help='Input Subtitle format.',
+    help='Input subtitle format.',
+)
+@click.option(
+    '-S',
+    '--split-sentence',
+    '--split_sentence',
+    is_flag=True,
+    default=False,
+    help='Re-segment subtitles by semantics.',
+)
+@click.option(
+    '-W',
+    '--word-level',
+    '--word_level',
+    is_flag=True,
+    default=False,
+    help='Include word-level alignment timestamps in output (for JSON, TextGrid, and subtitle formats).',
 )
 @click.option(
     '-D',
@@ -21,17 +44,22 @@ from lattifai.bin.cli_base import cli
     help='Device to use for inference.',
 )
 @click.option(
-    '-M', '--model_name_or_path', type=str, default='Lattifai/Lattice-1-Alpha', help='Lattifai model name or path'
+    '-M',
+    '--model-name-or-path',
+    '--model_name_or_path',
+    type=str,
+    default='Lattifai/Lattice-1-Alpha',
+    help='Model name or path for alignment.',
 )
 @click.option(
-    '-S',
-    '--split_sentence',
-    is_flag=True,
-    default=False,
-    help='Re-segment subtitles by semantics.',
+    '--api-key',
+    '--api_key',
+    type=str,
+    default=None,
+    help='API key for LattifAI.',
 )
 @click.argument(
-    'input_audio_path',
+    'input_media_path',
     type=click.Path(exists=True, dir_okay=False),
 )
 @click.argument(
@@ -43,24 +71,228 @@ from lattifai.bin.cli_base import cli
     type=click.Path(allow_dash=True),
 )
 def align(
-    input_audio_path: Pathlike,
+    input_media_path: Pathlike,
     input_subtitle_path: Pathlike,
     output_subtitle_path: Pathlike,
     input_format: str = 'auto',
+    split_sentence: bool = False,
+    word_level: bool = False,
     device: str = 'cpu',
     model_name_or_path: str = 'Lattifai/Lattice-1-Alpha',
+    api_key: str = None,
+):
+    """
+    Command used to align media(audio/video) with subtitles
+    """
+    try:
+        client = LattifAI(model_name_or_path=model_name_or_path, device=device, api_key=api_key)
+        client.alignment(
+            input_media_path,
+            input_subtitle_path,
+            format=input_format.lower(),
+            split_sentence=split_sentence,
+            return_details=word_level,
+            output_subtitle_path=output_subtitle_path,
+        )
+        click.echo(colorful.green(f'✅ Alignment completed successfully: {output_subtitle_path}'))
+    except Exception as e:
+        from lattifai.errors import LattifAIError
+
+        # Display error message
+        if isinstance(e, LattifAIError):
+            click.echo(colorful.red('❌ Alignment failed:'))
+            click.echo(e.get_message())
+            # Show support info
+            click.echo(e.get_support_info())
+        else:
+            click.echo(colorful.red(f'❌ Alignment failed: {str(e)}'))
+
+        raise click.ClickException('Alignment failed')
+
+
+@cli.command()
+@click.option(
+    '-M',
+    '--media-format',
+    '--media_format',
+    type=click.Choice(
+        [
+            # Audio formats
+            'mp3',
+            'wav',
+            'm4a',
+            'aac',
+            'flac',
+            'ogg',
+            'opus',
+            'aiff',
+            # Video formats
+            'mp4',
+            'webm',
+            'mkv',
+            'avi',
+            'mov',
+        ],
+        case_sensitive=False,
+    ),
+    default='mp3',
+    help='Media format for YouTube download (audio or video).',
+)
+@click.option(
+    '-S',
+    '--split-sentence',
+    '--split_sentence',
+    is_flag=True,
+    default=False,
+    help='Re-segment subtitles by semantics.',
+)
+@click.option(
+    '-W',
+    '--word-level',
+    '--word_level',
+    is_flag=True,
+    default=False,
+    help='Include word-level alignment timestamps in output (for JSON, TextGrid, and subtitle formats).',
+)
+@click.option(
+    '-O',
+    '--output-dir',
+    '--output_dir',
+    type=click.Path(file_okay=False, dir_okay=True, writable=True),
+    default='.',
+    help='Output directory (default: current directory).',
+)
+@click.option(
+    '-D',
+    '--device',
+    type=click.Choice(['cpu', 'cuda', 'mps'], case_sensitive=False),
+    default='cpu',
+    help='Device to use for inference.',
+)
+@click.option(
+    '-M',
+    '--model-name-or-path',
+    '--model_name_or_path',
+    type=str,
+    default='Lattifai/Lattice-1-Alpha',
+    help='Model name or path for alignment.',
+)
+@click.option(
+    '--api-key',
+    '--api_key',
+    type=str,
+    default=None,
+    help='API key for LattifAI.',
+)
+@click.option(
+    '--gemini-api-key',
+    '--gemini_api_key',
+    type=str,
+    default=None,
+    help='Gemini API key for transcription fallback when subtitles are unavailable.',
+)
+@click.option(
+    '-F',
+    '--output-format',
+    '--output_format',
+    type=click.Choice(OUTPUT_SUBTITLE_FORMATS, case_sensitive=False),
+    default='vtt',
+    help='Subtitle output format.',
+)
+@click.argument(
+    'yt_url',
+    type=str,
+)
+def youtube(
+    yt_url: str,
+    media_format: str = 'mp3',
     split_sentence: bool = False,
+    word_level: bool = False,
+    output_dir: str = '.',
+    device: str = 'cpu',
+    model_name_or_path: str = 'Lattifai/Lattice-1-Alpha',
+    api_key: str = None,
+    gemini_api_key: str = None,
+    output_format: str = 'vtt',
 ):
     """
-    Command used to align audio with subtitles
+    Download media and subtitles from YouTube for further alignment.
     """
-    from lattifai import LattifAI
-
-    client = LattifAI(model_name_or_path=model_name_or_path, device=device)
-    client.alignment(
-        input_audio_path,
-        input_subtitle_path,
-        format=input_format.lower(),
-        split_sentence=split_sentence,
-        output_subtitle_path=output_subtitle_path,
-    )
+    from lattifai.workflows.gemini import GeminiTranscriber
+    from lattifai.workflows.youtube import YouTubeDownloader, YouTubeSubtitleAgent
+
+    # Get Gemini API key
+    gemini_key = gemini_api_key or os.getenv('GEMINI_API_KEY')
+    if not gemini_key:
+        click.echo(
+            colorful.red(
+                '❌ Gemini API key is required. Set GEMINI_API_KEY environment variable or use --gemini-api-key option.'
+            )
+        )
+        raise click.ClickException('Missing Gemini API key')
+
+    async def _process():
+        # Initialize components with their configuration (only config, not runtime params)
+        downloader = YouTubeDownloader()
+        transcriber = GeminiTranscriber(api_key=gemini_key)
+        aligner = AsyncLattifAI(api_key=api_key, model_name_or_path=model_name_or_path, device=device)
+
+        # Create agent with initialized components
+        agent = YouTubeSubtitleAgent(
+            downloader=downloader,
+            transcriber=transcriber,
+            aligner=aligner,
+            max_retries=0,
+        )
+
+        result = await agent.process_youtube_url(
+            url=yt_url,
+            output_dir=output_dir,
+            media_format=media_format,
+            force_overwrite=False,
+            output_format=output_format,
+            split_sentence=split_sentence,
+            word_level=word_level,
+        )
+        return result
+
+    try:
+        result = asyncio.run(_process())
+
+        # Display results
+        click.echo(colorful.green('✅ Processing completed!'))
+        click.echo()
+
+        # Show metadata
+        metadata = result.get('metadata', {})
+        if metadata:
+            click.echo(f'🎬    Title: {metadata.get("title", "Unknown")}')
+            click.echo(f'⏱️  Duration: {metadata.get("duration", 0)} seconds')
+            click.echo()
+
+        # Show exported files
+        exported_files = result.get('exported_files', {})
+        if exported_files:
+            click.echo(colorful.green('📄 Generated subtitle files:'))
+            for format_name, file_path in exported_files.items():
+                click.echo(f'  {format_name}: {file_path}')
+            click.echo()
+
+        # Show subtitle count
+        subtitle_count = result.get('subtitle_count', 0)
+        click.echo(f'📝 Generated {subtitle_count} subtitle segments')
+
+    except Exception as e:
+        from lattifai.errors import LattifAIError
+
+        # Extract error message without support info (to avoid duplication)
+        if isinstance(e, LattifAIError):
+            # Use the get_message() method which includes proper formatting
+            click.echo(colorful.red('❌ Failed to process YouTube URL:'))
+            click.echo(e.get_message())
+            # Show support info once at the end
+            click.echo(e.get_support_info())
+        else:
+            click.echo(colorful.red(f'❌ Failed to process YouTube URL: {str(e)}'))
+
+        raise click.ClickException('Processing failed')

lattifai/bin/cli_base.py

@@ -17,3 +17,8 @@ def cli():
         format='%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s',
         level=logging.INFO,
     )
+
+    import os
+
+    os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
+    os.environ['TOKENIZERS_PARALLELISM'] = 'FALSE'

lattifai/bin/subtitle.py

@@ -1,12 +1,16 @@
+import asyncio
+from pathlib import Path
+
 import click
 from lhotse.utils import Pathlike
 
 from lattifai.bin.cli_base import cli
+from lattifai.io import SUBTITLE_FORMATS
 
 
 @cli.group()
 def subtitle():
-    """Group of commands used to convert subtitle format."""
+    """Commands for subtitle format conversion and management."""
     pass
 
 
@@ -26,7 +30,181 @@ def convert(
     """
     Convert subtitle file to another format.
     """
-    import pysubs2
+    if str(output_subtitle_path).lower().endswith('.TextGrid'.lower()):
+        from lattifai.io import SubtitleIO
+
+        alignments = SubtitleIO.read(input_subtitle_path)
+        SubtitleIO.write(alignments, output_subtitle_path)
+    else:
+        import pysubs2
+
+        subtitle = pysubs2.load(input_subtitle_path)
+
+        subtitle.save(output_subtitle_path)
+
+
+@subtitle.command()
+@click.argument('url', type=str, required=True)
+@click.option(
+    '--output-dir',
+    '--output_dir',
+    '-o',
+    type=click.Path(file_okay=False, dir_okay=True),
+    default='.',
+    help='Output directory for downloaded subtitle files (default: current directory).',
+)
+@click.option(
+    '--output-format',
+    '--output_format',
+    '-f',
+    type=click.Choice(SUBTITLE_FORMATS + ['best'], case_sensitive=False),
+    default='best',
+    help='Preferred subtitle format to download (default: best available).',
+)
+@click.option('--force-overwrite', '-F', is_flag=True, help='Overwrite existing files without prompting.')
+@click.option(
+    '--lang',
+    '-l',
+    '-L',
+    '--subtitle-lang',
+    '--subtitle_lang',
+    type=str,
+    help='Specific subtitle language/track to download (e.g., "en").',
+)
+def download(
+    url: str,
+    output_dir: str,
+    output_format: str,
+    force_overwrite: bool,
+    lang: str,
+):
+    """
+    Download subtitles from YouTube URL using yt-dlp.
+
+    URL should be a valid YouTube URL (e.g., https://www.youtube.com/watch?v=VIDEO_ID).
+    """
+    # Import here to avoid circular imports and keep startup fast
+    from lattifai.workflows.youtube import YouTubeDownloader
+
+    # Validate URL format
+    if not _is_valid_youtube_url(url):
+        click.echo(f'Error: Invalid YouTube URL format: {url}', err=True)
+        click.echo('Please provide a valid YouTube URL (e.g., https://www.youtube.com/watch?v=VIDEO_ID)', err=True)
+        raise click.Abort()
+
+    # Convert relative path to absolute
+    output_path = Path(output_dir).resolve()
+
+    # Create output directory if it doesn't exist
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    click.echo(f'Downloading subtitles from: {url}')
+    click.echo(f'          Output directory: {output_path}')
+    click.echo(f'         Preferred format: {output_format}')
+    if lang:
+        click.echo(f'       Subtitle language: {lang}')
+    else:
+        click.echo('       Subtitle language: All available')
+
+    # Initialize downloader and download
+    downloader = YouTubeDownloader()
+
+    async def download_subtitles():
+        try:
+            result = await downloader.download_subtitles(
+                url=url,
+                output_dir=str(output_path),
+                force_overwrite=force_overwrite,
+                subtitle_lang=lang,
+            )
+
+            if result:
+                click.echo('✅ Subtitles downloaded successfully!')
+                return result
+            else:
+                click.echo('⚠️  No subtitles available for this video')
+                return None
+
+        except Exception as e:
+            click.echo(f'❌ Error downloading subtitles: {str(e)}', err=True)
+            raise click.Abort()
+
+    # Run the async function
+    result = asyncio.run(download_subtitles())
+
+    if result:
+        if result == 'gemini':
+            click.echo('✨ Gemini transcription selected (use the agent command to transcribe)')
+        else:
+            click.echo(f'📄 Subtitle file saved to: {result}')
+
+
+@subtitle.command()
+@click.argument('url', type=str, required=True)
+def list_subs(url: str):
+    """
+    List available subtitle tracks for a YouTube video.
+
+    URL should be a valid YouTube URL (e.g., https://www.youtube.com/watch?v=VIDEO_ID)
+    """
+    # Import here to avoid circular imports and keep startup fast
+    from lattifai.workflows.youtube import YouTubeDownloader
+
+    # Validate URL format
+    if not _is_valid_youtube_url(url):
+        click.echo(f'Error: Invalid YouTube URL format: {url}', err=True)
+        click.echo('Please provide a valid YouTube URL (e.g., https://www.youtube.com/watch?v=VIDEO_ID)', err=True)
+        raise click.Abort()
+
+    click.echo(f'Listing available subtitles for: {url}')
+
+    # Initialize downloader
+    downloader = YouTubeDownloader()
+
+    async def list_available_subtitles():
+        try:
+            result = await downloader.list_available_subtitles(url)
+
+            if result:
+                click.echo('📋 Available subtitle tracks:')
+                for subtitle_info in result:
+                    click.echo(f'  🎬 Language: {subtitle_info["language"]} - {subtitle_info["name"]}')
+                    click.echo(f'     📄 Formats: {", ".join(subtitle_info["formats"])}')
+                    click.echo()
+
+                click.echo('💡 To download a specific track, use:')
+                click.echo(f'   lattifai subtitle download "{url}" --lang <language_code>')
+                click.echo('   Example: lattifai subtitle download "{}" --lang en-JkeT_87f4cc'.format(url))
+            else:
+                click.echo('⚠️  No subtitles available for this video')
+
+        except Exception as e:
+            click.echo(f'❌ Error listing subtitles: {str(e)}', err=True)
+            raise click.Abort()
+
+    # Run the async function
+    asyncio.run(list_available_subtitles())
+
+
+def _is_valid_youtube_url(url: str) -> bool:
+    """
+    Validate if the URL is a valid YouTube URL format.
+
+    Supports various YouTube URL formats:
+    - https://www.youtube.com/watch?v=VIDEO_ID
+    - https://youtu.be/VIDEO_ID
+    - https://www.youtube.com/shorts/VIDEO_ID
+    - https://m.youtube.com/watch?v=VIDEO_ID
+    """
+    import re
+
+    patterns = [
+        r'(?:youtube\.com/watch\?v=|youtu\.be/|youtube\.com/shorts/)([a-zA-Z0-9_-]{11})',
+        r'youtube\.com/embed/([a-zA-Z0-9_-]{11})',
+        r'youtube\.com/v/([a-zA-Z0-9_-]{11})',
+    ]
 
-    subtitle = pysubs2.load(input_subtitle_path)
-    subtitle.save(output_subtitle_path)
+    for pattern in patterns:
+        if re.search(pattern, url):
+            return True
+    return False