Update app.py

app.py

@@ -1,8 +1,7 @@
-import streamlit as st
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from diffusers import DiffusionPipeline
 import torch
-import accelerate
 
 # Load the models and tokenizers
 translation_model_name = "google/madlad400-3b-mt"
@@ -15,7 +14,7 @@ diffusion_model_name = "stabilityai/stable-diffusion-xl-base-1.0"
 diffusion_pipeline = DiffusionPipeline.from_pretrained(diffusion_model_name, torch_dtype=torch.float16)
 diffusion_pipeline = diffusion_pipeline.to("cuda")
 
-# Define the translation and transcription pipeline with accelerate
+# Define the translation and transcription pipeline
 translation_pipeline = pipeline("translation", model=translation_model, tokenizer=translation_tokenizer, device_map="auto")
 transcription_pipeline = pipeline("automatic-speech-recognition", model=transcription_model, device_map="auto")
 
@@ -28,22 +27,28 @@ def transcribe_and_translate_audio_fon(audio_path, num_images=1):
     translation_result = translation_pipeline(transcription_fon, source_lang="fon", target_lang="fr")
     translation_fr = translation_result[0]["translation_text"]
 
+    # Generate images based on the French translation using the diffusion model
     images = diffusion_pipeline(translation_fr, num_images_per_prompt=num_images)["images"]
 
     return images
 
-# Create a Streamlit app
-st.title("Fon Audio to Image Translation")
-
-# Upload audio file
-audio_file = st.file_uploader("Upload an audio file", type=["wav"])
-
-# Transcribe, translate and generate images
-if audio_file:
-    images = transcribe_and_translate_audio_fon(audio_file)
-    st.image(images[0])
-
-
-# Use Accelerate to distribute the computation across available GPUs
-#images = accelerate.launch(transcribe_and_translate_and_generate, audio_file="Fongbe_Speech_Dataset/Fongbe_Speech_Dataset/fongbe_speech_audio_files/wav/64_fongbe_6b36d45b77344caeb1c8d773303c9dcb_for_validation_2022-03-11-23-50-13.wav", num_images=2)
+# Create a Gradio interface
+def process_audio(audio, num_images):
+    images = transcribe_and_translate_audio_fon(audio, num_images)
+    return images
 
+# Define Gradio interface components
+audio_input = gr.Audio(source="upload", type="filepath", label="Upload an audio file")
+image_output = gr.Gallery(label="Generated Images").style(grid=2)
+num_images_input = gr.Slider(minimum=1, maximum=5, step=1, value=1, label="Number of Images")
+
+# Launch Gradio interface
+interface = gr.Interface(
+    fn=process_audio,
+    inputs=[audio_input, num_images_input],
+    outputs=image_output,
+    title="Fon Audio to Image Translation",
+    description="Upload an audio file in Fon, and the app will transcribe, translate to French, and generate related images."
+)
+
+interface.launch()