allow to use any pyannote pretrained pipeline

app.py

@@ -2,13 +2,18 @@ import gradio as gr
 from pyannote_viewer import PyannoteViewer
 from pyannote.audio import Pipeline
 import os
+from huggingface_hub import HfApi
 
-
-def apply_pipeline(audio: str) -> tuple:
+def apply_pipeline(audio: str, pipeline_name: str) -> tuple:
     pipeline = Pipeline.from_pretrained(
-        "pyannote/speech-separation-ami-1.0", use_auth_token=os.environ["HF_TOKEN"]
+        pipeline_name, use_auth_token=os.environ["HF_TOKEN"]
     )
-    return pipeline(audio)
+
+    outputs = pipeline(audio)
+    if isinstance(outputs, tuple):
+        return outputs
+    else:
+        return (outputs, audio)
 
 
 with gr.Blocks() as demo:
@@ -22,20 +27,29 @@ with gr.Blocks() as demo:
             )
         # space title and description
         with gr.Column(scale=10):
-            gr.Markdown('# Speaker diarization and speech separation pipeline')
+            gr.Markdown('# pyannote pretrained pipelines')
     
             gr.Markdown(
-                "This space is dedicated to showing the use of the speaker diarization and speech separation [pipeline](https://huggingface.co/pyannote/speech-separation-ami-1.0) integrated to `pyannote.audio`. To use this space:"
-                "\n - Load or record an audio"
-                "\n - Click on the apply pipeline button"
-                "\n - After pipeline processed the audio, you can then listen for each speaker separetely. Annotations on waveforms correspond to the speaker diarization produced by the pipeline, with one color per speaker."
+                "You like [pyannote.audio](https://github.com/pyannote/pyannote-audio)? Consider using [pyannoteAI](https://pyannote.ai/) for better and faster options.\n"
+                "\nGo [here](https://huggingface.co/pyannote) for more detail on each pipeline available in this space."
             )
+            
+            gr.Markdown()
+
 
+    gr.Markdown("#### Select a pretrained pipeline:")
+    available_pipelines = [p.modelId for p in HfApi().list_models(filter="pyannote-audio-pipeline")]
+    available_pipelines = list(filter(lambda p: p.startswith("pyannote/"), available_pipelines))
+    dropdown = gr.Dropdown(choices=available_pipelines, value=available_pipelines[0], interactive=True, label="Pretrained pipeline")
+
+    gr.Markdown("#### Upload or record an audio:")
     audio = gr.Audio(type="filepath")
-    btn = gr.Button("Apply separation pipeline")
+
+    btn = gr.Button("Apply pipeline")
+
     source_viewer = PyannoteViewer(interactive=False)
 
-    btn.click(fn=apply_pipeline, inputs=[audio], outputs=[source_viewer])
+    btn.click(fn=apply_pipeline, inputs=[audio, dropdown], outputs=[source_viewer])
 
 
 if __name__ == "__main__":

pyannote_viewer/README.md

@@ -1,20 +1,20 @@
 
-# `gradio_sourceviewer`
+# `pyannote_viewer`
 <img alt="Static Badge" src="https://img.shields.io/badge/version%20-%201.0.0%20-%20orange">  
 
-Python library for easily interacting with trained machine learning models
+Gradio custom component to visualize pyannote's pipelines outputs
 
 ## Installation
 
 ```bash
-pip install gradio_sourceviewer
+pip install pyannote-viewer
 ```
 
 ## Usage
 
 ```python
 import gradio as gr
-from gradio_sourceviewer import SourceViewer
+from pyannote_viewer import PyannoteViewer
 from pyannote.audio import Pipeline
 import os
 
@@ -29,9 +29,9 @@ def apply_pipeline(audio: str) -> tuple:
 with gr.Blocks() as demo:
     audio = gr.Audio(type="filepath")
     btn = gr.Button("Apply separation pipeline")
-    source_viewer = SourceViewer(interactive=False)
+    pyannote_viewer = PyannoteViewer(interactive=False)
 
-    btn.click(fn=apply_pipeline, inputs=[audio], outputs=[source_viewer])
+    btn.click(fn=apply_pipeline, inputs=[audio], outputs=[pyannote_viewer])
 
 
 if __name__ == "__main__":
@@ -39,7 +39,7 @@ if __name__ == "__main__":
 
 ```
 
