LiKenun's picture
Switch to use GPU instead of inference client
5c395b2
raw
history blame
3.94 kB
from dotenv import load_dotenv
from os import getenv
import gradio as gr
from automatic_speech_recognition import create_asr_tab
from chatbot import create_chatbot_tab
from image_classification import create_image_classification_tab
from image_to_text import create_image_to_text_tab
from text_to_image import create_text_to_image_tab
from text_to_speech import create_text_to_speech_tab
from translation import create_translation_tab
class App:
"""Main application class for the AI Building Blocks Gradio interface.
This class orchestrates the entire application by creating the Gradio UI
and integrating all the individual building block tabs.
"""
def __init__(
self,
text_to_image_model: str,
image_to_text_model: str,
image_classification_model: str,
text_to_speech_model: str,
audio_transcription_model: str,
chat_model: str,
fallback_translation_model: str
):
"""Initialize the App with model IDs.
Args:
text_to_image_model: Model ID for text-to-image generation.
image_to_text_model: Model ID for image captioning.
image_classification_model: Model ID for image classification.
text_to_speech_model: Model ID for text-to-speech.
audio_transcription_model: Model ID for automatic speech recognition.
chat_model: Model ID for chatbot.
fallback_translation_model: Fallback translation model ID for languages
without specific translation models.
"""
self.text_to_image_model = text_to_image_model
self.image_to_text_model = image_to_text_model
self.image_classification_model = image_classification_model
self.text_to_speech_model = text_to_speech_model
self.audio_transcription_model = audio_transcription_model
self.chat_model = chat_model
self.fallback_translation_model = fallback_translation_model
def run(self):
"""Launch the Gradio application with all building block tabs.
Creates a Gradio Blocks interface with multiple tabs, each representing
a different AI building block. The application will block until the
interface is closed.
"""
with gr.Blocks(title="AI Building Blocks") as demo:
gr.Markdown("# AI Building Blocks")
gr.Markdown("A gallery of building blocks for building AI applications")
with gr.Tabs():
with gr.Tab("Text-to-image Generation"):
create_text_to_image_tab(self.text_to_image_model)
with gr.Tab("Image-to-text or Image Captioning"):
create_image_to_text_tab(self.image_to_text_model)
with gr.Tab("Image Classification"):
create_image_classification_tab(self.image_classification_model)
with gr.Tab("Text-to-speech (TTS)"):
create_text_to_speech_tab(self.text_to_speech_model)
with gr.Tab("Automatic Speech Recognition (ASR)"):
create_asr_tab(self.audio_transcription_model)
with gr.Tab("Chat"):
create_chatbot_tab(self.chat_model)
with gr.Tab("Translation to English"):
create_translation_tab(self.fallback_translation_model)
demo.launch()
if __name__ == "__main__":
load_dotenv()
app = App(
text_to_image_model=getenv("TEXT_TO_IMAGE_MODEL"),
image_to_text_model=getenv("IMAGE_TO_TEXT_MODEL"),
image_classification_model=getenv("IMAGE_CLASSIFICATION_MODEL"),
text_to_speech_model=getenv("TEXT_TO_SPEECH_MODEL"),
audio_transcription_model=getenv("AUDIO_TRANSCRIPTION_MODEL"),
chat_model=getenv("CHAT_MODEL"),
fallback_translation_model=getenv("FALLBACK_TRANSLATION_MODEL")
)
app.run()