Create app.py

app.py

@@ -0,0 +1,224 @@
+import torch
+import torchaudio
+from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, BitsAndBytesConfig
+import gradio as gr
+import os
+import time
+
+# --- Configuration ---
+model_name = "ibm-granite/granite-speech-3.2-8b"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# --- Load Model and Processor (runs only once on startup) ---
+print("Loading processor...")
+speech_granite_processor = AutoProcessor.from_pretrained(
+    model_name, trust_remote_code=True)
+tokenizer = speech_granite_processor.tokenizer
+print("Processor loaded.")
+
+print("Configuring quantization...")
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",  # TODO: Try fp4 as an alternative.
+    bnb_4bit_use_double_quant=True
+)
+print("Quantization configured.")
+
+print("Loading model...")
+speech_granite = AutoModelForSpeechSeq2Seq.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    device_map="auto",
+    trust_remote_code=True
+)
+speech_granite.eval()
+print("Model loaded.")
+
+# --- Core Transcription Function ---
+def transcribe_audio(audio_input):
+    """
+    Transcribes audio using the loaded Granite model.
+
+    Args:
+        audio_input (tuple or str): Audio data from Gradio.
+            If from microphone: A tuple (sample_rate, numpy_array).
+            If from file upload: A string filepath.
+
+    Returns:
+        str: The transcribed text.
+        float: Processing time in seconds.
+    """
+    start_time = time.time()
+
+    if audio_input is None:
+        return "Error: No audio provided.", 0.0
+
+    print(f"Received audio input type: {type(audio_input)}")
+
+    # --- Load and Preprocess Audio ---
+    try:
+        if isinstance(audio_input, str): # File upload
+            audio_path = audio_input
+            wav, sr = torchaudio.load(audio_path, normalize=True)
+            print(f"Loaded from file: {audio_path}")
+        elif isinstance(audio_input, tuple): # Microphone input
+            sr, wav_np = audio_input
+            wav = torch.from_numpy(wav_np).float().unsqueeze(0) # Convert numpy to tensor [1, N]
+            # Normalize microphone input (assuming it's not normalized)
+            wav = wav / torch.max(torch.abs(wav))
+            print(f"Loaded from microphone input. Sample rate: {sr}, Shape: {wav.shape}")
+        else:
+             return f"Error: Unsupported audio input type: {type(audio_input)}.", 0.0
+
+        print(f"Original sample rate: {sr}, Channels: {wav.shape[0] if wav.dim() > 1 else 1}")
+
+        # Convert to mono if stereo
+        if wav.dim() > 1 and wav.shape[0] > 1:
+            wav = torch.mean(wav, dim=0, keepdim=True)
+            print("Converted stereo to mono")
+
+        # Ensure it's 2D [1, N]
+        if wav.dim() == 1:
+            wav = wav.unsqueeze(0)
+
+        # Resample to 16kHz if necessary
+        if sr != 16000:
+            print(f"Resampling from {sr}Hz to 16000Hz...")
+            resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)
+            wav = resampler(wav)
+            sr = 16000
+            print(f"Resampled to {sr}Hz")
+
+        print(f"Final audio: sample rate {sr}Hz, shape {wav.shape}")
+        assert wav.shape[0] == 1 and sr == 16000, "Audio preprocessing failed"
+
+    except Exception as e:
+        print(f"Error during audio loading/processing: {e}")
+        return f"Error processing audio: {e}", 0.0
+
+    # --- Prepare Prompt ---
+    chat = [
+        {
+            "role": "system",
+            "content": "Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant",
+        },
+        {
+            "role": "user",
+            "content": "<|audio|>can you transcribe the speech into a written format?",
+        }
+    ]
+    text = tokenizer.apply_chat_template(
+        chat, tokenize=False, add_generation_prompt=True
+    )
+
+    # --- Process and Generate ---
+    try:
+        print("Processing inputs...")
+        # Send audio tensor (wav) directly, not the filepath
+        model_inputs = speech_granite_processor(
+            text,
+            audios=wav.squeeze(0).numpy(),
+            sampling_rate=sr,
+            device=device, # Compute embeddings on target device
+            return_tensors="pt",
+        ).to(device) # Move tensors to target device (GPU/CPU)
+        print("Inputs processed.")
+
+
+        print("Generating transcription...")
+        # Generate on the same device as the model
+        model_outputs = speech_granite.generate(
+            **model_inputs,
+            max_new_tokens=1000,
+            num_beams=4,
+            do_sample=False,
+            min_length=1,
+            top_p=1.0,
+            repetition_penalty=3.0,
+            length_penalty=1.0,
+            temperature=1.0,
+            bos_token_id=tokenizer.bos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        print("Generation complete.")
+
+        # --- Decode Output ---
+        num_input_tokens = model_inputs["input_ids"].shape[-1]
+        # Ensure output tensor is on CPU for decoding if necessary
+        new_tokens = model_outputs[0, num_input_tokens:].cpu() # Move to CPU before decoding
+
+        output_text = tokenizer.batch_decode(
+            [new_tokens], # Wrap in a list for batch_decode
+            add_special_tokens=False,
+            skip_special_tokens=True
+        )
+        transcription = output_text[0].strip().upper() # Get first item, strip whitespace, uppercase
+        print(f"Raw output: {output_text[0]}")
+        print(f"Final Transcription: {transcription}")
+
+    except Exception as e:
+        print(f"Error during generation/decoding: {e}")
+        import traceback
+        traceback.print_exc() # Print full traceback for debugging
+        return f"Error during transcription: {e}", 0.0
+
+    end_time = time.time()
+    processing_time = round(end_time - start_time, 2)
+    print(f"Processing time: {processing_time} seconds")
+
+    # Clean up temporary file if it was created by Gradio upload
+    # NOTE: Gradio typically handles cleanup, but belt-and-suspenders approach
+    if isinstance(audio_input, str) and os.path.exists(audio_input):
+         try:
+             # Check if it looks like a temp file before deleting
+             if "gradio" in audio_input or "tmp" in audio_input:
+                 # os.remove(audio_input) # Be cautious enabling this
+                 print(f"Skipping deletion of temp file: {audio_input}")
+                 pass
+         except OSError as e:
+             print(f"Warning: Could not delete temp file {audio_input}: {e}")
+
+
+    return transcription, processing_time
+
+# --- Gradio Interface Definition ---
+# Download example files (replace with actual URLs if needed, or use local paths if packaged)
+# Example using librispeech sample from HF datasets
+try:
+    from datasets import load_dataset
+    ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+    example_audio_path = ds[0]["file"] # Use the path directly if possible
+    example_list = [[example_audio_path]] # Gradio expects list of lists for examples
+except Exception as e:
+    print(f"Could not load example dataset: {e}. Examples will be empty.")
+    example_list = []
+
+
+title = "IBM Granite Speech-to-Text (8B Quantized)"
+description = """
+Transcribe speech audio using the `ibm-granite/granite-speech-3.2-8b` model (4-bit quantized).
+Upload an audio file or use your microphone. The model expects **English** speech.
+Processing might take some time depending on the audio length and hardware (especially on CPU or less powerful GPUs).
+"""
+
+# Define inputs and outputs
+audio_in = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Input Audio") # Use filepath for torchaudio
+text_out = gr.Textbox(label="Transcription", lines=5)
+time_out = gr.Number(label="Processing Time (s)")
+
+# Create and launch the interface
+iface = gr.Interface(
+    fn=transcribe_audio,
+    inputs=audio_in,
+    outputs=[text_out, time_out],
+    title=title,
+    description=description,
+    examples=example_list,
+    cache_examples=False # Disable caching if examples change or have issues
+)
+
+if __name__ == "__main__":
+    iface.launch(debug=True) # Set debug=True for more detailed logs locally