revamped code

app.py

@@ -58,34 +58,42 @@
 # if __name__ == "__main__":
 #     interface.launch(share=True)
 
+
 import gradio as gr
 from translate import translate_text, SUPPORTED_LANGUAGES
+from image_generator import generate_image
 
-def chatbot(input_text, src_lang_name):
-    src_lang_code = SUPPORTED_LANGUAGES[src_lang_name]
-    translated_text = translate_text(input_text, src_lang=src_lang_code)
-    return translated_text
-
-# Gradio Interface
-with gr.Blocks() as app:
-    gr.Markdown("## LinguaVision - Multilingual Translation Chatbot")
+def chatbot(input_text, src_lang):
+    """Process user input, translate it, and generate an image."""
+    # Get the language code from the selected language
+    src_lang_code = SUPPORTED_LANGUAGES[src_lang]
     
-    with gr.Row():
-        language_dropdown = gr.Dropdown(
-            choices=list(SUPPORTED_LANGUAGES.keys()),
-            label="Select Input Language",
-            value="Auto Detect"
-        )
-        input_box = gr.Textbox(label="Enter Text")
+    # Translate input to English
+    translated_text = translate_text(input_text, src_lang_code, tgt_lang_code="en_XX")
     
-    output_box = gr.Textbox(label="Translated Text")
-    translate_button = gr.Button("Translate")
+    # Generate an image based on the translated text
+    image_path = generate_image(f"A scene depicting: {translated_text}")
     
-    translate_button.click(
-        chatbot,
-        inputs=[input_box, language_dropdown],
-        outputs=output_box
-    )
+    return translated_text, image_path
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=chatbot,
+    inputs=[
+        gr.Textbox(label="Enter text"),
+        gr.Dropdown(
+            label="Select Input Language",
+            choices=list(SUPPORTED_LANGUAGES.keys()),
+            value="English",
+        ),
+    ],
+    outputs=[
+        gr.Textbox(label="Translated Text"),
+        gr.Image(label="Generated Image"),
+    ],
+    title="LinguaVision - Multilingual Chatbot with Image Generation",
+    description="Enter text and select the input language to translate it into English and generate an image based on the text.\n NOTE: This tool takes approximately 20 minutes to execute.",
+)
 
-# Run the app
-app.launch()
+if __name__ == "__main__":
+    interface.launch()

translate.py

@@ -13,20 +13,6 @@
 #     print("Translated Text:", translate_text(input_text, src_lang="es", tgt_lang="en"))
 
 
-# from transformers import MarianMTModel, MarianTokenizer
-
-# # Preload the translation model globally
-# model_name = "facebook/mbart-large-50-many-to-many-mmt"
-# tokenizer = MarianTokenizer.from_pretrained(model_name)
-# translation_model = MarianMTModel.from_pretrained(model_name)
-
-# def translate_text(text, src_lang="auto", tgt_lang="en"):
-#     """Translate text from any language to English."""
-#     inputs = tokenizer(text, return_tensors="pt", padding=True)
-#     translated = translation_model.generate(**inputs)
-#     return tokenizer.decode(translated[0], skip_special_tokens=True)
-
-import spaces
 from transformers import MarianMTModel, MarianTokenizer
 
 # Preload the translation model globally
@@ -34,41 +20,8 @@ model_name = "Helsinki-NLP/opus-mt-mul-en"
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 translation_model = MarianMTModel.from_pretrained(model_name)
 
-# Supported languages for the dropdown
-SUPPORTED_LANGUAGES = {
-    "Auto Detect": "auto",
-    "English": "en",
-    "Spanish": "es",
-    "French": "fr",
-    "German": "de",
-    "Italian": "it",
-    "Japanese": "ja",
-    "Chinese": "zh",
-    "Korean": "ko",
-    "Hindi": "hi",
-    "Arabic": "ar",
-    "Russian": "ru",
-    "Portuguese": "pt",
-    "Dutch": "nl",
-    "Bengali": "bn",
-    "Turkish": "tr",
-    "Vietnamese": "vi",
-    "Indonesian": "id",
-    "Malay": "ms",
-}
-@spaces.GPU
 def translate_text(text, src_lang="auto", tgt_lang="en"):
-    """Translate text from the selected source language to English."""
-    # Add language prefix for the source language
-    if src_lang != "auto":
-        text = f">>{src_lang}<< {text}"
-    
-    # Tokenize the input text
+    """Translate text from any language to English."""
     inputs = tokenizer(text, return_tensors="pt", padding=True)
-    
-    # Perform translation
     translated = translation_model.generate(**inputs)
-    
-    # Decode and return the translated text
     return tokenizer.decode(translated[0], skip_special_tokens=True)
-