-## `SourceViewer`
+## `PyannoteViewer`
 
 ### Initialization
 
@@ -362,16 +362,16 @@ WaveformOptions | dict | None
 
 | name | description |
 |:-----|:------------|
-| `stream` | This listener is triggered when the user streams the SourceViewer. |
-| `change` | Triggered when the value of the SourceViewer changes either because of user input (e.g. a user types in a textbox) OR because of a function update (e.g. an image receives a value from the output of an event trigger). See `.input()` for a listener that is only triggered by user input. |
-| `clear` | This listener is triggered when the user clears the SourceViewer using the X button for the component. |
-| `play` | This listener is triggered when the user plays the media in the SourceViewer. |
-| `pause` | This listener is triggered when the media in the SourceViewer stops for any reason. |
-| `stop` | This listener is triggered when the user reaches the end of the media playing in the SourceViewer. |
-| `start_recording` | This listener is triggered when the user starts recording with the SourceViewer. |
-| `pause_recording` | This listener is triggered when the user pauses recording with the SourceViewer. |
-| `stop_recording` | This listener is triggered when the user stops recording with the SourceViewer. |
-| `upload` | This listener is triggered when the user uploads a file into the SourceViewer. |
+| `stream` | This listener is triggered when the user streams the PyannoteViewer. |
+| `change` | Triggered when the value of the PyannoteViewer changes either because of user input (e.g. a user types in a textbox) OR because of a function update (e.g. an image receives a value from the output of an event trigger). See `.input()` for a listener that is only triggered by user input. |
+| `clear` | This listener is triggered when the user clears the PyannoteViewer using the X button for the component. |
+| `play` | This listener is triggered when the user plays the media in the PyannoteViewer. |
+| `pause` | This listener is triggered when the media in the PyannoteViewer stops for any reason. |
+| `stop` | This listener is triggered when the user reaches the end of the media playing in the PyannoteViewer. |
+| `start_recording` | This listener is triggered when the user starts recording with the PyannoteViewer. |
+| `pause_recording` | This listener is triggered when the user pauses recording with the PyannoteViewer. |
+| `stop_recording` | This listener is triggered when the user stops recording with the PyannoteViewer. |
+| `upload` | This listener is triggered when the user uploads a file into the PyannoteViewer. |
 
 
 
@@ -391,7 +391,8 @@ The code snippet below is accurate in cases where the component is used as both
  def predict(
      value: str | tuple[int, numpy.ndarray] | None
  ) -> tuple[
-        pyannote.core.annotation.Annotation, numpy.ndarray
+        pyannote.core.annotation.Annotation,
+        numpy.ndarray | pathlib.Path | str,
     ]
     | None:
      return value

pyannote_viewer/backend/pyannote_viewer/pyannote_viewer.py

@@ -19,6 +19,7 @@ from gradio.events import Events
 from gradio.exceptions import Error
 
 from pyannote.core.annotation import Annotation
+from pyannote.core.feature import SlidingWindowFeature
 
 
 @dataclasses.dataclass
