#!/usr/bin/env python3 """ Voice Designer Stem Captioning Script Processes vocal stem records with Gemini captioning, supporting parallel execution and appending to existing stems_captions field. """ import os import json import time import tempfile import logging import multiprocessing as mp import random import re from datetime import datetime from pathlib import Path import argparse from tqdm import tqdm from concurrent.futures import ProcessPoolExecutor, as_completed import uuid # Gemini imports from google import genai from google.genai import types # Configuration GOOGLE_CREDS_PATH = "/home/vibert/data/credentials/google-3.json" GEMINI_PROJECT_ID = "ml-gemini-455703" GEMINI_LOCATION = "us-central1" DEFAULT_GEMINI_MODEL = "gemini-2.5-flash" MAX_GEMINI_OUTPUT_TOKENS = 500 # Comprehensive list of vocal-related keywords for tag extraction VOCAL_KEYWORDS = [ "vocal", "vox", "sing", "voice", "harmon", "choir", "chorus", "croon", "rap", "spoken word", "speech", "talk", "whisper", "hum", "melody", "lyric", "verse", "hook", "ad-lib", "growl", "scream", "shout", "yell", "chant", "duet", "trio", "quartet", "ensemble", "acapella", "tenor", "soprano", "alto", "baritone", "contralto", "mezzo", "male", "female", "child", "boy", "girl", "operatic", "belt", "rasp", "breath", "airy", "nasal", "accent", "dialect", "soulful", "powerful", "gentle", "aggressive", "smooth", "rough", "clean", "distort", "auto-tun", "pitch", "autotune", ] def validate_comma_separated_keywords(text: str, min_count: int = 15) -> tuple[bool, int]: """ Validate that text contains comma-separated keywords. Args: text: The text to validate min_count: Minimum number of keywords required Returns: Tuple of (is_valid: bool, keyword_count: int) Raises: ValueError: If format is invalid for comma-separated keywords """ if not text or not text.strip(): raise ValueError("Empty or whitespace-only response") # Clean the text cleaned_text = text.strip() # Check for basic comma separation (should contain at least one comma) if "," not in cleaned_text: raise ValueError("Response does not contain comma-separated format") # Split by commas and clean each keyword keywords = [keyword.strip() for keyword in cleaned_text.split(",")] # Filter out empty keywords keywords = [kw for kw in keywords if kw] # Check for reasonable keyword format (reasonable length only) # for keyword in keywords: # # Skip if keyword is too short or too long # if len(keyword) < 2 or len(keyword) > 50: # raise ValueError(f"Invalid keyword length: '{keyword}'") keyword_count = len(keywords) is_valid = keyword_count >= min_count return is_valid, keyword_count def select_model_probabilistically() -> str: """ Select Gemini model probabilistically: 80% Flash, 20% Pro Returns: Selected model name """ models = ["gemini-2.5-flash", "gemini-2.5-pro"] weights = [0.8, 0.2] selected = random.choices(models, weights=weights, k=1)[0] return selected class VoiceDesignerCaptioner: def __init__( self, prompts_file, tmpdir=None, max_retries=5, gemini_model=None, min_keyword_count=20, use_expansion=False, ): """Initialize the voice designer captioner""" self.prompts_file = prompts_file # self.tmpdir = tmpdir or f"/home/vibert/tmp/voice_captioning_{uuid.uuid4().hex[:8]}" self.max_retries = max_retries self.gemini_model = gemini_model or DEFAULT_GEMINI_MODEL self.min_keyword_count = min_keyword_count self.use_expansion = use_expansion # Create tmpdir # os.makedirs(self.tmpdir, exist_ok=True) # Initialize Gemini client self.gemini_client = None self.setup_gemini_client() # Load all prompts (no filtering) self.prompts = self.load_all_prompts() def setup_gemini_client(self): """Setup Gemini client""" try: os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = GOOGLE_CREDS_PATH self.gemini_client = genai.Client( vertexai=True, project=GEMINI_PROJECT_ID, location=GEMINI_LOCATION ) return True except Exception as e: logging.error(f"Gemini setup error: {e}") return False def load_all_prompts(self): """Load all prompts from the prompts file without filtering""" try: with open(self.prompts_file, "r") as f: prompts_data = json.load(f) # Get all prompts without filtering prompts = [] if "prompts" in prompts_data: for prompt_type, prompt_info in prompts_data["prompts"].items(): prompts.append({"type": prompt_type, "prompt": prompt_info["prompt"]}) logging.info(f"Loaded {len(prompts)} prompts: {[p['type'] for p in prompts]}") return prompts except Exception as e: logging.error(f"Could not load prompts from {self.prompts_file}: {e}") return [] def _load_audio_as_prompt(self, local_fp): """Convert local audio file to Gemini-compatible format""" with tempfile.TemporaryDirectory() as tempdir: if local_fp.endswith(".opus"): out_local_fp = local_fp else: out_local_fp = os.path.join(tempdir, "audio.opus") import subprocess subprocess.run( [ "ffmpeg", "-i", local_fp, "-c:a", "libopus", "-ar", "48000", "-ac", "2", "-y", out_local_fp, ], check=True, capture_output=True, ) with open(out_local_fp, "rb") as f: audio_bytes = f.read() return types.Part.from_bytes(data=audio_bytes, mime_type="audio/mpeg") def generate_caption(self, audio_part, prompt_text, prompt_type, record=None): """Generate caption using Gemini with retry logic and keyword accumulation""" if not self.gemini_client: return {"error": "Gemini client not initialized"} base_delay = 2.0 accumulated_keywords = [] total_attempts = 0 max_total_attempts = self.max_retries + 3 # Extra attempts for accumulation while total_attempts < max_total_attempts: try: # Select model probabilistically for each caption attempt current_model = select_model_probabilistically() # Configure thinking settings based on model thinking_config = None if "flash" in current_model.lower(): thinking_config = types.ThinkingConfig( includeThoughts=False, thinkingBudget=0, ) elif "pro" in current_model.lower(): # Pro model requires non-zero thinking budget thinking_config = types.ThinkingConfig( includeThoughts=False, thinkingBudget=4000, # Use minimum non-zero budget ) config_params = { "candidateCount": 1, "audioTimestamp": False, "temperature": 0.0, "max_output_tokens": MAX_GEMINI_OUTPUT_TOKENS, "frequency_penalty": 0.1, } if thinking_config: config_params["thinkingConfig"] = thinking_config # Enhance prompt with vocal tags context if available if record: vocal_tags = self.extract_vocal_tags(record) enhanced_prompt = self.enhance_prompt_with_vocal_context(prompt_text, vocal_tags) if vocal_tags: logging.info( f" 🎤 Enhanced prompt with vocal context: {', '.join(vocal_tags[:3])}{'...' if len(vocal_tags) > 3 else ''}" ) else: enhanced_prompt = prompt_text # Modify prompt for accumulation if we already have keywords and expansion is enabled if accumulated_keywords and self.use_expansion: expansion_prompt = f"{enhanced_prompt} Add more descriptive keywords to expand this list: {', '.join(accumulated_keywords)}" current_prompt = expansion_prompt else: current_prompt = enhanced_prompt response = self.gemini_client.models.generate_content( model=current_model, contents=[audio_part, current_prompt], config=types.GenerateContentConfig(**config_params), ) if ( response.candidates and response.candidates[0].content and response.candidates[0].content.parts ): result = response.candidates[0].content.parts[0].text # Validate comma-separated format and count keywords try: is_sufficient, keyword_count = validate_comma_separated_keywords( result, self.min_keyword_count ) # Parse keywords from this response (remove trailing periods and markdown formatting) current_keywords = [ kw.strip().rstrip(".").lstrip("*").rstrip("*") for kw in result.split(",") if kw.strip() ] if accumulated_keywords: # Merge with existing keywords, avoiding duplicates (case-insensitive) all_keywords = accumulated_keywords.copy() existing_lower = {kw.lower() for kw in all_keywords} for keyword in current_keywords: if keyword.lower() not in existing_lower: all_keywords.append(keyword) existing_lower.add(keyword.lower()) accumulated_keywords = all_keywords else: accumulated_keywords = current_keywords # Deduplicate final keyword list (case-insensitive, preserve first occurrence) seen = set() deduped_keywords = [] for keyword in accumulated_keywords: keyword_lower = keyword.lower() if keyword_lower not in seen: deduped_keywords.append(keyword) seen.add(keyword_lower) accumulated_keywords = deduped_keywords # Check if we have enough keywords now final_count = len(accumulated_keywords) if final_count >= self.min_keyword_count: final_result = ", ".join(accumulated_keywords) logging.info( f" ✅ Accumulated {final_count} keywords successfully (deduplicated)" ) return {"caption": final_result, "error": None} else: # Need more keywords - continue accumulating logging.info( f" 📝 Got {final_count} keywords, need {self.min_keyword_count}. Accumulating more..." ) total_attempts += 1 continue except ValueError as ve: # Format validation failed - treat as retryable error logging.warning(f" ⚠️ Validation failed: {ve}") raise Exception(f"Validation failed: {ve}") else: # Treat "No content in response" as a retryable error raise Exception("No content in response") except Exception as e: total_attempts += 1 error_msg = str(e).lower() is_last_attempt = total_attempts >= max_total_attempts if is_last_attempt: # Return partial result if we have some keywords if accumulated_keywords: partial_result = ", ".join(accumulated_keywords) logging.warning( f"❌ Partial result after {total_attempts} attempts: {len(accumulated_keywords)} keywords" ) return {"caption": partial_result, "error": f"Partial result: {str(e)}"} else: logging.error(f"❌ Failed after {total_attempts} attempts: {str(e)}") return {"error": f"Failed after {total_attempts} attempts: {str(e)}"} # Calculate delay with exponential backoff attempt_for_delay = min(total_attempts, self.max_retries) if ( "resource exhausted" in error_msg or "quota" in error_msg or "rate limit" in error_msg ): delay = base_delay * (3**attempt_for_delay) * 2 logging.warning( f"⚠️ Resource exhausted (attempt {total_attempts}), waiting {delay}s..." ) elif "internal error" in error_msg or "service unavailable" in error_msg: delay = base_delay * (2**attempt_for_delay) logging.warning(f"⚠️ Server error (attempt {total_attempts}), waiting {delay}s...") else: delay = base_delay + (1.5**attempt_for_delay) logging.warning( f"⚠️ API error {error_msg} (model: {current_model})(attempt {total_attempts}), waiting {delay}s..." ) time.sleep(delay) # Should not reach here, but return error if we do return {"error": f"Exhausted all {max_total_attempts} attempts"} def is_vocal_stem(self, stem_name): """Check if stem name contains vocal-related keywords""" vocal_keywords = ["vocal", "vox"] stem_name_lower = stem_name.lower() return any(keyword in stem_name_lower for keyword in vocal_keywords) def is_record_fully_processed(self, record, required_prompts, vocal_stems): """Check if record has successful captions for all vocal stems + prompt types""" stems_captions = record.get("stems_captions", {}) for stem_name in vocal_stems: if stem_name not in stems_captions: return False existing_captions = stems_captions[stem_name] completed_prompts = set() for caption_obj in existing_captions: prompt_type = caption_obj.get("prompt_type") caption_text = caption_obj.get("caption", "") error = caption_obj.get("error") # Consider successful if: has caption, no error, meets 10-char requirement if prompt_type and caption_text and not error and len(caption_text.strip()) >= 10: completed_prompts.add(prompt_type) # Check if all required prompts are completed for this stem required_prompt_types = {p["type"] for p in required_prompts} if not required_prompt_types.issubset(completed_prompts): return False return True def extract_vocal_tags(self, record): """Extract tags that contain vocal-related keywords""" tags = record.get("tags", []) if not tags: return [] vocal_tags = [] for tag in tags: tag_lower = str(tag).lower() if any(keyword in tag_lower for keyword in VOCAL_KEYWORDS): vocal_tags.append(str(tag)) return vocal_tags def enhance_prompt_with_vocal_context(self, base_prompt, vocal_tags): """Enhance the prompt with vocal tag context""" if not vocal_tags: return base_prompt # Create a clean, readable list of vocal characteristics vocal_context = ", ".join(vocal_tags) # Append vocal context to the prompt enhanced_prompt = f"""{base_prompt} Additional context - Here are some characteristics I already know about this vocal: {vocal_context} Please incorporate this information to provide more accurate and detailed keywords.""" return enhanced_prompt def process_record_stems(self, record, vocal_stems_only=True): """Process stems in a single record""" record_id = record.get("id", "unknown") # Initialize or get existing stems_captions if "stems_captions" not in record: record["stems_captions"] = {} stems = record.get("stems", {}) if not stems: logging.warning(f"Record {record_id} has no stems") return record # Filter stems based on vocal_stems_only setting if vocal_stems_only: target_stems = {name: path for name, path in stems.items() if self.is_vocal_stem(name)} stem_type_desc = "vocal stems" if not target_stems: logging.warning(f"Record {record_id} has no vocal stems") return record else: target_stems = stems stem_type_desc = "stems" logging.info( f"Processing record {record_id} with {len(target_stems)} {stem_type_desc}" + ( f" (out of {len(stems)} total stems)" if vocal_stems_only and len(target_stems) != len(stems) else "" ) ) for stem_name, stem_path in target_stems.items(): stem_label = "vocal stem" if vocal_stems_only else "stem" logging.info(f" → Processing {stem_label}: {stem_name}") # Initialize or get existing captions for this stem if stem_name not in record["stems_captions"]: record["stems_captions"][stem_name] = [] # Check if file exists if not os.path.exists(stem_path): logging.error(f" ❌ File not found: {stem_path}") continue # Process audio once per stem try: audio_part = self._load_audio_as_prompt(stem_path) logging.info(f" 🎵 Audio processed for {stem_name}") except Exception as e: logging.error(f" ❌ Audio processing failed: {e}") continue # Get existing prompt types to avoid duplicates existing_prompt_types = set() for existing_caption in record["stems_captions"][stem_name]: if "prompt_type" in existing_caption: existing_prompt_types.add(existing_caption["prompt_type"]) # Process all prompts for this stem for prompt_info in self.prompts: prompt_type = prompt_info["type"] prompt_text = prompt_info["prompt"] # If this prompt type already exists, append to the list if prompt_type in existing_prompt_types: logging.info(f" → Appending to existing prompt: {prompt_type}") else: logging.info(f" → New prompt: {prompt_type}") result = self.generate_caption(audio_part, prompt_text, prompt_type, record) # Create caption object caption_obj = { "prompt_type": prompt_type, "caption": result.get("caption", ""), "error": result.get("error"), } # Append to the existing list (whether new or existing prompt type) record["stems_captions"][stem_name].append(caption_obj) if result.get("error"): logging.error(f" ❌ Error: {result['error']}") else: caption_preview = ( result["caption"][:80] + "..." if len(result["caption"]) > 80 else result["caption"] ) logging.info(f" ✅ Success: {caption_preview}") logging.info(f" ✅ Record {record_id} completed") return record def process_record_worker(args): """Worker function for parallel processing""" ( record, prompts_file, tmpdir, vocal_stems_only, gemini_model, force_model, min_keyword_count, use_expansion, ) = args # Setup logging for this worker worker_id = os.getpid() logging.basicConfig(level=logging.INFO, format=f"[Worker {worker_id}] %(asctime)s - %(message)s") # Select model probabilistically unless forced if force_model: selected_model = gemini_model logging.info(f"Using forced model: {selected_model}") else: selected_model = select_model_probabilistically() logging.info(f"Probabilistically selected model: {selected_model}") # Create captioner for this worker captioner = VoiceDesignerCaptioner( prompts_file, tmpdir=tmpdir, gemini_model=selected_model, min_keyword_count=min_keyword_count, use_expansion=use_expansion, ) try: return captioner.process_record_stems(record, vocal_stems_only) except Exception as e: logging.error(f"Worker failed processing record {record.get('id', 'unknown')}: {e}") return record # Return original record on failure def load_completed_records(prev_output_file, prompts_file, vocal_stems_only=True): """Load set of fully completed record IDs from previous output file""" if not prev_output_file or not os.path.exists(prev_output_file): return set() # Load required prompts to check completion try: with open(prompts_file, "r") as f: prompts_data = json.load(f) required_prompts = [] if "prompts" in prompts_data: for prompt_type, prompt_info in prompts_data["prompts"].items(): required_prompts.append({"type": prompt_type, "prompt": prompt_info["prompt"]}) except Exception as e: logging.error(f"Could not load prompts from {prompts_file}: {e}") return set() completed_ids = set() total_records = 0 # Create temporary captioner to use helper methods temp_captioner = VoiceDesignerCaptioner(prompts_file) try: with open(prev_output_file, "r") as f: for line in f: if not line.strip(): continue try: record = json.loads(line.strip()) total_records += 1 # Get vocal stems for this record stems = record.get("stems", {}) if vocal_stems_only: vocal_stems = [ name for name, path in stems.items() if temp_captioner.is_vocal_stem(name) ] else: vocal_stems = list(stems.keys()) # Skip records with no relevant stems if not vocal_stems: continue # Check if this record is fully processed if temp_captioner.is_record_fully_processed(record, required_prompts, vocal_stems): completed_ids.add(record.get("id")) except json.JSONDecodeError: continue except Exception as e: logging.error(f"Error reading previous output file {prev_output_file}: {e}") return set() logging.info( f"Found {len(completed_ids)} completed records out of {total_records} total records in previous output" ) return completed_ids def run_parallel_captioning( input_jsonl, output_jsonl, prompts_file, num_workers=4, tmpdir=None, limit=None, vocal_stems_only=True, gemini_model=None, save_every=10, force_model=False, min_keyword_count=20, use_expansion=False, resume=False, prev_output_jsonl=None, ): """Run parallel captioning on stems""" gemini_model = gemini_model or DEFAULT_GEMINI_MODEL # Setup logging timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") log_file = f"/home/vibert/tmp/voice_captioning_{timestamp}.log" logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", handlers=[logging.FileHandler(log_file), logging.StreamHandler()], ) logging.info(f"Voice Designer Parallel Captioning Started") logging.info(f"Input: {input_jsonl}") logging.info(f"Output: {output_jsonl}") logging.info(f"Prompts: {prompts_file}") logging.info(f"Workers: {num_workers}") logging.info(f"Limit: {limit or 'No limit'}") logging.info(f"Process: {'Vocal stems only' if vocal_stems_only else 'All stems'}") logging.info(f"Model: {gemini_model}") # Load completed records if resuming completed_ids = set() if resume and prev_output_jsonl: logging.info(f"Resume mode enabled, loading completed records from: {prev_output_jsonl}") completed_ids = load_completed_records(prev_output_jsonl, prompts_file, vocal_stems_only) if completed_ids: logging.info( f"Loaded {len(completed_ids)} completed record IDs, will skip these during processing" ) else: logging.info("No completed records found or previous output file is empty") # Read all records and preserve order with index mapping all_records = [] record_index_map = {} total_input_records = 0 skipped_completed = 0 with open(input_jsonl, "r") as f: for line_num, line in enumerate(f): if limit and line_num >= limit: break try: record = json.loads(line.strip()) total_input_records += 1 # Only process records with stems if record.get("stems"): record_id = record.get("id", f"line_{line_num}") # Skip if this record is already completed (when resuming) if resume and record_id in completed_ids: skipped_completed += 1 continue all_records.append(record) record_index_map[record_id] = line_num # Store original position except json.JSONDecodeError: logging.warning(f"Skipping malformed line {line_num}") continue # Log resume statistics if resume: logging.info(f"Resume statistics:") logging.info(f" - Total input records examined: {total_input_records}") logging.info(f" - Records skipped (already completed): {skipped_completed}") logging.info(f" - Records remaining to process: {len(all_records)}") else: logging.info(f"Loaded {len(all_records)} records with stems") if not all_records: logging.error("No records with stems found") return # Prepare output file (clear existing content) output_path = Path(output_jsonl) output_path.parent.mkdir(parents=True, exist_ok=True) # Process in batches to maintain order and enable incremental saving total_processed = 0 processed_records = {} # Dictionary to store results by record_id for batch_start in range(0, len(all_records), save_every): batch_end = min(batch_start + save_every, len(all_records)) batch_records = all_records[batch_start:batch_end] logging.info( f"Processing batch {batch_start//save_every + 1}: records {batch_start+1} to {batch_end}" ) # Create worker arguments for this batch worker_args = [ ( record, prompts_file, tmpdir, vocal_stems_only, gemini_model, force_model, min_keyword_count, use_expansion, ) for record in batch_records ] # Process batch in parallel batch_results = {} with ProcessPoolExecutor(max_workers=num_workers) as executor: # Submit all jobs for this batch future_to_record = { executor.submit(process_record_worker, args): args[0] for args in worker_args } # Collect batch results for future in tqdm( as_completed(future_to_record), total=len(future_to_record), desc=f"Batch {batch_start//save_every + 1}", ): try: processed_record = future.result() record_id = processed_record.get("id", "unknown") batch_results[record_id] = processed_record except Exception as e: original_record = future_to_record[future] record_id = original_record.get("id", "unknown") logging.error(f"Failed to process record {record_id}: {e}") batch_results[record_id] = original_record # Add batch results to main results dictionary processed_records.update(batch_results) total_processed += len(batch_results) # Write accumulated results in original order (incremental save) write_ordered_results(processed_records, all_records, output_jsonl, record_index_map) logging.info( f"Saved batch {batch_start//save_every + 1}. Total processed: {total_processed}/{len(all_records)}" ) logging.info(f"✅ Processing complete! Results saved to: {output_jsonl}") return total_processed def write_ordered_results(processed_records, all_records, output_jsonl, record_index_map): """Write processed records to file in original order""" with open(output_jsonl, "w") as f: for record in all_records: record_id = record.get("id", "unknown") if record_id in processed_records: # Write processed record f.write(json.dumps(processed_records[record_id]) + "\n") else: # Write original record if not yet processed f.write(json.dumps(record) + "\n") def main(): parser = argparse.ArgumentParser(description="Voice Designer Vocal Stem Captioning with Gemini") parser.add_argument( "--input-jsonl", type=str, required=True, help="Path to input JSONL file with vocal stems" ) parser.add_argument("--output-jsonl", type=str, required=True, help="Path to output JSONL file") parser.add_argument( "--prompts-file", type=str, required=True, help="Path to JSON file containing prompts" ) parser.add_argument( "--num-workers", type=int, default=4, help="Number of parallel workers (default: 4)" ) parser.add_argument("--tmpdir", type=str, default=None, help="Temporary directory for processing") parser.add_argument( "--limit", type=int, default=None, help="Limit number of records to process (for testing)" ) parser.add_argument( "--all-stems", action="store_true", help="Process all stems instead of vocal stems only (default: vocal stems only)", ) parser.add_argument( "--model", type=str, choices=["gemini-2.5-flash", "gemini-2.5-pro"], default="gemini-2.5-flash", help="Gemini model to use (default: gemini-2.5-flash)", ) parser.add_argument( "--save-every", type=int, default=10, help="Save results every N processed records (default: 10)" ) parser.add_argument( "--force-model", action="store_true", help="Force use of specified model instead of probabilistic selection (20%% pro, 80%% flash)", ) parser.add_argument( "--min-keywords", type=int, default=20, help="Minimum number of keywords required per caption (default: 20)", ) parser.add_argument( "--use-expansion", action="store_true", help="Enable keyword expansion prompts when accumulating keywords (default: False)", ) parser.add_argument( "--resume", action="store_true", help="Resume from previous output file, processing only remaining records (default: False)", ) parser.add_argument( "--prev-output-jsonl", type=str, help="Path to previous output JSONL file to resume from (required when --resume is used)", ) args = parser.parse_args() # Validate inputs if not Path(args.input_jsonl).exists(): print(f"Error: Input file {args.input_jsonl} does not exist") return 1 if not Path(args.prompts_file).exists(): print(f"Error: Prompts file {args.prompts_file} does not exist") return 1 # Validate resume parameters if args.resume: if not args.prev_output_jsonl: print(f"Error: --prev-output-jsonl is required when --resume is enabled") return 1 if not Path(args.prev_output_jsonl).exists(): print(f"Error: Previous output file {args.prev_output_jsonl} does not exist") return 1 if Path(args.output_jsonl).exists(): print(f"Warning: Output file {args.output_jsonl} already exists and will be overwritten") if args.prev_output_jsonl == args.output_jsonl: print(f"Error: Previous output file and new output file cannot be the same") return 1 vocal_stems_only = not args.all_stems print(f"🎵 Voice Designer Stem Captioning") print(f"Input: {args.input_jsonl}") print(f"Output: {args.output_jsonl}") if args.resume: print(f"Resume from: {args.prev_output_jsonl}") print(f"Prompts: {args.prompts_file}") print(f"Workers: {args.num_workers}") print(f"Limit: {args.limit or 'No limit'}") print(f"Process: {'Vocal stems only' if vocal_stems_only else 'All stems'}") print(f"Model: {args.model}") print(f"Selection: {'Forced' if args.force_model else 'Probabilistic (80% Flash, 20% Pro)'}") print(f"Min keywords: {args.min_keywords}") print(f"Use expansion: {args.use_expansion}") print(f"Save every: {args.save_every} records") if args.resume: print(f"Resume mode: ENABLED") print() try: processed_count = run_parallel_captioning( input_jsonl=args.input_jsonl, output_jsonl=args.output_jsonl, prompts_file=args.prompts_file, num_workers=args.num_workers, tmpdir=args.tmpdir, limit=args.limit, vocal_stems_only=vocal_stems_only, gemini_model=args.model, save_every=args.save_every, force_model=args.force_model, min_keyword_count=args.min_keywords, use_expansion=args.use_expansion, resume=args.resume, prev_output_jsonl=args.prev_output_jsonl, ) print(f"✅ Successfully processed {processed_count} records") return 0 except Exception as e: print(f"❌ Error during processing: {e}") return 1 if __name__ == "__main__": exit(main())