Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
import torch | |
from transformers import AutoTokenizer, AutoModelForCausalLM, SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan | |
from datasets import load_dataset # To get a speaker embedding for TTS | |
import os | |
import spaces # Import the spaces library for GPU decorator | |
import tempfile # For creating temporary audio files | |
import soundfile as sf # To save audio files | |
# --- Configuration for Language Model (LLM) --- | |
# IMPORTANT: When deploying to Hugging Face Spaces, it's best to use the Hugging Face model ID | |
# rather than a local path ('.'), as the Space will fetch it from the Hub. | |
HUGGINGFACE_MODEL_ID = "HuggingFaceH4/Qwen2.5-1.5B-Instruct-gkd" | |
# You might need to adjust TORCH_DTYPE based on your GPU and model support | |
# torch.float16 (FP16) is common for inference, torch.bfloat16 for newer GPUs | |
# For ZeroGPU/H200, bfloat16 is often preferred if the model supports it and GPU allows. | |
TORCH_DTYPE = torch.bfloat16 # Use bfloat16 for optimal H200 performance | |
# Generation parameters for the LLM (can be adjusted for different response styles) | |
MAX_NEW_TOKENS = 512 | |
DO_SAMPLE = True | |
TEMPERATURE = 0.7 | |
TOP_K = 50 | |
TOP_P = 0.95 | |
# --- Configuration for Text-to-Speech (TTS) --- | |
TTS_MODEL_ID = "microsoft/speecht5_tts" | |
TTS_VOCODER_ID = "microsoft/speecht5_hifigan" | |
# --- Global variables for models and tokenizers --- | |
tokenizer = None | |
llm_model = None # Renamed to avoid conflict with tts_model | |
tts_processor = None | |
tts_model = None | |
tts_vocoder = None | |
speaker_embeddings = None # Global for TTS speaker embedding | |
# --- Load Models and Tokenizers Function --- | |
# Decorate with @spaces.GPU to signal this function needs GPU access | |
def load_models(): | |
""" | |
Loads the language model, tokenizer, TTS models, and speaker embeddings | |
from Hugging Face Hub. This function will be called once when the Gradio app starts up. | |
""" | |
global tokenizer, llm_model, tts_processor, tts_model, tts_vocoder, speaker_embeddings | |
if tokenizer is not None and llm_model is not None and tts_model is not None: | |
print("All models and tokenizers already loaded.") | |
return | |
# When deploying to HF Spaces, you generally don't need an explicit HF_TOKEN | |
# for public models, but it's good practice for private models or if | |
# rate limits are hit. | |
hf_token = os.environ.get("HF_TOKEN") # Access HF_TOKEN from Space secrets if set | |
# Load Language Model (LLM) | |
print(f"Loading LLM tokenizer from: {HUGGINGFACE_MODEL_ID}") | |
try: | |
tokenizer = AutoTokenizer.from_pretrained(HUGGINGFACE_MODEL_ID, token=hf_token) | |
if tokenizer.pad_token is None: | |
tokenizer.pad_token = tokenizer.eos_token | |
print(f"Set tokenizer.pad_token to tokenizer.eos_token ({tokenizer.pad_token_id})") | |
print(f"Loading LLM model from: {HUGGINGFACE_MODEL_ID}...") | |
llm_model = AutoModelForCausalLM.from_pretrained( | |
HUGGINGFACE_MODEL_ID, | |
torch_dtype=TORCH_DTYPE, | |
device_map="auto", # Automatically maps model to GPU if available, else CPU | |
token=hf_token # Pass token if loading private model | |
) | |
llm_model.eval() # Set model to evaluation mode | |
print("LLM model loaded successfully.") | |
except Exception as e: | |
print(f"Error loading LLM model or tokenizer: {e}") | |
print("Please ensure the LLM model ID is correct and you have an internet connection for initial download, or the local path is valid.") | |
tokenizer = None | |
llm_model = None | |
raise RuntimeError("Failed to load LLM model. Check your model ID/path and internet connection.") | |
# Load TTS models | |
print(f"Loading TTS processor, model, and vocoder from: {TTS_MODEL_ID}, {TTS_VOCODER_ID}") | |
try: | |
tts_processor = SpeechT5Processor.from_pretrained(TTS_MODEL_ID, token=hf_token) | |
tts_model = SpeechT5ForTextToSpeech.from_pretrained(TTS_MODEL_ID, token=hf_token) | |
tts_vocoder = SpeechT5HifiGan.from_pretrained(TTS_VOCODER_ID, token=hf_token) | |
# Load a speaker embedding (essential for SpeechT5 TTS) | |
# Using a sample from a public dataset for demonstration | |
print("Loading speaker embeddings for TTS...") | |
embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation", token=hf_token) | |
# Using a specific speaker embedding (you can experiment with different indices) | |
speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0) | |
# Move TTS components to the same device as the LLM model | |
device = llm_model.device if llm_model else 'cpu' | |
tts_model.to(device) | |
tts_vocoder.to(device) | |
speaker_embeddings = speaker_embeddings.to(device) | |
print(f"TTS models and speaker embeddings loaded successfully to device: {device}.") | |
except Exception as e: | |
print(f"Error loading TTS models or speaker embeddings: {e}") | |
print("Please ensure TTS model IDs are correct and you have an internet connection.") | |
tts_processor = None | |
tts_model = None | |
tts_vocoder = None | |
speaker_embeddings = None | |
raise RuntimeError("Failed to load TTS components. Check model IDs and internet connection.") | |
# --- Generate Response and Audio Function --- | |
# Decorate with @spaces.GPU as this function performs GPU-intensive inference | |
def generate_response_and_audio( | |
message: str, # Current user message | |
history: list # Gradio Chatbot history format (list of dictionaries with 'role' and 'content') | |
) -> tuple: # Returns (updated_history, audio_file_path) | |
""" | |
Generates a text response from the loaded LLM and then converts it to audio | |
using the loaded TTS model. | |
""" | |
global tokenizer, llm_model, tts_processor, tts_model, tts_vocoder, speaker_embeddings | |
# Initialize all models if not already loaded | |
if tokenizer is None or llm_model is None or tts_model is None: | |
load_models() | |
if tokenizer is None or llm_model is None: # Check LLM loading status | |
history.append({"role": "user", "content": message}) | |
history.append({"role": "assistant", "content": "Error: Chatbot LLM not loaded. Please check logs."}) | |
return history, None | |
# --- 1. Generate Text Response (LLM) --- | |
# Format messages for the model's chat template | |
messages = history # Use history directly as it's already in the correct format | |
messages.append({"role": "user", "content": message}) # Add current user message | |
# Apply the chat template and tokenize | |
try: | |
input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
except Exception as e: | |
print(f"Error applying chat template: {e}") | |
# Fallback for models without explicit chat templates | |
input_text = "" | |
for item in history: | |
if item["role"] == "user": | |
input_text += f"User: {item['content']}\n" | |
elif item["role"] == "assistant": | |
input_text += f"Assistant: {item['content']}\n" | |
input_text += f"User: {message}\nAssistant:" | |
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(llm_model.device) | |
# Generate response | |
with torch.no_grad(): # Disable gradient calculations for inference | |
output_ids = llm_model.generate( | |
input_ids, | |
max_new_tokens=MAX_NEW_TOKENS, | |
do_sample=DO_SAMPLE, | |
temperature=TEMPERATURE, | |
top_k=TOP_K, | |
top_p=TOP_P, | |
pad_token_id=tokenizer.eos_token_id # Important for generation to stop cleanly | |
) | |
# Decode the generated text, excluding the input prompt part | |
generated_token_ids = output_ids[0][input_ids.shape[-1]:] | |
generated_text = tokenizer.decode(generated_token_ids, skip_special_tokens=True).strip() | |
# --- 2. Generate Audio from Response (TTS) --- | |
audio_path = None | |
if tts_processor and tts_model and tts_vocoder and speaker_embeddings is not None: | |
try: | |
# Ensure TTS components are on the correct device | |
device = llm_model.device if llm_model else 'cpu' | |
tts_model.to(device) | |
tts_vocoder.to(device) | |
speaker_embeddings = speaker_embeddings.to(device) | |
tts_inputs = tts_processor( | |
text=generated_text, | |
return_tensors="pt", | |
max_length=550, # Set a max length to prevent excessively long audio | |
truncation=True # Enable truncation if text exceeds max_length | |
).to(device) | |
with torch.no_grad(): | |
speech = tts_model.generate_speech(tts_inputs["input_ids"], speaker_embeddings, vocoder=tts_vocoder) | |
# Create a temporary file to save the audio | |
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file: | |
audio_path = tmp_file.name | |
# Ensure audio data is on CPU before saving with soundfile | |
sf.write(audio_path, speech.cpu().numpy(), samplerate=16000) | |
print(f"Audio saved to: {audio_path}") | |
except Exception as e: | |
print(f"Error generating audio: {e}") | |
audio_path = None # Return None if audio generation fails | |
else: | |
print("TTS components not loaded. Skipping audio generation.") | |
# --- 3. Update Chat History --- | |
# Append the latest generated response to the history with its role | |
history.append({"role": "assistant", "content": generated_text}) | |
return history, audio_path | |
# --- Gradio Interface --- | |
with gr.Blocks() as demo: | |
gr.Markdown( | |
""" | |
# HuggingFaceH4/Qwen2.5-1.5B-Instruct-gkd chat bot | |
Type your message below and chat with the model! | |
""" | |
) | |
# Set type='messages' for the chatbot to use OpenAI-style dictionaries | |
chatbot = gr.Chatbot(label="Conversation", type='messages') | |
with gr.Row(): | |
text_input = gr.Textbox( | |
label="Your message", | |
placeholder="Type your message here...", | |
scale=4 | |
) | |
submit_button = gr.Button("Send", scale=1) | |
audio_output = gr.Audio( | |
label="Listen to Response", | |
autoplay=True, # Automatically play audio | |
interactive=False # Don't allow user to interact with this audio component | |
) | |
# Link the text input and button to the generation function | |
# Outputs now include both the chatbot history and the audio file path | |
submit_button.click( | |
fn=generate_response_and_audio, | |
inputs=[text_input, chatbot], | |
outputs=[chatbot, audio_output], | |
queue=True # Queue requests for better concurrency | |
) | |
text_input.submit( # Also trigger on Enter key | |
fn=generate_response_and_audio, | |
inputs=[text_input, chatbot], | |
outputs=[chatbot, audio_output], | |
queue=True | |
) | |
# Clear button | |
def clear_chat(): | |
# Clear history, text input, and audio output | |
return [], "", None | |
clear_button = gr.Button("Clear Chat") | |
clear_button.click(clear_chat, inputs=None, outputs=[chatbot, text_input, audio_output]) | |
# Load all models when the app starts up | |
load_models() | |
# Launch the Gradio app | |
demo.queue().launch() | |