@@ -249,7 +250,7 @@ class PyannoteViewer(
             )
 
     def postprocess(
-        self, value: Tuple[Annotation, np.ndarray] | None
+        self, value: Tuple[Annotation, np.ndarray | Path | str] | None
     ) -> FileData | bytes | None:
         """
         Parameters:
@@ -260,7 +261,8 @@ class PyannoteViewer(
         if value is None:
             return None
 
-        annotations, sources = value
+        annotations, audio = value
+
         labels = annotations.labels()
 
         # format diarization output
@@ -271,19 +273,28 @@ class PyannoteViewer(
                 Segment(start=segment.start, end=segment.end, channel=label_idx)
             )
 
-        # save sources in cache
-        source_filepath = processing_utils.save_audio_to_cache(
-            data=sources.data,
-            sample_rate=16_000,
-            format=self.format,
-            cache_dir=self.GRADIO_CACHE,
-        )
-        orig_name = Path(source_filepath).name
+        if isinstance(audio, SlidingWindowFeature):
+            # save sources in cache
+            audio_filepath = processing_utils.save_audio_to_cache(
+                data=audio.data,
+                sample_rate=16_000,
+                format=self.format,
+                cache_dir=self.GRADIO_CACHE,
+            )
+            multichannel = True
+        elif isinstance(audio, (Path, str)):
+            audio_filepath = audio
+            multichannel = False
+        else:
+            raise ValueError("Unknown type for audio value")
+
+        orig_name = Path(audio_filepath).name
 
         return {
             "segments": segments,
             "labels": labels,
-            "sources_file": FileData(path=source_filepath, orig_name=orig_name),
+            "multichannel": multichannel,
+            "audio_file": FileData(path=audio_filepath, orig_name=orig_name),
         }
 
     def stream_output(

pyannote_viewer/frontend/Index.svelte

@@ -10,13 +10,13 @@
 	import InteractiveAudio from "./interactive/InteractiveAudio.svelte";
 	import { StatusTracker } from "@gradio/statustracker";
 	import { Block, UploadText } from "@gradio/atoms";
-	import type { WaveformOptions, Segment } from "./shared/types";
+	import type { WaveformOptions, PipelineOutput } from "./shared/types";
 
 	export let elem_id = "";
 	export let elem_classes: string[] = [];
 	export let visible = true;
 	export let interactive: boolean;
-	export let value: null | {"segments": Segment[], "labels" : string[], "sources_file": FileData} = null;
+	export let value: PipelineOutput | null = null;
 	export let sources:
 		| ["microphone"]
 		| ["upload"]
@@ -106,7 +106,6 @@
 		minPxPerSec: 20,
 		mediaControls: waveform_options.show_controls,
 		sampleRate: waveform_options.sample_rate || 44100,
-		splitChannels: true,
 	};
 
 

pyannote_viewer/frontend/interactive/InteractiveAudio.svelte

@@ -16,9 +16,9 @@
 	import AudioRecorder from "../recorder/AudioRecorder.svelte";
 	import StreamAudio from "../streaming/StreamAudio.svelte";
 	import { SelectSource } from "@gradio/atoms";
-	import type { WaveformOptions, Segment } from "../shared/types";
+	import type { WaveformOptions, PipelineOutput } from "../shared/types";
 
-	export let value: null | {"segments": Segment[], "labels": string[], "sources_file": FileData} = null;
+	export let value: PipelineOutput | null = null;
 	export let label: string;
 	export let root: string;
 	export let show_label = true;
@@ -74,7 +74,7 @@
 
 	const dispatch = createEventDispatcher<{
 		change: typeof value;
-		stream: FileData;
+		stream: typeof value;
 		edit: never;
 		play: never;
 		pause: never;
@@ -95,7 +95,7 @@
 	): Promise<void> => {
 		let _audio_blob = new File(blobs, "audio.wav");
 		const val = await prepare_files([_audio_blob], event === "stream");
-		value.sources_file = (
+		value.audio_file = (
 			(await upload(val, root, undefined, upload_fn))?.filter(
 				Boolean
 			) as FileData[]
@@ -192,8 +192,8 @@
 	}
 
 	function handle_load({ detail }: { detail: FileData }): void {
-		value = {"segments": [], "labels": [], "sources_file": null}
-		value.sources_file = detail;
+		value = {"segments": [], "labels": [], "multichannel": false, "audioFile": null}
+		value.audio_file = detail;
 		dispatch("change", value);
 		dispatch("upload", detail);
 	}
@@ -264,7 +264,7 @@
 			{i18n}
 			on:clear={clear}
 			on:edit={() => (mode = "edit")}
-			download={show_download_button ? value.sources_file.url : null}
+			download={show_download_button ? value.audio_file.url : null}
 			absolute={true}
 		/>
 

pyannote_viewer/frontend/player/AudioPlayer.svelte

@@ -3,15 +3,15 @@
 	import { Music } from "@gradio/icons";
 	import { format_time, type I18nFormatter } from "@gradio/utils";
 	import WaveSurfer from "wavesurfer.js";
-	import RegionsPlugin, {type Region} from "wavesurfer.js/dist/plugins/regions";
-	import { skip_audio, process_audio } from "../shared/utils";
+	import RegionsPlugin from "wavesurfer.js/dist/plugins/regions";
+	import { skip_audio } from "../shared/utils";
 	import WaveformControls from "../shared/WaveformControls.svelte";
 	import { Empty } from "@gradio/atoms";
-	import type { FileData } from "@gradio/client";
-	import type { WaveformOptions, Segment } from "../shared/types";
+	import type { WaveformOptions, PipelineOutput } from "../shared/types";
 	import { createEventDispatcher } from "svelte";
+    import Color from "@gradio/icons/src/Color.svelte";
 
-	export let value: null | {"segments": Segment[], "labels": string[], "sources_file": FileData}= null;
+	export let value: PipelineOutput | null = null;
 	export let label: string;
 	export let root: string;
 	export let i18n: I18nFormatter;
@@ -50,7 +50,7 @@
 	}>();
 
 	const create_waveform = (): void => {
-		const audio = new Audio(root + `/file=${value.sources_file.path}`)
+		const audio = new Audio(root + `/file=${value.audio_file.path}`)
 		audio.crossOrigin = "anonymous"
 
 		audioContext = new AudioContext();
@@ -58,6 +58,7 @@
 		waveform = WaveSurfer.create({
 			container: container,
 			media: audio,
+			splitChannels: value.multichannel,
 			...waveform_settings
 		});
 	};
@@ -72,6 +73,7 @@
 	$: waveform?.on("decode", (duration: any) => {
 		audioDecoded = true;
 		const numChannels = waveform.getDecodedData().numberOfChannels;
+		console.log(numChannels);
 		audio_duration = duration;
 		durationRef && (durationRef.textContent = format_time(duration));
 	
@@ -80,6 +82,10 @@
 		splitter = audioContext.createChannelSplitter(numChannels);
 		mediaNode.connect(splitter);
 
+		if(!value.multichannel){
+			splitter.connect(audioContext.destination, 0);
+		}
+
 		// add diarization annotation on each source:
 		if(!wsRegion){
 			wsRegion = waveform.registerPlugin(RegionsPlugin.create())
@@ -87,13 +93,13 @@
 				const region = wsRegion.addRegion({
 					start: segment.start,
 					end: segment.end,
-					channelIdx: segment.channel,
+					channelIdx: value.multichannel ? segment.channel : 0,
 					drag: false,
 					resize: false,
 					color: colors[segment.channel % colors.length],
 				});
 
-				const regionHeight = 100 / numChannels;
+				const regionHeight = 100 / (value.multichannel ? numChannels : 1);
 				region.element.style.cssText += `height: ${regionHeight}% !important;`;
 				// TODO: Can we do better than force region color ?
 				region.element.style.cssText += `background-color: ${region.color} !important;`;
@@ -144,10 +150,10 @@
 	<Empty size="small">
 		<Music />
 	</Empty>
-{:else if value.sources_file.is_stream}
+{:else if value.audio_file.is_stream}
 	<audio
 		class="standard-player"
-		src={value.sources_file.url}
+		src={value.audio_file.url}
 		controls
 		autoplay={waveform_settings.autoplay}
 	/>
@@ -159,20 +165,30 @@
 	<div class="viewer">
 		<div class="source-selection">
 			{#if audioDecoded}
-				{#each [...Array(waveform.getDecodedData().numberOfChannels).keys()] as channelIdx}
-					<label class="source" style={`height: ${waveform_settings.height}px`}>
-						<input 
-							type="radio" 
-							name="channels" 
-							value={`${channelIdx}`}
-							on:change={(ev) => {
-								splitter.disconnect()
-								splitter.connect(audioContext.destination, Number(ev.target.value), 0);
-							}}
-						/>
-						{value.labels[channelIdx]}
-					</label>
-				{/each}
+				{#if value.multichannel}
+					<!-- Separation pipeline case -->
+					{#each [...Array(waveform.getDecodedData().numberOfChannels).keys()] as channelIdx}
+						<label style={`height: ${waveform_settings.height}px; background-color: ${colors[channelIdx % colors.length]}`}>
+							<input 
+								type="radio" 
+								name="channels" 
+								value={`${channelIdx}`}
+								on:change={(ev) => {
+									splitter.disconnect()
+									splitter.connect(audioContext.destination, Number(ev.target.value), 0);
+								}}
+							/>
+							{value.labels[channelIdx]}
+						</label>
+					{/each}
+				{:else}
+						{#each [...Array(value.labels.length)].keys() as labelIdx}
+							<label style={`background-color: ${colors[labelIdx % colors.length]};`}>
+								<input type="hidden">
+								{value.labels[labelIdx]}
+							</label>
+						{/each}
+				{/if}
 			{/if}
 		</div>
 		<div class="waveform-container">
@@ -225,6 +241,15 @@
 		background-color: var(--color-accent);
 	}
 
+	label {
+		display: flex;
+		align-items: center;
+		margin-bottom: 0.25em;
+		padding-left: 0.5em;
+		padding-right: 0.5em;
+	}
+
+
 	.component-wrapper {
 		padding: var(--size-3);
 		width: 100%;
@@ -240,11 +265,6 @@
 		margin-right: 1em;
 	}
 
-	.source {
-		display: flex;
-		align-items: center;
-	}
-
 	:global(::part(wrapper)) {
 		margin-bottom: var(--size-2);
 	}

pyannote_viewer/frontend/shared/types.ts

@@ -1,3 +1,5 @@
+import type { FileData } from "@gradio/client";
+
 export type WaveformOptions = {
 	waveform_color?: string;
 	waveform_progress_color?: string;
@@ -13,3 +15,10 @@ export type Segment = {
 	end: number;
 	channel: number;
 }
+
+export type PipelineOutput = {
+	segments: Segment[];
+	labels: string[];
+	multichannel: boolean;
+	audio_file: FileData;
+}

pyannote_viewer/frontend/static/StaticAudio.svelte

@@ -6,11 +6,10 @@
 	import type { I18nFormatter } from "@gradio/utils";
 	import AudioPlayer from "../player/AudioPlayer.svelte";
 	import { createEventDispatcher } from "svelte";
-	import type { FileData } from "@gradio/client";
 	import { DownloadLink } from "@gradio/wasm/svelte";
-	import type { WaveformOptions, Segment } from "../shared/types";
+	import type { WaveformOptions, PipelineOutput } from "../shared/types";
 
-	export let value: null | {"segments": Segment[], "labels":string[], "sources_file": FileData} = null;
+	export let value: PipelineOutput | null = null;
 	export let label: string;
 	export let root: string;
 	export let show_label = true;
@@ -42,7 +41,7 @@
 {#if value !== null}
 	<div class="icon-buttons">
 		{#if show_download_button}
-			<DownloadLink href={value.sources_file.url} download={value.sources_file.orig_name || value.sources_file.path}>
+			<DownloadLink href={value.audio_file.url} download={value.audio_file.orig_name || value.audio_file.path}>
 				<IconButton Icon={Download} label={i18n("common.download")} />
 			</DownloadLink>
 		{/if}

requirements.txt

@@ -1,6 +1,6 @@
 asteroid==0.7.0
 fastapi==0.112.4
-gradio-sourceviewer==1.0.0
+pyannote-viewer==1.0.0
 pyannote-audio==3.3.2
 spaces==0.30.2
 transformers==4.46.3