init Spirit LM

.gitignore

@@ -0,0 +1 @@
+venv

app.py

@@ -1,7 +1,88 @@
 import gradio as gr
+from spiritlm.model.spiritlm_model import Spiritlm, OutputModality, GenerationInput, ContentType
+from transformers import GenerationConfig
+import torchaudio
+import torch
+import tempfile
+import os
+import numpy as np
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Initialize the Spirit LM model with the modified class
+spirit_lm = Spiritlm("spirit-lm-base-7b")
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def generate_output(input_type, input_content_text, input_content_audio, output_modality, temperature, top_p, max_new_tokens, do_sample, speaker_id):
+    generation_config = GenerationConfig(
+        temperature=temperature,
+        top_p=top_p,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+    )
+
+    if input_type == "text":
+        interleaved_inputs = [GenerationInput(content=input_content_text, content_type=ContentType.TEXT)]
+    elif input_type == "audio":
+        # Load audio file
+        waveform, sample_rate = torchaudio.load(input_content_audio)
+        interleaved_inputs = [GenerationInput(content=waveform.squeeze(0), content_type=ContentType.SPEECH)]
+    else:
+        raise ValueError("Invalid input type")
+
+    outputs = spirit_lm.generate(
+        interleaved_inputs=interleaved_inputs,
+        output_modality=OutputModality[output_modality.upper()],
+        generation_config=generation_config,
+        speaker_id=speaker_id,  # Pass the selected speaker ID
+    )
+
+    text_output = ""
+    audio_output = None
+
+    for output in outputs:
+        if output.content_type == ContentType.TEXT:
+            text_output = output.content
+        elif output.content_type == ContentType.SPEECH:
+            # Ensure output.content is a NumPy array
+            if isinstance(output.content, np.ndarray):
+                # Debugging: Print shape and dtype of the audio data
+                print("Audio data shape:", output.content.shape)
+                print("Audio data dtype:", output.content.dtype)
+
+                # Ensure the audio data is in the correct format
+                if len(output.content.shape) == 1:
+                    # Mono audio data
+                    audio_data = torch.from_numpy(output.content).unsqueeze(0)
+                else:
+                    # Stereo audio data
+                    audio_data = torch.from_numpy(output.content)
+
+                # Save the audio content to a temporary file
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio_file:
+                    torchaudio.save(temp_audio_file.name, audio_data, 16000)
+                    audio_output = temp_audio_file.name
+            else:
+                raise TypeError("Expected output.content to be a NumPy array, but got {}".format(type(output.content)))
+
+    return text_output, audio_output
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=generate_output,
+    inputs=[
+        gr.Radio(["text", "audio"], label="Input Type", value="text"),
+        gr.Textbox(label="Input Content (Text)"),
+        gr.Audio(label="Input Content (Audio)", type="filepath"),
+        gr.Radio(["TEXT", "SPEECH", "ARBITRARY"], label="Output Modality", value="SPEECH"),
+        gr.Slider(0, 1, step=0.1, value=0.9, label="Temperature"),
+        gr.Slider(0, 1, step=0.05, value=0.95, label="Top P"),
+        gr.Slider(1, 800, step=1, value=500, label="Max New Tokens"),
+        gr.Checkbox(value=True, label="Do Sample"),
+        gr.Dropdown(choices=[0, 1, 2, 3], value=0, label="Speaker ID"), 
+    ],
+    outputs=[gr.Textbox(label="Generated Text"), gr.Audio(label="Generated Audio")],
+    title="Spirit LM WebUI Demo",
+    description="Demo for generating text or audio using the Spirit LM model.",
+    flagging_mode="never",
+)
+
+# Launch the interface
+iface.launch()

requirements.txt

@@ -0,0 +1,101 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.6.2.post1
+audioread==3.0.1
+certifi==2024.8.30
+cffi==1.17.1
+charset-normalizer==3.4.0
+click==8.1.7
+decorator==5.1.1
+einops==0.8.0
+encodec==0.1.1
+exceptiongroup==1.2.2
+fairscale==0.4.13
+fastapi==0.115.4
+ffmpy==0.4.0
+filelock==3.16.1
+fsspec==2024.10.0
+gradio==5.4.0
+gradio_client==1.4.2
+h11==0.14.0
+httpcore==1.0.6
+httpx==0.27.2
+huggingface-hub==0.26.2
+idna==3.10
+Jinja2==3.1.4
+joblib==1.4.2
+lazy_loader==0.4
+librosa==0.10.2.post1
+llvmlite==0.43.0
+local-attention==1.9.15
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+mpmath==1.3.0
+msgpack==1.1.0
+networkx==3.4.2
+numba==0.60.0
+numpy==2.0.2
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
+omegaconf==2.3.0
+orjson==3.10.10
+packaging==24.1
+pandas==2.2.3
+pillow==11.0.0
+platformdirs==4.3.6
+pooch==1.8.2
+pyarrow==18.0.0
+pycparser==2.22
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydub==0.25.1
+Pygments==2.18.0
+python-dateutil==2.9.0.post0
+python-multipart==0.0.12
+pytz==2024.2
+PyYAML==6.0.2
+regex==2024.9.11
+requests==2.32.3
+rich==13.9.3
+ruff==0.7.1
+safehttpx==0.1.1
+safetensors==0.4.5
+scikit-learn==1.5.2
+scipy==1.14.1
+semantic-version==2.10.0
+sentencepiece==0.2.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+soundfile==0.12.1
+soxr==0.5.0.post1
+spiritlm==0.1.0
+starlette==0.41.2
+sympy==1.13.1
+threadpoolctl==3.5.0
+tokenizers==0.20.1
+tomlkit==0.12.0
+torch==2.5.0
+torchaudio==2.5.0
+torchfcpe==0.0.4
+tqdm==4.66.6
+transformers==4.46.0
+triton==3.1.0
+typer==0.12.5
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+uvicorn==0.32.0
+websockets==12.0