changed the whole code

app.py

@@ -1,25 +1,58 @@
+# import gradio as gr
+# from translate import translate_text
+# from image_generator import generate_image
+
+# def chatbot(input_text, src_lang="auto"):
+#     # Translate input to English
+#     translated_text = translate_text(input_text, src_lang, "en")
+#     # Generate image based on translated text
+#     image_path = generate_image(f"A scene depicting: {translated_text}")
+#     return translated_text, image_path
+
+# interface = gr.Interface(
+#     fn=chatbot,
+#     inputs=[
+#         gr.Textbox(label="Input Text"),
+#         gr.Dropdown(choices=["auto", "es", "fr", "de"], label="Source Language")
+#     ],
+#     outputs=[
+#         gr.Textbox(label="Translated Text"),
+#         gr.Image(label="Generated Image")
+#     ],
+#     title="English Language Learning Chatbot"
+# )
+
+# if __name__ == "__main__":
+#     interface.launch()
+
+
 import gradio as gr
 from translate import translate_text
 from image_generator import generate_image
 
 def chatbot(input_text, src_lang="auto"):
+    """Process user input, translate it, and generate an image."""
     # Translate input to English
     translated_text = translate_text(input_text, src_lang, "en")
-    # Generate image based on translated text
+    
+    # Generate an image based on the translated text
     image_path = generate_image(f"A scene depicting: {translated_text}")
+    
     return translated_text, image_path
 
+# Gradio Interface
 interface = gr.Interface(
     fn=chatbot,
     inputs=[
-        gr.Textbox(label="Input Text"),
-        gr.Dropdown(choices=["auto", "es", "fr", "de"], label="Source Language")
+        gr.Textbox(label="Enter text in any language"),
+        gr.Textbox(label="Source Language (optional, e.g., 'es' for Spanish)", value="auto"),
     ],
     outputs=[
         gr.Textbox(label="Translated Text"),
-        gr.Image(label="Generated Image")
+        gr.Image(label="Generated Image"),
     ],
-    title="English Language Learning Chatbot"
+    title="Multilingual Chatbot with Image Generation",
+    description="Enter text in any language to translate it into English and generate an image based on the text.",
 )
 
 if __name__ == "__main__":

image_generator.py

@@ -12,20 +12,39 @@
 #     print("Generated Image Path:", generate_image(prompt))
 
 
+# from diffusers import StableDiffusionPipeline
+# import torch
+
+# def generate_image(prompt):
+#     model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+    
+#     # Use GPU if available, otherwise fallback to CPU
+#     device = "cuda" if torch.cuda.is_available() else "cpu"
+#     model.to(device)
+    
+#     image = model(prompt).images[0]
+#     image.save("output.png")
+#     return "output.png"
+
+# if __name__ == "__main__":
+#     prompt = "A friendly person saying 'How are you?'"
+#     print("Generated Image Path:", generate_image(prompt))
+
+
 from diffusers import StableDiffusionPipeline
 import torch
 
+# Preload the model globally
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1-base",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+model.to(device)
+
 def generate_image(prompt):
-    model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    
-    # Use GPU if available, otherwise fallback to CPU
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
-    
+    """Generate an image from a text prompt."""
     image = model(prompt).images[0]
-    image.save("output.png")
-    return "output.png"
-
-if __name__ == "__main__":
-    prompt = "A friendly person saying 'How are you?'"
-    print("Generated Image Path:", generate_image(prompt))
+    output_path = "output.png"
+    image.save(output_path)
+    return output_path

translate.py

@@ -1,13 +1,27 @@
+# from transformers import MarianMTModel, MarianTokenizer
+
+# def translate_text(text, src_lang="es", tgt_lang="en"):
+#     model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
+#     tokenizer = MarianTokenizer.from_pretrained(model_name)
+#     model = MarianMTModel.from_pretrained(model_name)
+#     inputs = tokenizer(text, return_tensors="pt", padding=True)
+#     translated = model.generate(**inputs)
+#     return tokenizer.decode(translated[0], skip_special_tokens=True)
+
+# if __name__ == "__main__":
+#     input_text = "¿Cómo estás?"
+#     print("Translated Text:", translate_text(input_text, src_lang="es", tgt_lang="en"))
+
+
 from transformers import MarianMTModel, MarianTokenizer
 
-def translate_text(text, src_lang="es", tgt_lang="en"):
-    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
+# Preload the translation model globally
+model_name = "Helsinki-NLP/opus-mt-mul-en"
+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 = model.generate(**inputs)
+    translated = translation_model.generate(**inputs)
     return tokenizer.decode(translated[0], skip_special_tokens=True)
-
-if __name__ == "__main__":
-    input_text = "¿Cómo estás?"
-    print("Translated Text:", translate_text(input_text, src_lang="es", tgt_lang="en"))