Added revisions to the translation app

2026-01-12 08:52:38 -05:00
parent 64add85920
commit 8181bebaa0
44 changed files with 993 additions and 375 deletions
@@ -0,0 +1,168 @@
+import whisper
+import os
+import sys
+import subprocess
+import torch
+
+def check_gpu_health():
+    """
+    Performs a robust check for GPU availability and prints detailed troubleshooting 
+    info if issues are detected, specific to Bazzite/VS Code environments.
+    """
+    print("Checking GPU health...")
+    
+    # 1. Check if the OS/Driver sees the GPU
+    nvidia_smi_ok = False
+    try:
+        subprocess.run(["nvidia-smi"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
+        nvidia_smi_ok = True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        nvidia_smi_ok = False
+
+    # 2. Check if PyTorch sees the GPU
+    torch_cuda_ok = torch.cuda.is_available()
+
+    if torch_cuda_ok:
+        print(f"✅ GPU is accessible: {torch.cuda.get_device_name(0)}")
+        print(f"   CUDA Version: {torch.version.cuda}")
+        return True
+
+    # --- Troubleshooting Block ---
+    print("\n⚠️  WARNING: GPU not detected by PyTorch. Falling back to CPU.")
+    print("   Transcription will be significantly slower.\n")
+
+    print("--- Diagnostic Report ---")
+    if nvidia_smi_ok:
+        print("1. [OK] 'nvidia-smi' command works. The system driver is installed and visible.")
+        print("2. [FAIL] PyTorch cannot see the GPU.")
+        print("   -> Likely Cause: You might have installed the CPU-only version of PyTorch.")
+        print("   -> Solution: Reinstall PyTorch with CUDA support:")
+        print("      pip uninstall torch torchvision torchaudio")
+        print("      pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118")
+    else:
+        print("1. [FAIL] 'nvidia-smi' command failed or not found.")
+        print("   -> Likely Cause: Nvidia drivers are missing, or the container/sandbox cannot access the GPU.")
+        
+        print("\n   --- Bazzite / VS Code / Container Specific Checks ---")
+        print("   a. If you are running inside a dev container (DevBox/Distrobox/Toolbox):")
+        print("      Ensure the container was created with nvidia support.")
+        print("      (Bazzite usually handles this for 'distrobox', but check your config).")
+        print("   b. If you are using VS Code Flatpak:")
+        print("      Flatpak might be restricting access. Check Flatseal permissions for VS Code.")
+        print("   c. Driver Check:")
+        print("      Run 'rpm -qa | grep nvidia' in your host terminal to verify drivers are installed.")
+
+    print("-------------------------\n")
+    return False
+
+def get_vram_gb():
+    """Returns the total VRAM in GB of the first CUDA device, or 0 if no CUDA."""
+    if not torch.cuda.is_available():
+        return 0
+    try:
+        # returns bytes
+        total_mem = torch.cuda.get_device_properties(0).total_memory 
+        return total_mem / (1024 ** 3)
+    except Exception:
+        return 0
+
+def get_optimal_model_size():
+    """
+    Determines the best Whisper model based on available VRAM.
+    Rough estimates for VRAM usage (fp16):
+    - large: ~10 GB
+    - medium: ~5 GB
+    - small: ~2 GB
+    - base: ~1 GB
+    - tiny: ~1 GB
+    """
+    vram = get_vram_gb()
+    if vram == 0:
+        return "base"
+    
+    print(f"Detected GPU with {vram:.2f} GB VRAM.")
+
+    if vram >= 11:
+        return "large"
+    elif vram >= 6:
+        return "medium"
+    elif vram >= 3:
+        return "small"
+    else:
+        return "base"
+
+def format_timestamp(seconds: float):
+    """Converts seconds to SRT timestamp format (HH:MM:SS,mmm)."""
+    whole_seconds = int(seconds)
+    milliseconds = int((seconds - whole_seconds) * 1000)
+    
+    hours = whole_seconds // 3600
+    minutes = (whole_seconds % 3600) // 60
+    seconds = whole_seconds % 60
+    
+    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
+
+def save_as_srt(result, output_path):
+    """Saves the Whisper transcription result as an SRT file."""
+    with open(output_path, "w", encoding="utf-8") as f:
+        for i, segment in enumerate(result["segments"], start=1):
+            start = format_timestamp(segment["start"])
+            end = format_timestamp(segment["end"])
+            text = segment["text"].strip()
+            
+            f.write(f"{i}\n")
+            f.write(f"{start} --> {end}\n")
+            f.write(f"{text}\n\n")
+    print(f"SRT saved to: {output_path}")
+
+def transcribe_audio(audio_path, model_size="auto", language=None):
+    """
+    Transcribes an audio file using OpenAI's Whisper model.
+
+    Args:
+        audio_path (str): Path to the input audio file.
+        model_size (str): Size of the Whisper model to use. If "auto", selects based on VRAM.
+        language (str, optional): Language code (e.g., "en", "fr", "es"). If None, auto-detects.
+
+    Returns:
+        dict: The full transcription result containing segments and text.
+    """
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Run health check once
+    check_gpu_health()
+
+    # Determine model size if auto
+    if model_size == "auto":
+        model_size = get_optimal_model_size()
+        print(f"Auto-selected model: '{model_size}'")
+
+    print(f"Loading Whisper model ('{model_size}')...")
+    
+    # Check for GPU availability
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    try:
+        model = whisper.load_model(model_size, device=device)
+    except RuntimeError as e:
+        if "out of memory" in str(e).lower():
+            print("Error: GPU Out of Memory. Try using a smaller model size.")
+        else:
+            print(f"Error loading model: {e}")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        sys.exit(1)
+
+    print(f"Transcribing {audio_path}...")
+    try:
+        # fp16=False is needed for CPU, but we can let whisper handle defaults usually.
+        # language=None allows auto-detection.
+        result = model.transcribe(audio_path, language=language)
+        print("Transcription complete.")
+        return result
+    except Exception as e:
+        print(f"Error during transcription: {e}")
+        sys.exit(1)