Update app.py

app.py

@@ -5,12 +5,13 @@ import gradio as gr
 from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
 import os
 
-# Load models and tokenizers globally to avoid reloading them for every request
+# Load the translation model
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 translation_model = MarianMTModel.from_pretrained(model_name)
 translation_tokenizer = MarianTokenizer.from_pretrained(model_name)
 
-gpt_model_name = "EleutherAI/gpt-neo-1.3B"
+# Load GPT-2 model and tokenizer (smaller and faster than GPT-Neo)
+gpt_model_name = "gpt2"
 gpt_tokenizer = AutoTokenizer.from_pretrained(gpt_model_name)
 gpt_model = AutoModelForCausalLM.from_pretrained(gpt_model_name)
 
@@ -20,10 +21,10 @@ def translate_text(tamil_text):
     translation = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translation
 
-def query_gpt_neo(translated_text):
+def query_gpt_2(translated_text):
     prompt = f"Continue the story based on the following text: {translated_text}"
     inputs = gpt_tokenizer(prompt, return_tensors="pt")
-    outputs = gpt_model.generate(inputs['input_ids'], max_length=50, num_return_sequences=1)  # Reduced max_length
+    outputs = gpt_model.generate(inputs['input_ids'], max_length=50, num_return_sequences=1)  # Reduced max_length for speed
     creative_text = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return creative_text
 
@@ -46,8 +47,8 @@ def process_input(tamil_input):
         # Translate the input text
         translated_output = translate_text(tamil_input)
         
-        # Generate creative text using GPT-Neo
-        creative_output = query_gpt_neo(translated_output)
+        # Generate creative text using GPT-2
+        creative_output = query_gpt_2(translated_output)
         
         # Generate an image using Hugging Face's FLUX model
         image_bytes = query_image({"inputs": translated_output})