import os import json import requests import subprocess import shutil import time import sys import threading from typing import Dict, List, Optional, Any from fastapi import FastAPI, HTTPException from fastapi.responses import JSONResponse import uvicorn import torch import librosa from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline # Fix Unicode encoding for Windows if sys.platform == 'win32': import io sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8') # Initialize FastAPI app app = FastAPI(title="Audio Transcriber", description="Audio transcription and upload service") # ==== CONFIGURATION ==== # The new backend URL for state management and transcription upload # It is now read from an environment variable, falling back to the default if not set. BACKEND_URL = os.environ.get("BACKEND_URL", "https://samfredoly-acp.hf.space") # The original Hugging Face repo IDs are still needed for fetching the audio files # and the reference file list, as the backend only handles transcription storage. SOURCE_REPO_ID = "Samfredoly/BG_Vid" # Fetch audio files from here TARGET_REPO_ID = "samfred2/A_Text" # Target repo ID is now a constant for the backend REFERENCE_REPO_ID = "Fred808/BG3" # Reference repo to match audio filenames # Path Configuration DOWNLOAD_FOLDER = "downloads_audio" TRANSCRIPTIONS_FOLDER = "transcriptions" LOCAL_STATE_FOLDER = ".state_audio" os.makedirs(DOWNLOAD_FOLDER, exist_ok=True) os.makedirs(TRANSCRIPTIONS_FOLDER, exist_ok=True) os.makedirs(LOCAL_STATE_FOLDER, exist_ok=True) # Whisper Model Setup (using transformers) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32 WHISPER_MODEL_ID = f"openai/whisper-small" # Global model cache _whisper_model = None _whisper_processor = None _whisper_pipeline = None def get_whisper_pipeline(): """Get or initialize the Whisper pipeline.""" global _whisper_model, _whisper_processor, _whisper_pipeline if _whisper_pipeline is not None: return _whisper_pipeline try: log_message(f"Loading Whisper model {WHISPER_MODEL_ID}...", "INFO") model = AutoModelForSpeechSeq2Seq.from_pretrained( WHISPER_MODEL_ID, torch_dtype=TORCH_DTYPE, low_cpu_mem_usage=True, use_safetensors=True ) model = model.to(DEVICE) processor = AutoProcessor.from_pretrained(WHISPER_MODEL_ID) _whisper_pipeline = pipeline( "automatic-speech-recognition", model=model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, torch_dtype=TORCH_DTYPE, device=DEVICE ) log_message(f"βœ… Whisper model loaded successfully on {DEVICE.upper()}", "INFO") return _whisper_pipeline except Exception as e: log_message(f"❌ Failed to load Whisper model: {str(e)}", "ERROR") raise # State Files FAILED_FILES_LOG = "failed_audio_files.log" HF_STATE_FILE = "processing_audio_state.json" # This is the filename the backend uses # Processing Parameters PROCESSING_DELAY = 2 MAX_RETRIES = 3 MIN_FREE_SPACE_GB = 1 WHISPER_MODEL = "small" # Whisper model size # NOTE: The Hugging Face API is still required for listing files in SOURCE_REPO_ID and REFERENCE_REPO_ID from huggingface_hub import HfApi, hf_hub_url HF_TOKEN = os.environ.get("HF_TOKEN", "") hf_api = HfApi(token=HF_TOKEN) # Global State processing_status = { "is_running": False, "current_file": None, "total_files": 0, "processed_files": 0, "failed_files": 0, "transcribed_files": 0, "last_update": None, "logs": [] } def log_message(message: str, level: str = "INFO"): """Log messages with timestamp""" timestamp = time.strftime("%Y-%m-%d %H:%M:%S") log_entry = f"[{timestamp}] {level}: {message}" print(log_entry) processing_status["logs"].append(log_entry) processing_status["last_update"] = timestamp if len(processing_status["logs"]) > 100: processing_status["logs"] = processing_status["logs"][-100:] def log_failed_file(filename: str, error: str): """Log failed files to persistent file""" with open(FAILED_FILES_LOG, "a") as f: f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} - {filename}: {error}\n") def get_disk_usage(path: str) -> Dict[str, float]: """Get disk usage statistics in GB""" statvfs = os.statvfs(path) total = statvfs.f_frsize * statvfs.f_blocks / (1024**3) free = statvfs.f_frsize * statvfs.f_bavail / (1024**3) used = total - free return {"total": total, "free": free, "used": used} def check_disk_space(path: str = ".") -> bool: """Check if there's enough disk space""" disk_info = get_disk_usage(path) if disk_info["free"] < MIN_FREE_SPACE_GB: log_message(f'⚠️ Low disk space: {disk_info["free"]:.2f}GB free, {disk_info["used"]:.2f}GB used') return False return True def cleanup_temp_files(): """Clean up temporary files to free space""" log_message("🧹 Cleaning up temporary files...", "INFO") current_file = processing_status.get("current_file") for file in os.listdir(DOWNLOAD_FOLDER): if file != current_file and file.endswith((".wav", ".mp3")): try: os.remove(os.path.join(DOWNLOAD_FOLDER, file)) log_message(f"πŸ—‘οΈ Removed old download: {file}", "INFO") except: pass # Helper function to save state locally def save_json_state(file_path: str, data: Dict[str, Any]): """Save state to JSON file""" with open(file_path, "w") as f: json.dump(data, f, indent=2) # --- NEW API FUNCTIONS FOR STATE MANAGEMENT AND UPLOAD --- def download_state_from_api() -> Dict[str, Any]: """Downloads the state file from the backend API.""" url = f"{BACKEND_URL}/state/" default_state = {"next_download_index": 0, "file_states": {}} try: response = requests.get(url, timeout=10) response.raise_for_status() # The API returns {"state": {...}} state_data = response.json().get("state", default_state) # Ensure the structure is correct (migration logic from original load_json_state) if "file_states" not in state_data or not isinstance(state_data["file_states"], dict): state_data["file_states"] = {} if "next_download_index" not in state_data: state_data["next_download_index"] = 0 log_message(f"βœ… Downloaded state: next_download_index={state_data['next_download_index']}, processed_files={len([f for f,s in state_data['file_states'].items() if s=='processed'])}", "INFO") return state_data except requests.exceptions.RequestException as e: log_message(f"⚠️ Failed to download state from API ({url}): {str(e)}. Starting from default state.", "WARNING") return default_state def upload_state_to_api(state: Dict[str, Any]) -> bool: """ Saves the state locally and uploads it to the backend API's /upload/ endpoint. This simulates the original HF state upload for locking/unlocking. """ local_path = os.path.join(LOCAL_STATE_FOLDER, HF_STATE_FILE) url = f"{BACKEND_URL}/upload/" try: # 1. Save the current state locally save_json_state(local_path, state) # 2. Upload the state file to the backend with open(local_path, "rb") as f: files = {'file': (HF_STATE_FILE, f, 'application/json')} response = requests.post(url, files=files, timeout=30) response.raise_for_status() log_message(f"βœ… Successfully uploaded state file to API: {HF_STATE_FILE}", "INFO") return True except requests.exceptions.HTTPError as e: if hasattr(e, 'response') and e.response.status_code == 409: log_message(f"⚠️ State file already exists on server (409 Conflict) - Treating as success.", "INFO") return True log_message(f"❌ Failed to upload state file to API ({url}): {str(e)}", "ERROR") return False except requests.exceptions.RequestException as e: log_message(f"❌ Failed to upload state file to API ({url}): {str(e)}", "ERROR") return False except Exception as e: log_message(f"❌ An unexpected error occurred during API state upload: {str(e)}", "ERROR") return False def upload_transcription_to_api(json_output_path: str, matched_filename: str) -> bool: """Uploads the transcription JSON file to the backend API's /upload/ endpoint.""" url = f"{BACKEND_URL}/upload/" try: with open(json_output_path, "rb") as f: files = {'file': (os.path.basename(json_output_path), f, 'application/json')} response = requests.post(url, files=files, timeout=30) response.raise_for_status() log_message(f"βœ… Successfully uploaded transcription to API: {os.path.basename(json_output_path)}", "INFO") return True except requests.exceptions.HTTPError as e: if hasattr(e, 'response') and e.response.status_code == 409: log_message(f"⚠️ File already exists on server (409 Conflict) - Treating as success.", "INFO") return True log_message(f"❌ Failed to upload transcription to API ({url}): {str(e)}", "ERROR") return False except requests.exceptions.RequestException as e: log_message(f"❌ Failed to upload transcription to API ({url}): {str(e)}", "ERROR") return False except Exception as e: log_message(f"❌ An unexpected error occurred during API upload: {str(e)}", "ERROR") return False def lock_file_for_processing(wav_filename: str, state: Dict[str, Any]) -> bool: """Marks a file as 'processing' in the state file and uploads the lock via API.""" log_message(f"πŸ”’ Attempting to lock file: {wav_filename} (Marking as 'processing')", "INFO") state["file_states"][wav_filename] = "processing" if upload_state_to_api(state): log_message(f"βœ… Successfully locked file: {wav_filename} via API state upload", "INFO") return True else: log_message(f"❌ Failed to upload lock for file: {wav_filename}. Aborting processing.", "ERROR") # Revert local state change if upload fails if wav_filename in state["file_states"]: del state["file_states"][wav_filename] return False def unlock_file_as_processed(wav_filename: str, state: Dict[str, Any], next_index: int) -> bool: """Marks a file as 'processed', updates the index, and uploads the state via API.""" log_message(f"πŸ”“ Attempting to unlock file: {wav_filename} (Marking as 'processed')", "INFO") state["file_states"][wav_filename] = "processed" state["next_download_index"] = next_index if upload_state_to_api(state): log_message(f"βœ… Successfully unlocked and marked as processed: {wav_filename} via API state upload", "INFO") return True else: log_message(f"❌ Failed to upload final state for file: {wav_filename}.", "ERROR") return False # --- END NEW API FUNCTIONS --- def download_with_retry(url: str, dest_path: str, max_retries: int = 3) -> bool: """Download file with retry logic and disk space checking""" if not check_disk_space(): cleanup_temp_files() if not check_disk_space(): log_message("❌ Insufficient disk space even after cleanup", "ERROR") return False try: os.makedirs(os.path.dirname(dest_path), exist_ok=True) except Exception as e: log_message(f"❌ Failed to create directory for download path {os.path.dirname(dest_path)}: {str(e)}", "ERROR") return False # The original code used HF_TOKEN for authorization headers, which is only needed # if the source repo is private. We keep it for compatibility. headers = {"Authorization": f"Bearer {HF_TOKEN}"} for attempt in range(max_retries): try: with requests.get(url, headers=headers, stream=True) as r: r.raise_for_status() with open(dest_path, "wb") as f: for chunk in r.iter_content(chunk_size=8192): if chunk: f.write(chunk) log_message(f"βœ… Download successful: {os.path.basename(dest_path)}", "INFO") return True except requests.exceptions.RequestException as e: log_message(f"⚠️ Download attempt {attempt + 1}/{max_retries} failed for {url}: {str(e)}", "WARNING") if attempt < max_retries - 1: time.sleep(2 ** attempt) # Exponential backoff else: log_message(f"❌ Download failed after {max_retries} attempts for {url}", "ERROR") return False except Exception as e: log_message(f"❌ An unexpected error occurred during download: {str(e)}", "ERROR") return False return False def get_reference_map(reference_repo_id: str) -> Dict[str, str]: """ Downloads the reference file list from the Hugging Face repo and creates a map from audio filename (without extension) to the reference filename. """ log_message(f"Fetching reference file list from {reference_repo_id}...", "INFO") # This is a placeholder for the actual logic to get the file list. # Assuming the reference repo contains a list of files that match the audio files. # In a real scenario, this would involve listing files in the repo. # For now, we'll assume a simple list of files can be retrieved. try: # Use HfApi to list files in the reference repo repo_files = hf_api.list_repo_files(repo_id=reference_repo_id, repo_type="dataset") reference_map = {} for file in repo_files: # Assuming the reference files are named like 'audio_file_name.txt' # and we want to map the audio file name (e.g., 'audio_file_name.wav') to it. base_name, ext = os.path.splitext(file) if ext.lower() in ['.txt', '.json']: # Only consider text/json files as reference # The key is the audio file name without extension reference_map[base_name] = file log_message(f"βœ… Successfully created reference map with {len(reference_map)} entries.", "INFO") return reference_map except Exception as e: log_message(f"❌ Failed to fetch reference map from Hugging Face: {str(e)}", "ERROR") return {} def find_matching_filename(audio_filename: str, reference_map: Dict[str, str]) -> Optional[str]: """Finds the matching reference filename for a given audio filename.""" base_name, _ = os.path.splitext(audio_filename) return reference_map.get(base_name) def get_next_file_to_process(source_repo_id: str, state: Dict[str, Any]) -> Optional[Dict[str, Any]]: """ Determines the next file to process based on the current state and the file list from the source Hugging Face repository. """ log_message(f"Determining next file to process from {source_repo_id}...", "INFO") try: # 1. Get the list of all files in the source repo repo_files = hf_api.list_repo_files(repo_id=source_repo_id, repo_type="dataset") # Filter for audio files (e.g., .wav, .mp3) audio_files = sorted([f for f in repo_files if f.lower().endswith(('.wav', '.mp3'))]) processing_status["total_files"] = len(audio_files) if not audio_files: log_message("No audio files found in the source repository.", "INFO") return None # 2. Get the next index from the state next_index = state.get("next_download_index", 0) file_states = state.get("file_states", {}) # 3. Skip forward past all processed and processing files starting from next_index # This ensures we don't repeatedly find files that have already been handled current_index = next_index while current_index < len(audio_files): filename = audio_files[current_index] status = file_states.get(filename, "unprocessed") # If this file is processed or currently processing, skip it if status in ["processed", "processing"]: current_index += 1 continue # If this file failed, we can retry it, so return it if status == "failed": file_url = hf_hub_url(repo_id=source_repo_id, filename=filename, repo_type="dataset") log_message(f"Found failed file for retry at index {current_index}: {filename}", "INFO") return { "filename": filename, "url": file_url, "index": current_index } # If this file is unprocessed, we found our next file file_url = hf_hub_url(repo_id=source_repo_id, filename=filename, repo_type="dataset") log_message(f"Found next file at index {current_index}: {filename}", "INFO") return { "filename": filename, "url": file_url, "index": current_index } log_message("All files have been processed or are locked. Checking for any failed files from the start.", "INFO") # 4. If we've processed all files from next_index to end, check from beginning for failed files for i in range(0, next_index): filename = audio_files[i] status = file_states.get(filename, "unprocessed") if status == "failed": file_url = hf_hub_url(repo_id=source_repo_id, filename=filename, repo_type="dataset") log_message(f"Found failed file for retry at index {i}: {filename}", "INFO") return { "filename": filename, "url": file_url, "index": i } log_message("All files have been processed. Waiting for new files...", "INFO") return None except Exception as e: log_message(f"❌ Failed to get next file to process: {str(e)}", "ERROR") return None def run_whisper_transcription(audio_path: str, output_dir: str, model: str) -> Optional[str]: """ Runs Whisper transcription using the transformers library. Returns the path to the generated JSON file on success. No ffmpeg dependency required. """ log_message(f"πŸŽ™οΈ Starting transcription for {os.path.basename(audio_path)} with model {model}...", "INFO") try: # Get the Whisper pipeline pipe = get_whisper_pipeline() # Load audio using librosa log_message(f"Loading audio file: {audio_path}", "INFO") audio_data, sample_rate = librosa.load(audio_path, sr=16000) # Run transcription log_message(f"Running transcription...", "INFO") result = pipe( audio_data, chunk_length_s=30, batch_size=8, return_timestamps=True ) # Extract text and chunks transcription_text = result.get("text", "") chunks = result.get("chunks", []) log_message(f"βœ… Transcription successful: {len(transcription_text)} characters", "INFO") # Prepare output JSON structure output_json = { "text": transcription_text, "chunks": chunks, "language": result.get("language", "en") } # Save to JSON file base_name, _ = os.path.splitext(os.path.basename(audio_path)) json_output_path = os.path.join(output_dir, f"{base_name}.json") with open(json_output_path, "w", encoding="utf-8") as f: json.dump(output_json, f, indent=2, ensure_ascii=False) log_message(f"βœ… Saved transcription to: {json_output_path}", "INFO") return json_output_path except Exception as e: log_message(f"❌ An error occurred during transcription: {str(e)}", "ERROR") import traceback log_message(f"Traceback: {traceback.format_exc()}", "ERROR") return None def process_audio_file(audio_path: str, reference_map: Dict[str, str], output_filename: str) -> bool: """ Transcribes the audio file, renames the output JSON to match the reference, and uploads the result to the API. """ # 1. Run transcription json_output_path = run_whisper_transcription(audio_path, TRANSCRIPTIONS_FOLDER, WHISPER_MODEL) if not json_output_path: return False # 2. Rename the JSON file to the matched filename # The output_filename already includes the correct extension (e.g., .txt or .json) # We assume the reference map provides the full target filename. # The whisper output is a JSON file named after the audio file. # We need to rename it to the target filename (which should be a JSON file for the backend). # The output_filename is the matched filename from the reference map (e.g., 'audio_file_name.txt') # The backend expects a JSON file. Let's assume the matched filename should be used as the base # but with a .json extension for the upload. # Let's stick to the original logic: the backend expects a JSON file with the name # of the audio file (or the matched reference file) with a .json extension. # Since the whisper output is already a JSON file, we just need to rename it # to the desired final name. # The output_filename passed here is the base name of the audio file or the matched reference file. # If it's a reference file name (e.g., 'file.txt'), we should probably use 'file.json'. # For simplicity and to match the backend's expectation (which handles JSON), # we will rename the whisper output JSON to the base name of the audio file # and ensure it has a .json extension. base_name, _ = os.path.splitext(output_filename) final_json_filename = f"{base_name}.json" final_json_path = os.path.join(TRANSCRIPTIONS_FOLDER, final_json_filename) try: if json_output_path != final_json_path: shutil.move(json_output_path, final_json_path) log_message(f"βœ… Renamed transcription to: {final_json_filename}", "INFO") except Exception as e: log_message(f"❌ Failed to rename transcription file: {str(e)}", "ERROR") return False # 3. Upload transcription to API if upload_transcription_to_api(final_json_path, final_json_filename): processing_status["transcribed_files"] += 1 # Clean up the local transcription file after successful upload try: os.remove(final_json_path) log_message(f"πŸ—‘οΈ Cleaned up local transcription file: {final_json_path}", "INFO") except Exception as e: log_message(f"❌ Failed to clean up transcription file: {str(e)}", "ERROR") return True else: log_message(f"❌ Failed to upload transcription to API: {final_json_filename}", "ERROR") return False def main_processing_loop(): """The main loop that continuously checks for and processes new audio files.""" global processing_status if processing_status["is_running"]: log_message("Processing loop is already running.", "WARNING") return processing_status["is_running"] = True log_message("πŸš€ Audio transcription processing loop started.", "INFO") # 1. Get the reference map once reference_map = get_reference_map(REFERENCE_REPO_ID) if not reference_map: log_message("❌ Could not get reference map. Stopping loop.", "CRITICAL") processing_status["is_running"] = False return try: while processing_status["is_running"]: time.sleep(PROCESSING_DELAY) # 1. Download FRESH state from the API at the start of each iteration # This ensures we respect the next_download_index that other workers may have set current_state = download_state_from_api() next_file_info = get_next_file_to_process(SOURCE_REPO_ID, current_state) if next_file_info is None: log_message("πŸ’€ No new audio files to process. Sleeping for a while...", "INFO") time.sleep(PROCESSING_DELAY * 5) continue target_file = next_file_info['filename'] audio_url = next_file_info['url'] target_index = next_file_info['index'] processing_status["current_file"] = target_file success = False matched_filename = None try: # 2. Lock file by updating state on the API # IMPORTANT: Update next_download_index when locking to prevent other workers from picking same file old_index = current_state["next_download_index"] current_state["next_download_index"] = target_index + 1 log_message(f"πŸ“ Incrementing next_download_index from {old_index} to {current_state['next_download_index']}", "INFO") if not lock_file_for_processing(target_file, current_state): log_message(f"❌ Failed to lock file {target_file}. Skipping.", "ERROR") time.sleep(PROCESSING_DELAY) continue local_wav_path = os.path.join(DOWNLOAD_FOLDER, os.path.basename(target_file)) log_message(f"⬇️ Downloading audio file: {target_file}", "INFO") if download_with_retry(audio_url, local_wav_path): # Extract base filename for matching base_filename = os.path.basename(target_file) matched_filename = find_matching_filename(base_filename, reference_map) # Use matched filename if found, otherwise use original filename output_filename = matched_filename if matched_filename else base_filename # 3. Process and Upload transcription to API if process_audio_file(local_wav_path, reference_map, output_filename): success = True log_message(f"βœ… Finished processing: {target_file}", "INFO") else: log_message(f"❌ Processing failed for: {target_file}", "ERROR") else: log_message(f"❌ Download failed for: {target_file}", "ERROR") except Exception as e: log_message(f"πŸ”₯ An unhandled error occurred while processing {target_file}: {str(e)}", "ERROR") log_failed_file(target_file, str(e)) finally: # 4. Unlock/Mark as processed by updating state on the API # IMPORTANT: Keep the incremented next_download_index from locking if success: # Mark as processed and keep the incremented index, then upload state unlock_file_as_processed(target_file, current_state, current_state["next_download_index"]) processing_status["processed_files"] += 1 else: # Mark as failed but keep the incremented index so next worker can proceed log_message(f"⚠️ File {target_file} failed. Marking as 'failed' and updating state.", "WARNING") current_state["file_states"][target_file] = "failed" # Keep the incremented next_download_index - don't change it upload_state_to_api(current_state) # Clean up the downloaded audio file regardless of success try: if os.path.exists(local_wav_path): os.remove(local_wav_path) log_message(f"πŸ—‘οΈ Cleaned up local audio file: {local_wav_path}", "INFO") except Exception as e: log_message(f"❌ Failed to clean up audio file: {str(e)}", "ERROR") processing_status["current_file"] = None time.sleep(PROCESSING_DELAY) except Exception as e: log_message(f"πŸ”₯ Critical error in main processing loop: {str(e)}", "CRITICAL") finally: processing_status["is_running"] = False log_message("πŸ›‘ Audio transcription processing loop stopped.", "INFO") # --- FastAPI Endpoints (Unchanged) --- # Add to configuration section AUTO_START_PROCESSING = os.environ.get("AUTO_START_PROCESSING", "true").lower() == "true" @app.on_event("startup") async def startup_event(): """Conditionally start processing based on environment variable.""" if AUTO_START_PROCESSING: log_message("πŸš€ AUTO_START_PROCESSING enabled - Starting processing loop...", "INFO") thread = threading.Thread(target=main_processing_loop, daemon=True) thread.start() log_message("βœ… Background processing thread started", "INFO") else: log_message("⏸️ AUTO_START_PROCESSING disabled - Use /start endpoint to begin", "INFO") @app.get("/") async def root(): """Root endpoint to check service status.""" return {"message": "Audio Transcriber Service is running", "status": processing_status} @app.get("/status") async def get_status(): """Get the current processing status.""" return processing_status @app.post("/start") async def start_processing(): """Start the background processing loop.""" if processing_status["is_running"]: return JSONResponse(status_code=200, content={"message": "Processing already running."}) thread = threading.Thread(target=main_processing_loop) thread.start() return JSONResponse(status_code=200, content={"message": "Processing started in background."}) @app.post("/stop") async def stop_processing(): """Stop the background processing loop.""" if not processing_status["is_running"]: return JSONResponse(status_code=200, content={"message": "Processing is not running."}) processing_status["is_running"] = False return JSONResponse(status_code=200, content={"message": "Processing stop requested. Will stop after current file."}) # --- Main Execution --- if __name__ == "__main__": # This block is for local testing and won't be used in the final sandbox execution # but is good practice for a runnable script. uvicorn.run(app, host="0.0.0.0", port=8000)