from __future__ import annotations import math from dataclasses import asdict, dataclass from pathlib import Path from .catalog import load_catalog from .config import VIDEO_EXTENSIONS from .media import normalize_uploads, probe_duration DURATION_MODE_MANUAL = "Manual (fixed seconds)" DURATION_MODE_SEMI_AUTO = "Semi-auto (input/model estimate)" DURATION_MODES = [DURATION_MODE_SEMI_AUTO, DURATION_MODE_MANUAL] ZERO_GPU_MIN_SECONDS = 30 ZERO_GPU_MAX_SECONDS = 300 DEFAULT_MANUAL_SECONDS = 180 DEFAULT_BASE_SECONDS_PER_MINUTE = 12.0 DEFAULT_SAFETY_MULTIPLIER = 1.25 DEFAULT_SAFE_MODE_ENABLED = False SAFE_MODE_MULTIPLIER = 1.30 # Calibration runs showed that the first inference in a fresh ZeroGPU callback # can pay a sizeable CUDA/ONNX kernel warm-up cost that is not proportional to # source duration. Keep this as an explicit, inspectable term rather than hiding # it in the per-minute multiplier. ONNX_FIRST_INFERENCE_WARMUP_SECONDS = 30.0 NON_ONNX_FIRST_INFERENCE_WARMUP_SECONDS = 20.0 ADDITIONAL_MODEL_WARMUP_SECONDS = 8.0 ADDITIONAL_BATCH_ITEM_DISPATCH_SECONDS = 1.5 DURATION_CALIBRATION_REVISION = "cold-start-batch-20260718" @dataclass(frozen=True) class DurationEstimate: mode: str seconds: int unclamped_seconds: float base_unclamped_seconds: float total_input_seconds: float file_count: int video_count: int model_count: int model_factor_sum: float parameter_factor: float base_seconds_per_minute: float safety_multiplier: float safe_mode_requested: bool safe_mode_applied: bool safe_mode_multiplier: float fixed_overhead_seconds: float first_inference_warmup_seconds: float additional_batch_item_overhead_seconds: float onnx_model_count: int non_onnx_model_count: int calibration_revision: str unprobed_files: int def to_dict(self) -> dict: return asdict(self) def _clamp_seconds(value: float) -> int: return int(max(ZERO_GPU_MIN_SECONDS, min(ZERO_GPU_MAX_SECONDS, math.ceil(value)))) def _input_stats(files) -> tuple[float, int, int, int]: total_seconds = 0.0 file_count = 0 video_count = 0 unprobed = 0 for path in normalize_uploads(files): file_count += 1 if path.suffix.lower() in VIDEO_EXTENSIONS: video_count += 1 try: seconds = float(probe_duration(path)) if not math.isfinite(seconds) or seconds <= 0: raise ValueError("duration unavailable") except Exception: # Keep the duration callable non-fatal. One minute is a conservative # fallback for an upload that ffprobe cannot inspect before dispatch. seconds = 60.0 unprobed += 1 total_seconds += seconds if file_count == 0: total_seconds = 60.0 file_count = 1 unprobed = 1 return total_seconds, file_count, video_count, unprobed def _catalog_factor(model_id: str) -> float: item = load_catalog().get(str(model_id)) if not item: return 1.15 factor = 1.0 category = str(item.get("category", "")).lower() text = " ".join( str(item.get(key, "")) for key in ("name", "model_filename", "architecture") ).lower() if "4-stem" in category or "4stem" in text or "6stem" in text: factor *= 1.45 if "scnet" in text: factor *= 1.15 if any(token in text for token in (" xl", "large", "super_big", "super big", "big_beta")): factor *= 1.15 return factor def _package_factor(filename: str) -> float: suffix = Path(str(filename)).suffix.lower() if suffix == ".onnx": return 0.75 if suffix == ".pth": return 0.95 return 1.0 def _model_stats(catalog_model_ids, package_model_filenames, custom_source: str) -> tuple[int, float]: catalog_ids = list(catalog_model_ids or []) package_files = list(package_model_filenames or []) factors = [_catalog_factor(model_id) for model_id in catalog_ids] factors.extend(_package_factor(filename) for filename in package_files) if str(custom_source or "None") in {"Hugging Face", "GitHub"}: factors.append(1.25) if not factors: factors = [1.0] return len(factors), float(sum(factors)) def _model_runtime_profile(catalog_model_ids, package_model_filenames, custom_source: str) -> tuple[int, int]: """Return ONNX/non-ONNX model counts for cold-start calibration.""" filenames: list[str] = [] catalog = load_catalog() for model_id in list(catalog_model_ids or []): item = catalog.get(str(model_id)) or {} filename = str(item.get("model_filename") or item.get("filename") or "") filenames.append(filename) filenames.extend(str(value or "") for value in list(package_model_filenames or [])) custom_count = 1 if str(custom_source or "None") in {"Hugging Face", "GitHub"} else 0 onnx_count = sum(1 for filename in filenames if Path(filename).suffix.lower() == ".onnx") non_onnx_count = max(0, len(filenames) - onnx_count) + custom_count if onnx_count + non_onnx_count == 0: non_onnx_count = 1 return onnx_count, non_onnx_count def calculate_duration_estimate( files, catalog_model_ids, package_model_filenames, custom_source, ensemble_algorithm, output_format, pitch_shift, chunk_duration, duration_mode, manual_gpu_seconds, semi_auto_base_seconds, semi_auto_safety_multiplier, semi_auto_safe_mode=DEFAULT_SAFE_MODE_ENABLED, total_input_seconds_override=None, ) -> DurationEstimate: mode = str(duration_mode or DURATION_MODE_SEMI_AUTO) total_seconds, file_count, video_count, unprobed = _input_stats(files) if total_input_seconds_override is not None: try: override = float(total_input_seconds_override) except (TypeError, ValueError): override = 0.0 if math.isfinite(override) and override > 0: total_seconds = override unprobed = 0 model_count, model_factor_sum = _model_stats( catalog_model_ids, package_model_filenames, str(custom_source or "None") ) onnx_model_count, non_onnx_model_count = _model_runtime_profile( catalog_model_ids, package_model_filenames, str(custom_source or "None") ) manual = float(manual_gpu_seconds or DEFAULT_MANUAL_SECONDS) base = max(1.0, float(semi_auto_base_seconds or DEFAULT_BASE_SECONDS_PER_MINUTE)) safety = max(0.5, float(semi_auto_safety_multiplier or DEFAULT_SAFETY_MULTIPLIER)) safe_requested = bool(semi_auto_safe_mode) if mode == DURATION_MODE_MANUAL: seconds = _clamp_seconds(manual) return DurationEstimate( mode=mode, seconds=seconds, unclamped_seconds=manual, base_unclamped_seconds=manual, total_input_seconds=total_seconds, file_count=file_count, video_count=video_count, model_count=model_count, model_factor_sum=model_factor_sum, parameter_factor=1.0, base_seconds_per_minute=base, safety_multiplier=safety, safe_mode_requested=safe_requested, safe_mode_applied=False, safe_mode_multiplier=1.0, fixed_overhead_seconds=0.0, first_inference_warmup_seconds=0.0, additional_batch_item_overhead_seconds=0.0, onnx_model_count=onnx_model_count, non_onnx_model_count=non_onnx_model_count, calibration_revision=DURATION_CALIBRATION_REVISION, unprobed_files=unprobed, ) parameter_factor = 1.0 if int(pitch_shift or 0) != 0: parameter_factor *= 1.12 if int(chunk_duration or 0) > 0: parameter_factor *= 1.08 if str(output_format or "FLAC").upper() in {"MP3", "OGG", "OPUS", "M4A"}: parameter_factor *= 1.06 algorithm = str(ensemble_algorithm or "avg_wave").lower() if "fft" in algorithm: parameter_factor *= 1.08 elif algorithm.startswith("uvr_") or algorithm == "ensemble_wav": parameter_factor *= 1.12 if model_count > 1: parameter_factor *= 1.0 + min(0.25, 0.05 * (model_count - 1)) fixed_overhead = 18.0 + 4.0 * model_count + 2.0 * file_count + 4.0 * video_count first_inference_warmup = ( ONNX_FIRST_INFERENCE_WARMUP_SECONDS if onnx_model_count > 0 else NON_ONNX_FIRST_INFERENCE_WARMUP_SECONDS ) if model_count > 1: first_inference_warmup += ADDITIONAL_MODEL_WARMUP_SECONDS * (model_count - 1) additional_batch_item_overhead = ADDITIONAL_BATCH_ITEM_DISPATCH_SECONDS * max(0, file_count - 1) compute = (total_seconds / 60.0) * base * model_factor_sum * parameter_factor base_raw = ( fixed_overhead + first_inference_warmup + additional_batch_item_overhead + compute ) * safety safe_applied = safe_requested effective_safe_multiplier = SAFE_MODE_MULTIPLIER if safe_applied else 1.0 raw = base_raw * effective_safe_multiplier seconds = _clamp_seconds(raw) return DurationEstimate( mode=mode, seconds=seconds, unclamped_seconds=raw, base_unclamped_seconds=base_raw, total_input_seconds=total_seconds, file_count=file_count, video_count=video_count, model_count=model_count, model_factor_sum=model_factor_sum, parameter_factor=parameter_factor, base_seconds_per_minute=base, safety_multiplier=safety, safe_mode_requested=safe_requested, safe_mode_applied=safe_applied, safe_mode_multiplier=effective_safe_multiplier, fixed_overhead_seconds=fixed_overhead, first_inference_warmup_seconds=first_inference_warmup, additional_batch_item_overhead_seconds=additional_batch_item_overhead, onnx_model_count=onnx_model_count, non_onnx_model_count=non_onnx_model_count, calibration_revision=DURATION_CALIBRATION_REVISION, unprobed_files=unprobed, ) def _extract_duration_args(args, kwargs) -> tuple: def value(index: int, name: str, default=None): if len(args) > index: return args[index] return kwargs.get(name, default) return ( value(0, "files"), value(1, "catalog_model_ids", []), value(2, "package_model_filenames", []), value(3, "custom_source", "None"), value(11, "ensemble_algorithm", "avg_wave"), value(12, "output_format", "FLAC"), value(14, "pitch_shift", 0), value(15, "chunk_duration", 0), value(16, "duration_mode", DURATION_MODE_SEMI_AUTO), value(17, "manual_gpu_seconds", DEFAULT_MANUAL_SECONDS), value(18, "semi_auto_base_seconds", DEFAULT_BASE_SECONDS_PER_MINUTE), value(19, "semi_auto_safety_multiplier", DEFAULT_SAFETY_MULTIPLIER), value(20, "semi_auto_safe_mode", DEFAULT_SAFE_MODE_ENABLED), ) def estimate_gpu_seconds(*args, **kwargs) -> int: try: return calculate_duration_estimate(*_extract_duration_args(args, kwargs)).seconds except Exception: return DEFAULT_MANUAL_SECONDS def estimate_gpu_details(*args, **kwargs) -> DurationEstimate: return calculate_duration_estimate(*_extract_duration_args(args, kwargs)) def format_duration_estimate(*args, **kwargs) -> str: estimate = estimate_gpu_details(*args, **kwargs) if estimate.mode == DURATION_MODE_MANUAL: return f"**ZeroGPU duration: {estimate.seconds} seconds** — manual fixed value; Safe mode is not applied." fallback = ( f" · ffprobe fallback used for {estimate.unprobed_files} file(s)" if estimate.unprobed_files else "" ) safe_text = ( f"Safe mode ON: base {math.ceil(estimate.base_unclamped_seconds)}s × " f"{estimate.safe_mode_multiplier:.2f} → {estimate.seconds}s" if estimate.safe_mode_applied else f"Safe mode OFF: {estimate.seconds}s" ) return ( f"**ZeroGPU duration estimate: {estimate.seconds} seconds** — {safe_text} \n" f"Input {estimate.total_input_seconds:.1f}s / {estimate.file_count} file(s), " f"models {estimate.model_count} (factor sum {estimate.model_factor_sum:.2f}), " f"cold-start {estimate.first_inference_warmup_seconds:.0f}s, " f"Batch dispatch {estimate.additional_batch_item_overhead_seconds:.1f}s, " f"parameter factor {estimate.parameter_factor:.2f}, calibration {estimate.safety_multiplier:.2f}" f"{fallback}." )