from __future__ import annotations import json from pathlib import Path import gradio as gr from .catalog import catalog_choices, default_model_id from .config import ( ENSEMBLE_ALGORITHMS, JOB_TTL_SECONDS, OUTPUT_FORMATS, SAFE_PACKAGE_MODELS, STEM_CHOICES, ) from .dataset_samples import ( DEFAULT_DATASET_SAMPLE_ID, dataset_sample_choices, dataset_sample_markdown, ) from .duration import ( DEFAULT_BASE_SECONDS_PER_MINUTE, DEFAULT_MANUAL_SECONDS, DEFAULT_SAFETY_MULTIPLIER, DEFAULT_SAFE_MODE_ENABLED, DURATION_MODE_MANUAL, DURATION_MODE_SEMI_AUTO, DURATION_MODES, ) from .errors import StageError from .longform import ( CHUNK_MODE_AUTO, CHUNK_MODE_FIXED, CHUNK_MODES, DEFAULT_FIXED_CHUNK_SECONDS, DEFAULT_PREVIEW_SECONDS, RANGE_MODE_CUSTOM, RANGE_MODE_FULL, RANGE_MODE_PREVIEW, RANGE_MODES, ) from .observability import DEFAULT_LOG_LEVEL, LOG_LEVELS from .preparation import ( BITRATE_AUTO, BUNDLE_LAYOUT_BY_INPUT, BUNDLE_LAYOUT_FLAT, cleanup_prepared_state, prepare_job, ) from .service import ( EXECUTION_MODE_INITIAL, EXECUTION_MODE_RETRY_INCOMPLETE, request_prepared_job_cancel, run_prepared_job, ) OUTPUT_BITRATES = [BITRATE_AUTO, "96k", "128k", "192k", "256k", "320k"] OUTPUT_SAMPLE_RATES = [24000, 32000, 44100, 48000] def _result_has_retryable_items(result) -> bool: try: returned_files = result[3] or [] except Exception: return False for value in returned_files: path = Path(str(value)) if path.name != "batch_manifest.json" or not path.is_file(): continue try: payload = json.loads(path.read_text(encoding="utf-8")) except Exception: return False return bool(payload.get("retryable_input_indexes")) return False def run_prepared_separation( prepared_state_value, execution_mode=EXECUTION_MODE_INITIAL, progress=None, callback_timing=None, ): result = run_prepared_job( prepared_state_value, progress=progress, callback_timing=callback_timing, execution_mode=execution_mode, ) retryable = _result_has_retryable_items(result) return ( *result, gr.update(interactive=False), gr.update(interactive=retryable), gr.update(interactive=retryable), ) def request_cancel_with_ui(prepared_state_value): return request_prepared_job_cancel(prepared_state_value) def prepare_separation_with_progress(*args, progress=gr.Progress(track_tqdm=True)): try: if len(args) not in {34, 35}: raise ValueError(f"Expected 34 or 35 Preparation inputs, received {len(args)}.") batch_continue_on_item_error = args[34] if len(args) == 35 else True state, markdown, config_path, log_path = prepare_job( *args[:22], 0, *args[28:34], progress=progress, range_mode=args[22], range_start_seconds=args[23], range_end_seconds=args[24], preview_seconds=args[25], chunk_mode=args[26], fixed_chunk_seconds=args[27], batch_continue_on_item_error=batch_continue_on_item_error, ) return ( state, markdown, config_path, log_path, gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True), ) except StageError as exc: hint = f"\n\n{exc.hint}" if exc.hint else "" diagnostic = getattr(exc, "diagnostic_path", None) log_path = getattr(exc, "log_path", None) markdown = ( "### Preparation failed before GPU allocation\n" f"- Stage: `{exc.stage}`\n" f"- Error: `{exc.message}`{hint}" ) return ( "", markdown, diagnostic, log_path, gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), ) except Exception as exc: return ( "", f"### Preparation failed before GPU allocation\n`{type(exc).__name__}: {exc}`", None, None, gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), ) def _source_visibility(value: str): return ( gr.update(visible=value == "Hugging Face"), gr.update(visible=value == "GitHub"), ) def _range_visibility(value: str): return ( gr.update(visible=value != RANGE_MODE_FULL), gr.update(visible=value == RANGE_MODE_CUSTOM), gr.update(visible=value == RANGE_MODE_PREVIEW), ) def _chunk_visibility(value: str): return gr.update(visible=value == CHUNK_MODE_FIXED) def _duration_visibility(value: str): return ( gr.update(visible=value == DURATION_MODE_MANUAL), gr.update(visible=value != DURATION_MODE_MANUAL), ) def _duration_controls(label_prefix: str = ""): prefix = f"{label_prefix} " if label_prefix else "" duration_mode = gr.Radio( label=f"{prefix}Time request mode", choices=DURATION_MODES, value=DURATION_MODE_SEMI_AUTO, ) with gr.Group(visible=False) as manual_group: manual_gpu_seconds = gr.Slider( label=f"{prefix}Manual maximum GPU time (seconds)", minimum=30, maximum=300, step=5, value=DEFAULT_MANUAL_SECONDS, ) with gr.Group(visible=True) as semi_auto_group: semi_auto_safe_mode = gr.Checkbox( label=f"{prefix}Extra safety margin (+30%, Semi-auto only)", value=DEFAULT_SAFE_MODE_ENABLED, ) gr.Markdown( "The extra margin defaults OFF. A request that is too short may still finish, but it can be aborted after an overrun warning. " "The +30% margin affects only Semi-auto; Manual is unchanged." ) with gr.Accordion("Advanced Semi-auto calibration", open=False): semi_auto_base_seconds = gr.Number( label=f"{prefix}Base estimate: GPU seconds per source minute and model", value=DEFAULT_BASE_SECONDS_PER_MINUTE, minimum=1, maximum=120, precision=1, ) semi_auto_safety_multiplier = gr.Slider( label=f"{prefix}Calibration multiplier", minimum=0.5, maximum=3.0, step=0.05, value=DEFAULT_SAFETY_MULTIPLIER, ) duration_mode.change( _duration_visibility, inputs=duration_mode, outputs=[manual_group, semi_auto_group], show_progress="hidden", ) return ( duration_mode, manual_gpu_seconds, semi_auto_base_seconds, semi_auto_safety_multiplier, semi_auto_safe_mode, ) def _invalidate_preparation(): return ( "", "Settings changed. Run **1. Prepare / refresh plan** again before requesting GPU time.", gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False), ) def build_demo( submit_function, prepare_function=prepare_separation_with_progress, ): choices = catalog_choices() sample_choices = [("Use upload only", ""), *dataset_sample_choices()] with gr.Blocks(title="SESA Fast Separation", delete_cache=(3600, 21600)) as demo: gr.Markdown( "# SESA Fast Separation\n" "Upload one or more audio or video files, prepare a frozen job plan, then run separation on GPU. " "Preparation validates media, extracts the selected range, downloads models, and creates an isolated job before GPU time is requested." ) with gr.Row(): with gr.Column(scale=3): files = gr.File( label="Audio or video files", file_count="multiple", type="filepath", ) public_sample = gr.Dropdown( label="Or use a public sample", choices=sample_choices, value="", interactive=True, ) public_sample_info = gr.Markdown(dataset_sample_markdown("")) public_sample.change( dataset_sample_markdown, inputs=public_sample, outputs=public_sample_info, show_progress="hidden", ) gr.Markdown("Select up to 6 models total across curated, built-in, and custom sources.") catalog_models = gr.Dropdown( label="Curated models", choices=choices, value=[default_model_id()], multiselect=True, max_choices=6, filterable=True, ) package_models = gr.Dropdown( label="Additional built-in models", choices=SAFE_PACKAGE_MODELS, multiselect=True, max_choices=6, value=[], ) with gr.Column(scale=2): ensemble_algorithm = gr.Dropdown( label="Ensemble algorithm", choices=ENSEMBLE_ALGORITHMS, value="avg_wave", ) output_format = gr.Dropdown( label="Output format", choices=OUTPUT_FORMATS, value="FLAC" ) output_bitrate = gr.Dropdown( label="Output bitrate (compressed formats)", choices=OUTPUT_BITRATES, value=BITRATE_AUTO, ) output_sample_rate = gr.Dropdown( label="Output sample rate", choices=OUTPUT_SAMPLE_RATES, value=44100, ) normalization_threshold = gr.Slider( label="Normalization peak threshold", minimum=0.1, maximum=1.0, step=0.01, value=0.9, ) amplification_threshold = gr.Slider( label="Amplification minimum peak threshold", minimum=0.0, maximum=1.0, step=0.01, value=0.0, ) single_stem = gr.Dropdown( label="Output stems", choices=STEM_CHOICES, value="All stems" ) bundle_layout = gr.Radio( label="Result ZIP layout", choices=[BUNDLE_LAYOUT_FLAT, BUNDLE_LAYOUT_BY_INPUT], value=BUNDLE_LAYOUT_FLAT, ) allow_cpu_fallback = gr.Checkbox( label="Allow limited CPU fallback when CUDA is unavailable", value=False, ) pitch_shift = gr.Slider( label="MDXC pitch shift (semitones)", minimum=-12, maximum=12, step=1, value=0 ) with gr.Accordion("Range and long-file chunking", open=False): gr.Markdown( "Range extraction runs during Preparation, before ZeroGPU. The selected chunk plan is resolved and saved in the frozen job plan." ) range_mode = gr.Radio( label="Processing range", choices=RANGE_MODES, value=RANGE_MODE_FULL ) with gr.Group(visible=False) as range_start_group: range_start_seconds = gr.Number( label="Start position (seconds)", value=0, minimum=0, precision=3 ) with gr.Group(visible=False) as range_end_group: range_end_seconds = gr.Number( label="End position (seconds; 0 means source end)", value=0, minimum=0, precision=3 ) with gr.Group(visible=False) as preview_group: preview_seconds = gr.Slider( label="Preview range length (seconds)", minimum=5, maximum=600, step=5, value=DEFAULT_PREVIEW_SECONDS ) range_mode.change( _range_visibility, inputs=range_mode, outputs=[range_start_group, range_end_group, preview_group], show_progress="hidden", ) chunk_mode = gr.Radio( label="Long-file chunk policy", choices=CHUNK_MODES, value=CHUNK_MODE_AUTO ) with gr.Group(visible=False) as fixed_chunk_group: fixed_chunk_seconds = gr.Dropdown( label="Fixed chunk duration", choices=[("2 minutes", 120), ("5 minutes", 300), ("10 minutes", 600)], value=DEFAULT_FIXED_CHUNK_SECONDS, ) chunk_mode.change( _chunk_visibility, inputs=chunk_mode, outputs=fixed_chunk_group, show_progress="hidden", ) gr.Markdown( "Chunk results are concatenated by audio-separator without crossfade; rare boundary artifacts are possible." ) with gr.Accordion("Batch behavior", open=False): gr.Markdown( "Upload order is preserved. Models load once per prepared job and inputs run sequentially in the same GPU callback. " "The Batch manifest is returned with the outputs and is also included in the result ZIP. " "Stop is cooperative at item boundaries; failed or unprocessed items can be retried from the same frozen Preparation." ) batch_continue_on_item_error = gr.Checkbox( label="Continue with remaining files if one file fails", value=True, ) with gr.Accordion("Custom MDXC / Roformer model", open=False): custom_source = gr.Radio( label="Source", choices=["None", "Hugging Face", "GitHub"], value="None" ) with gr.Group(visible=False) as hf_group: hf_repo = gr.Textbox(label="Hugging Face repository", placeholder="owner/repository") hf_weight = gr.Textbox(label="Weight filename", placeholder="model.ckpt") hf_config = gr.Textbox(label="YAML filename", placeholder="config.yaml") hf_revision = gr.Textbox(label="Revision", value="main") with gr.Group(visible=False) as github_group: github_weight_url = gr.Textbox( label="GitHub release asset URL", placeholder="https://github.com/owner/repo/releases/download/tag/model.ckpt", ) github_config_url = gr.Textbox( label="GitHub YAML URL", placeholder="https://raw.githubusercontent.com/owner/repo/commit/config.yaml", ) github_sha256 = gr.Textbox(label="Weight SHA-256 (recommended)") custom_source.change( _source_visibility, inputs=custom_source, outputs=[hf_group, github_group], show_progress="hidden", ) with gr.Accordion("ZeroGPU time request", open=False): ( duration_mode, manual_gpu_seconds, semi_auto_base_seconds, semi_auto_safety_multiplier, semi_auto_safe_mode, ) = _duration_controls() with gr.Accordion("Progress and logging", open=False): log_level = gr.Radio( label="Log detail", choices=LOG_LEVELS, value=DEFAULT_LOG_LEVEL ) gr.Markdown( "Preparation and GPU execution use separate per-job logs. INFO is compact; DEBUG includes additional package progress detail." ) prepare_inputs = [ files, public_sample, catalog_models, package_models, custom_source, hf_repo, hf_weight, hf_config, hf_revision, github_weight_url, github_config_url, github_sha256, ensemble_algorithm, output_format, output_bitrate, output_sample_rate, normalization_threshold, amplification_threshold, single_stem, bundle_layout, allow_cpu_fallback, pitch_shift, range_mode, range_start_seconds, range_end_seconds, preview_seconds, chunk_mode, fixed_chunk_seconds, duration_mode, manual_gpu_seconds, semi_auto_base_seconds, semi_auto_safety_multiplier, semi_auto_safe_mode, log_level, batch_continue_on_item_error, ] prepared_state = gr.State( "", time_to_live=JOB_TTL_SECONDS, delete_callback=cleanup_prepared_state, ) with gr.Row(): prepare_button = gr.Button("1. Prepare / refresh plan", variant="secondary") run_button = gr.Button("2. Separate prepared job", variant="primary", interactive=False) with gr.Row(): retry_button = gr.Button("Retry incomplete Batch items", variant="secondary", interactive=False) stop_button = gr.Button("Stop after current Batch item", variant="stop", interactive=False) initial_mode = gr.State(EXECUTION_MODE_INITIAL) retry_mode = gr.State(EXECUTION_MODE_RETRY_INCOMPLETE) cancel_status = gr.Markdown() preflight = gr.Markdown("Prepare the job before requesting GPU time.") with gr.Row(): preflight_json = gr.File(label="Prepared job plan (JSON)") preparation_log = gr.File(label="Preparation log") prepare_button.click( prepare_function, inputs=prepare_inputs, outputs=[ prepared_state, preflight, preflight_json, preparation_log, run_button, retry_button, stop_button, ], api_name="prepare", concurrency_limit=2, concurrency_id="sesa_prepare", show_progress="full", ) # Any setting change invalidates the frozen plan. The existing job remains # isolated until TTL cleanup, but cannot be submitted through this session state. for component in prepare_inputs: component.change( _invalidate_preparation, inputs=None, outputs=[prepared_state, preflight, run_button, retry_button, stop_button], show_progress="hidden", ) status = gr.Markdown() with gr.Row(): preview_one = gr.Audio(label="Output preview 1", type="filepath") preview_two = gr.Audio(label="Output preview 2", type="filepath") output_files = gr.File(label="Output files and Batch manifest", file_count="multiple") output_zip = gr.File(label="Result ZIP") with gr.Row(): reproducibility_json = gr.File(label="Reproducibility record (JSON)") job_log_file = gr.File(label="Runtime log file") with gr.Accordion("Runtime log details", open=False): job_log_tail = gr.Textbox( label="Runtime log tail", lines=14, max_lines=28, interactive=False, autoscroll=True ) execution_outputs = [ status, preview_one, preview_two, output_files, output_zip, job_log_tail, job_log_file, reproducibility_json, run_button, retry_button, stop_button, ] run_event = run_button.click( submit_function, inputs=[prepared_state, initial_mode], outputs=execution_outputs, api_name="separate_prepared", concurrency_limit=1, concurrency_id="sesa_gpu", show_progress="full", ) retry_event = retry_button.click( submit_function, inputs=[prepared_state, retry_mode], outputs=execution_outputs, api_name="retry_prepared_batch", concurrency_limit=1, concurrency_id="sesa_gpu", show_progress="full", ) stop_button.click( request_cancel_with_ui, inputs=[prepared_state], outputs=[cancel_status], cancels=[run_event, retry_event], queue=False, api_name="request_batch_cancel", api_visibility="private", show_progress="hidden", ) gr.Markdown( "Only the model cache is shared. Inputs, outputs, logs, and job configuration are isolated per job and removed automatically after the retention period." ) return demo