Inference

def __enter__(self) -> Stage:
    return self

def __exit__(self, *_: object) -> None:
    return None

@classmethod
def from_pretrained(
    cls,
    *,
    config: IntoConfig,
    checkpoint_path: t.StrPath,
    overrides: ConfigOverrides = (),
    model_device: torch.device | str | None = None,
    io_device: torch.device | str | None = None,
    module_name: str | None = None,
    class_name: str | None = None,
    package_name: str | None = None,
) -> InferenceEngine:
    from .config import into_config
    from .io import load_weights
    from .models import ModelMetadata

    config = into_config(config, overrides=overrides)

    model_device_resolved = _resolve_device(
        model_device or config.inference.model_device,
        field_name="inference.model_device",
    )
    io_device_resolved = _resolve_device(
        io_device or config.inference.io_device,
        field_name="inference.io_device",
    )
    resolved_module, resolved_class = resolve_model_entrypoint(
        config.model_type, module_name, class_name
    )
    metadata = ModelMetadata.from_module(
        module_name=resolved_module,
        model_cls_name=resolved_class,
        model_type=config.model_type,
        package=package_name,
    )

    model_params = config.model.to_concrete(metadata.params)
    full_output_stems = tuple(config.model.output_stem_names)
    requested_stems = tuple(config.inference.requested_stems or full_output_stems)

    state_dict_transform = None
    if requested_stems != full_output_stems:
        # optional model-level optimization contract: models can choose to
        # provide a stem-selection plan that may mutate params and checkpoint
        # loading. if absent, we keep full model outputs and discard
        # unrelated stems immediately after each forward pass.
        from .models import SupportsStemSelection

        if isinstance(metadata.model, SupportsStemSelection):
            plan = metadata.model.__splifft_stem_selection_plan__(model_params, requested_stems)
            model_params = plan.model_params
            state_dict_transform = plan.state_dict_transform

    model = metadata.model(model_params)
    if (forced_dtype := config.inference.force_weights_dtype) is not None:
        model = model.to(_resolve_runtime_dtype(forced_dtype, device=model_device_resolved))
    model = load_weights(
        model,
        checkpoint_path,
        device=model_device_resolved,
        state_dict_transform=state_dict_transform,
    ).eval()

    # maybe we should to an explicit try_compile() method while emitting events but eh.
    # we shuold probably log since it can take extremely long
    if (compile_cfg := config.inference.compile_model) is not None:
        compiled_model = torch.compile(
            model,
            fullgraph=compile_cfg.fullgraph,
            dynamic=compile_cfg.dynamic,
            mode=compile_cfg.mode,
        )
        model = cast(nn.Module, compiled_model)

    return cls(
        config=config,
        model=model,
        model_params_concrete=model_params,
        model_device=model_device_resolved,
        io_device=io_device_resolved,
        model_input_dtype=core.get_model_floating_dtype(model),
    )

@classmethod
def from_registry(
    cls,
    model_id: str,
    *,
    model_device: torch.device | str | None = None,
    io_device: torch.device | str | None = None,
    overrides: ConfigOverrides = (),
    fetch_if_missing: bool = True,
    force_overwrite_config: bool = False,
    force_overwrite_model: bool = False,
    registry_path: Path = PATH_REGISTRY_DEFAULT,
) -> InferenceEngine:
    from .config import Registry
    from .io import get_model_paths

    model_paths = get_model_paths(
        model_id,
        fetch_if_missing=fetch_if_missing,
        force_overwrite_config=force_overwrite_config,
        force_overwrite_model=force_overwrite_model,
        registry=Registry.from_file(registry_path),
    )
    return cls.from_pretrained(
        config=model_paths.path_config,
        checkpoint_path=model_paths.path_checkpoint,
        overrides=overrides,
        model_device=model_device,
        io_device=io_device,
    )

def to_audio_tensor(
    self,
    mixture: t.StrPath | t.BytesPath | t.RawAudioTensor | core.Audio[t.RawAudioTensor],
) -> core.Audio[t.RawAudioTensor]:
    if isinstance(mixture, core.Audio):
        return mixture
    elif isinstance(mixture, torch.Tensor):
        return core.Audio(
            data=t.RawAudioTensor(mixture),
            sample_rate=self.config.audio_io.target_sample_rate,
        )
    else:
        from .io import read_audio

        return read_audio(
            mixture,  # type: ignore[arg-type]
            self.config.audio_io.target_sample_rate,
            self.config.audio_io.force_channels,
            device=self.io_device,
        )

def run(
    self,
    mixture: t.StrPath | t.BytesPath | t.RawAudioTensor | core.Audio[t.RawAudioTensor],
) -> InferenceOutput:
    for event in self.stream(mixture):
        if isinstance(event, InferenceOutput):
            return event
    raise RuntimeError("inference stream finished without outputs")

def stream(
    self,
    mixture: t.StrPath | t.BytesPath | t.RawAudioTensor | core.Audio[t.RawAudioTensor],
) -> Generator[InferenceEvent, None, None]:
    archetype = self.config.validate_inference_contract(self.model_params_concrete)

    audio_tensor = self.to_audio_tensor(mixture)
    raw_mixture_data = t.RawAudioTensor(audio_tensor.data.to(self.io_device))
    mixture_data: t.RawAudioTensor | t.NormalizedAudioTensor = raw_mixture_data
    mixture_stats: core.NormalizationStats | None = None

    if self.config.normalization.enabled:
        with Stage("normalize") as s:
            yield s.started
            normalized = core.normalize_audio(
                core.Audio(data=raw_mixture_data, sample_rate=audio_tensor.sample_rate)
            )
            mixture_data = normalized.audio.data
            mixture_stats = normalized.stats
            yield s.completed

    mixture_data = yield from self._adapt_input_channels(mixture_data)

    if archetype == "sequence_labeling":
        requested_stems = self._requested_output_stem_names()
        requested_stem_indices = self._requested_output_stem_indices()
        sequence_outputs = yield from self._stream_sequence_labeling(
            mixture_data,
            requested_stems=requested_stems,
            output_indices=requested_stem_indices,
        )
        yield InferenceOutput(outputs=sequence_outputs, sample_rate=audio_tensor.sample_rate)
        return

    requested_stems = self._requested_output_stem_names()
    requested_stem_indices = self._requested_output_stem_indices()
    separated_data = yield from self._stream_waveform_pipeline(
        mixture_data,
        archetype,
        output_indices=requested_stem_indices,
        num_stems=len(requested_stems),
    )

    denormalized_stems: dict[t.ModelOutputStemName, t.RawAudioTensor] = {}
    with Stage("collect_outputs") as s:
        yield s.started
        for i, stem_name in enumerate(requested_stems):
            stem_data = separated_data[i, ...]
            if mixture_stats is not None:
                stem_data = core.denormalize_audio(
                    audio_data=t.NormalizedAudioTensor(stem_data),
                    stats=mixture_stats,
                )
            denormalized_stems[stem_name] = t.RawAudioTensor(stem_data)
        yield s.completed

    output_stems: dict[StemName, t.RawAudioTensor] = denormalized_stems
    if derived_stems_cfg := self.config.derived_stems:
        with Stage("derive_stems") as s:
            yield s.started
            output_stems = core.derive_stems(
                denormalized_stems,
                raw_mixture_data,
                derived_stems_cfg,
            )
            yield s.completed

    yield InferenceOutput(outputs=output_stems, sample_rate=audio_tensor.sample_rate)