#!/usr/bin/env python import argparse import json import logging import os import sys import time import uuid from importlib import metadata from audio_separator.remote import AudioSeparatorAPIClient # Cloud Run hard-limits request bodies to 32 MiB. Use 30 MiB threshold so a # little request overhead won't push us over. Larger files go via GCS. GCS_UPLOAD_THRESHOLD_BYTES = 30 * 1024 * 1024 DEFAULT_GCS_INPUT_BUCKET = "nomadkaraoke-audio-separator-outputs" GCS_INPUT_PREFIX = "cli-uploads" def upload_to_gcs(file_path: str, bucket_name: str, logger: logging.Logger) -> str: """Upload a local file to GCS and return its gs:// URI. Requires `google-cloud-storage` and Application Default Credentials (run `gcloud auth application-default login` on the laptop). """ try: from google.cloud import storage except ImportError as e: raise RuntimeError( "google-cloud-storage is required to upload files larger than " f"{GCS_UPLOAD_THRESHOLD_BYTES // (1024 * 1024)} MiB. " "Install it with: pip install google-cloud-storage" ) from e filename = os.path.basename(file_path) blob_path = f"{GCS_INPUT_PREFIX}/{uuid.uuid4()}-{filename}" gcs_uri = f"gs://{bucket_name}/{blob_path}" size_mib = os.path.getsize(file_path) / (1024 * 1024) logger.info(f"Uploading {size_mib:.1f} MiB to {gcs_uri} (server fetches from GCS, bypasses Cloud Run 32 MiB limit)") client = storage.Client() bucket = client.bucket(bucket_name) blob = bucket.blob(blob_path) blob.upload_from_filename(file_path) logger.info(f"Upload complete: {gcs_uri}") return gcs_uri def delete_from_gcs(gcs_uri: str, logger: logging.Logger) -> None: """Best-effort delete of a GCS object. Logs but doesn't raise on failure.""" try: from google.cloud import storage without_prefix = gcs_uri[len("gs://"):] slash_idx = without_prefix.index("/") bucket_name = without_prefix[:slash_idx] blob_path = without_prefix[slash_idx + 1:] client = storage.Client() bucket = client.bucket(bucket_name) bucket.blob(blob_path).delete() logger.info(f"Cleaned up uploaded input: {gcs_uri}") except Exception as e: logger.warning(f"Failed to delete {gcs_uri}: {e} (bucket lifecycle will reclaim it)") def main(): """Main entry point for the remote CLI.""" logger = logging.getLogger(__name__) log_handler = logging.StreamHandler() log_formatter = logging.Formatter(fmt="%(asctime)s.%(msecs)03d - %(levelname)s - %(module)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S") log_handler.setFormatter(log_formatter) logger.addHandler(log_handler) parser = argparse.ArgumentParser(description="Separate audio files using a remote audio-separator API.", formatter_class=lambda prog: argparse.RawTextHelpFormatter(prog, max_help_position=60)) # Get package version package_version = metadata.distribution("audio-separator").version # Main command subparsers subparsers = parser.add_subparsers(dest="command", help="Available commands") # Separate command separate_parser = subparsers.add_parser("separate", help="Separate audio files") separate_parser.add_argument("audio_files", nargs="+", help="Audio file paths to separate") # Model selection (mutually exclusive: preset, single model, or multiple models) model_group = separate_parser.add_mutually_exclusive_group() model_group.add_argument("-p", "--preset", help="Ensemble preset name (e.g. instrumental_clean, karaoke, vocal_balanced)") model_group.add_argument("-m", "--model", help="Single model to use for separation") model_group.add_argument("--models", nargs="+", help="Multiple models to use for separation") separate_parser.add_argument("--timeout", type=int, default=600, help="Timeout in seconds for polling (default: 600)") separate_parser.add_argument("--poll_interval", type=int, default=10, help="Polling interval in seconds (default: 10)") # Output parameters output_group = separate_parser.add_argument_group("Output Parameters") output_group.add_argument("--output_format", default="flac", help="Output format for separated files (default: %(default)s)") output_group.add_argument("--output_bitrate", help="Output bitrate for separated files") output_group.add_argument("--normalization", type=float, default=0.9, help="Max peak amplitude to normalize audio to (default: %(default)s)") output_group.add_argument("--amplification", type=float, default=0.0, help="Min peak amplitude to amplify audio to (default: %(default)s)") output_group.add_argument("--single_stem", help="Output only single stem (e.g. Vocals, Instrumental)") output_group.add_argument("--invert_spect", action="store_true", help="Invert secondary stem using spectrogram") output_group.add_argument("--sample_rate", type=int, default=44100, help="Sample rate of output audio (default: %(default)s)") output_group.add_argument("--use_soundfile", action="store_true", help="Use soundfile for output writing") output_group.add_argument("--use_autocast", action="store_true", help="Use PyTorch autocast for faster inference") output_group.add_argument("--custom_output_names", type=json.loads, help="Custom output names in JSON format") # MDX parameters mdx_group = separate_parser.add_argument_group("MDX Architecture Parameters") mdx_group.add_argument("--mdx_segment_size", type=int, default=256, help="MDX segment size (default: %(default)s)") mdx_group.add_argument("--mdx_overlap", type=float, default=0.25, help="MDX overlap (default: %(default)s)") mdx_group.add_argument("--mdx_batch_size", type=int, default=1, help="MDX batch size (default: %(default)s)") mdx_group.add_argument("--mdx_hop_length", type=int, default=1024, help="MDX hop length (default: %(default)s)") mdx_group.add_argument("--mdx_enable_denoise", action="store_true", help="Enable MDX denoising") # VR parameters vr_group = separate_parser.add_argument_group("VR Architecture Parameters") vr_group.add_argument("--vr_batch_size", type=int, default=1, help="VR batch size (default: %(default)s)") vr_group.add_argument("--vr_window_size", type=int, default=512, help="VR window size (default: %(default)s)") vr_group.add_argument("--vr_aggression", type=int, default=5, help="VR aggression (default: %(default)s)") vr_group.add_argument("--vr_enable_tta", action="store_true", help="Enable VR Test-Time-Augmentation") vr_group.add_argument("--vr_high_end_process", action="store_true", help="Enable VR high end processing") vr_group.add_argument("--vr_enable_post_process", action="store_true", help="Enable VR post processing") vr_group.add_argument("--vr_post_process_threshold", type=float, default=0.2, help="VR post process threshold (default: %(default)s)") # Demucs parameters demucs_group = separate_parser.add_argument_group("Demucs Architecture Parameters") demucs_group.add_argument("--demucs_segment_size", default="Default", help="Demucs segment size (default: %(default)s)") demucs_group.add_argument("--demucs_shifts", type=int, default=2, help="Demucs shifts (default: %(default)s)") demucs_group.add_argument("--demucs_overlap", type=float, default=0.25, help="Demucs overlap (default: %(default)s)") demucs_group.add_argument("--demucs_segments_enabled", type=bool, default=True, help="Enable Demucs segments (default: %(default)s)") # MDXC parameters mdxc_group = separate_parser.add_argument_group("MDXC Architecture Parameters") mdxc_group.add_argument("--mdxc_segment_size", type=int, default=256, help="MDXC segment size (default: %(default)s)") mdxc_group.add_argument("--mdxc_override_model_segment_size", action="store_true", help="Override MDXC model segment size") mdxc_group.add_argument("--mdxc_overlap", type=int, default=None, help="MDXC overlap (default: model config, falling back to 8)") mdxc_group.add_argument("--mdxc_batch_size", type=int, default=None, help="MDXC batch size (default: model config, falling back to 1)") mdxc_group.add_argument("--mdxc_pitch_shift", type=int, default=0, help="MDXC pitch shift (default: %(default)s)") # Status command status_parser = subparsers.add_parser("status", help="Check job status") status_parser.add_argument("task_id", help="Task ID to check status for") # Models command models_parser = subparsers.add_parser("models", help="List available models") models_parser.add_argument("--format", choices=["pretty", "json"], default="pretty", help="Output format") models_parser.add_argument("--filter", help="Filter models by name, type, or stem") # Download command download_parser = subparsers.add_parser("download", help="Download specific files from a job") download_parser.add_argument("task_id", help="Task ID") download_parser.add_argument("filenames", nargs="+", help="Filenames to download") # Global options parser.add_argument("-v", "--version", action="store_true", help="Show version information") parser.add_argument("-d", "--debug", action="store_true", help="Enable debug logging") parser.add_argument("--log_level", default="info", help="Log level (default: info)") parser.add_argument("--api_url", help="API URL (overrides AUDIO_SEPARATOR_API_URL env var)") parser.add_argument( "--gcs-bucket", help=( f"GCS bucket for uploading files >{GCS_UPLOAD_THRESHOLD_BYTES // (1024 * 1024)} MiB " f"(overrides AUDIO_SEPARATOR_GCS_INPUT_BUCKET env var, default: {DEFAULT_GCS_INPUT_BUCKET})" ), ) args = parser.parse_args() # Set up logging if args.debug: log_level = logging.DEBUG else: log_level = getattr(logging, args.log_level.upper()) logger.setLevel(log_level) # Handle version command if args.version: print(f"Client version: {package_version}") # Try to get server version if API URL is available api_url = args.api_url or os.environ.get("AUDIO_SEPARATOR_API_URL") if api_url: api_url = api_url.rstrip("/") api_client = AudioSeparatorAPIClient(api_url, logger) try: server_version = api_client.get_server_version() print(f"Server version: {server_version}") except Exception as e: logger.warning(f"Could not retrieve server version: {e}") else: logger.warning("API URL not provided. Set AUDIO_SEPARATOR_API_URL environment variable or use --api_url to get server version") sys.exit(0) # Get API URL api_url = args.api_url or os.environ.get("AUDIO_SEPARATOR_API_URL") if not api_url: logger.error("API URL not provided. Set AUDIO_SEPARATOR_API_URL environment variable or use --api_url") sys.exit(1) # Remove trailing slash api_url = api_url.rstrip("/") # Create API client api_client = AudioSeparatorAPIClient(api_url, logger) # Resolve GCS bucket for large-file uploads gcs_bucket = args.gcs_bucket or os.environ.get("AUDIO_SEPARATOR_GCS_INPUT_BUCKET", DEFAULT_GCS_INPUT_BUCKET) # Handle commands if args.command == "separate": handle_separate_command(args, api_client, logger, gcs_bucket) elif args.command == "status": handle_status_command(args, api_client, logger) elif args.command == "models": handle_models_command(args, api_client, logger) elif args.command == "download": handle_download_command(args, api_client, logger) else: parser.print_help() sys.exit(1) def handle_separate_command(args, api_client: AudioSeparatorAPIClient, logger: logging.Logger, gcs_bucket: str): """Handle the separate command.""" for audio_file in args.audio_files: logger.info(f"Processing '{audio_file}'...") # Decide upload path: small files go via multipart POST, large files via GCS # to bypass the Cloud Run 32 MiB request body limit. uploaded_gcs_uri = None try: file_size = os.path.getsize(audio_file) use_gcs = file_size > GCS_UPLOAD_THRESHOLD_BYTES except OSError as e: logger.error(f"❌ Cannot read '{audio_file}': {e}") continue try: if use_gcs: logger.info( f"File is {file_size / (1024 * 1024):.1f} MiB (>{GCS_UPLOAD_THRESHOLD_BYTES // (1024 * 1024)} MiB), " "uploading via GCS" ) uploaded_gcs_uri = upload_to_gcs(audio_file, gcs_bucket, logger) source_kwargs = {"file_path": None, "gcs_uri": uploaded_gcs_uri} else: source_kwargs = {"file_path": audio_file, "gcs_uri": None} # Prepare parameters for separation kwargs = { **source_kwargs, "model": args.model, "models": args.models, "preset": args.preset, "timeout": args.timeout, "poll_interval": args.poll_interval, "download": True, # Always download in CLI "output_dir": None, # Use current directory # Output parameters "output_format": args.output_format, "output_bitrate": args.output_bitrate, "normalization_threshold": args.normalization, "amplification_threshold": args.amplification, "output_single_stem": args.single_stem, "invert_using_spec": args.invert_spect, "sample_rate": args.sample_rate, "use_soundfile": args.use_soundfile, "use_autocast": args.use_autocast, "custom_output_names": args.custom_output_names, # MDX parameters "mdx_segment_size": args.mdx_segment_size, "mdx_overlap": args.mdx_overlap, "mdx_batch_size": args.mdx_batch_size, "mdx_hop_length": args.mdx_hop_length, "mdx_enable_denoise": args.mdx_enable_denoise, # VR parameters "vr_batch_size": args.vr_batch_size, "vr_window_size": args.vr_window_size, "vr_aggression": args.vr_aggression, "vr_enable_tta": args.vr_enable_tta, "vr_high_end_process": args.vr_high_end_process, "vr_enable_post_process": args.vr_enable_post_process, "vr_post_process_threshold": args.vr_post_process_threshold, # Demucs parameters "demucs_segment_size": args.demucs_segment_size, "demucs_shifts": args.demucs_shifts, "demucs_overlap": args.demucs_overlap, "demucs_segments_enabled": args.demucs_segments_enabled, # MDXC parameters "mdxc_segment_size": args.mdxc_segment_size, "mdxc_override_model_segment_size": args.mdxc_override_model_segment_size, "mdxc_overlap": args.mdxc_overlap, "mdxc_batch_size": args.mdxc_batch_size, "mdxc_pitch_shift": args.mdxc_pitch_shift, } # Use the convenience method that handles everything result = api_client.separate_audio_and_wait(**kwargs) if result["status"] == "completed": if "downloaded_files" in result: logger.info(f"✅ Separation completed! Downloaded {len(result['downloaded_files'])} files:") for file_path in result["downloaded_files"]: logger.info(f" - {file_path}") else: logger.info(f"✅ Separation completed! Files available for download: {result['files']}") else: logger.error(f"❌ Separation failed: {result.get('error', 'Unknown error')}") except Exception as e: logger.error(f"❌ Error processing '{audio_file}': {e}") finally: if uploaded_gcs_uri: delete_from_gcs(uploaded_gcs_uri, logger) def handle_status_command(args, api_client: AudioSeparatorAPIClient, logger: logging.Logger): """Handle the status command.""" try: status = api_client.get_job_status(args.task_id) logger.info(f"Job Status: {status['status']}") if "progress" in status: progress_info = f"Progress: {status['progress']}%" if "current_model_index" in status and "total_models" in status: model_info = f" (Model {status['current_model_index'] + 1}/{status['total_models']})" progress_info += model_info logger.info(progress_info) if "original_filename" in status: logger.info(f"Original File: {status['original_filename']}") if "models_used" in status: logger.info(f"Models Used: {', '.join(status['models_used'])}") if status["status"] == "error" and "error" in status: logger.error(f"Error: {status['error']}") elif status["status"] == "completed" and "files" in status: logger.info("Output Files:") for filename in status["files"]: logger.info(f" - {filename}") except Exception as e: logger.error(f"❌ Error getting status: {e}") def handle_models_command(args, api_client: AudioSeparatorAPIClient, logger: logging.Logger): """Handle the models command.""" try: models = api_client.list_models(args.format, args.filter) if args.format == "json": print(json.dumps(models, indent=2)) else: print(models["text"]) except Exception as e: logger.error(f"❌ Error listing models: {e}") def handle_download_command(args, api_client: AudioSeparatorAPIClient, logger: logging.Logger): """Handle the download command.""" try: for filename in args.filenames: logger.info(f"📂 Downloading: {filename}") output_path = api_client.download_file(args.task_id, filename) logger.info(f"✅ Downloaded: {output_path}") except Exception as e: logger.error(f"❌ Error downloading files: {e}") def poll_for_completion(task_id: str, api_client: AudioSeparatorAPIClient, logger: logging.Logger, timeout: int = 600, poll_interval: int = 10) -> bool: """Poll for job completion.""" start_time = time.time() last_progress = -1 while time.time() - start_time < timeout: try: status = api_client.get_job_status(task_id) current_status = status.get("status", "unknown") # Show progress if it changed if "progress" in status and status["progress"] != last_progress: progress_info = f"📊 Progress: {status['progress']}%" if "current_model_index" in status and "total_models" in status: model_info = f" (Model {status['current_model_index'] + 1}/{status['total_models']})" progress_info += model_info logger.info(progress_info) last_progress = status["progress"] # Check if completed if current_status == "completed": logger.info("✅ Job completed!") return True elif current_status == "error": logger.error(f"❌ Job failed: {status.get('error', 'Unknown error')}") return False # Wait before next poll time.sleep(poll_interval) except Exception as e: logger.warning(f"Error polling status: {e}") time.sleep(poll_interval) logger.error(f"❌ Job polling timed out after {timeout} seconds") return False def download_files(task_id: str, files: list, api_client: AudioSeparatorAPIClient, logger: logging.Logger): """Download all files from a completed job.""" if not files: logger.warning("No files to download") return logger.info(f"📥 Downloading {len(files)} output files...") downloaded_count = 0 for filename in files: try: logger.info(f" 📂 Downloading: {filename}") output_path = api_client.download_file(task_id, filename) logger.info(f" ✅ Downloaded: {output_path}") downloaded_count += 1 except Exception as e: logger.error(f" ❌ Failed to download {filename}: {e}") if downloaded_count > 0: logger.info(f"🎉 Successfully downloaded {downloaded_count} files to current directory!") else: logger.error("❌ No files were successfully downloaded") if __name__ == "__main__": main()