from transformers import pipeline
import torch
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
import gradio as gr
import gradio as gr

device = 0 if torch.cuda.is_available() else "cpu"

MODEL_ID = "wav2vec2-base-vinyl_condition"  

pipe = pipeline(
    task="audio-classification",
    model=MODEL_ID,
    chunk_length_s=30,
    device=device,
)

def get_vinyl_condition(filepath):
    output = pipe(
        filepath,
        max_new_tokens=256,
        chunk_length_s=30,
        batch_size=8,
    )
    return output[0]["label"]


demo = gr.Blocks()



demo = gr.Blocks()

file_transcribe = gr.Interface(
    fn=get_vinyl_condition,
    inputs=[
        gr.inputs.Audio(source="upload", optional=True, label="Audio file", type="filepath"),
    ],
    outputs="text",
    layout="horizontal",
    theme="huggingface",
    title="Vinyl Demo: Get Vinyl Condition",
    description=(
        "Get your vinyl ocndition based on the golmine grading starndard! Demo uses the"
        f" checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to get the condition of audio files"
        " of arbitrary length."
    ),
    examples=[
        ["./example.flac"],
    ],
    cache_examples=True,
    allow_flagging="never",
)

with demo:
    gr.TabbedInterface([file_transcribe], ["Transcribe Audio File"])

demo.launch(enable_queue=True)