Update app.py

app.py

@@ -4,19 +4,19 @@ import requests
 import io
 from PIL import Image
 import os
-import torch  # For text generation models
+import torch
 
 # Load the translation model and tokenizer
 model_name = "facebook/mbart-large-50-many-to-one-mmt"
 tokenizer = MBart50Tokenizer.from_pretrained(model_name)
 model = MBartForConditionalGeneration.from_pretrained(model_name)
 
-# Use GPT-2 for text generation instead of restricted models
-text_gen_model = "gpt2"
+# Use GPT-Neo for more powerful text generation
+text_gen_model = "EleutherAI/gpt-neo-1.3B"  # Use EleutherAI/gpt-neo-2.7B for even better results
 pipe = pipeline(
     "text-generation", 
     model=text_gen_model, 
-    torch_dtype=torch.bfloat16, 
+    torch_dtype=torch.float32, 
     device_map="auto"
 )
 
@@ -32,8 +32,17 @@ def translate_and_generate_image(tamil_text):
     translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
     translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
 
-    # Step 2: Generate descriptive English text using GPT-2
-    generated_text = pipe(translated_text, max_length=50, num_return_sequences=1, truncation=True)[0]['generated_text']
+    # Step 2: Generate descriptive English text using GPT-Neo
+    prompt = f"Create a detailed description based on the following text: {translated_text}"
+    generated_text = pipe(
+        prompt, 
+        max_length=100, 
+        num_return_sequences=1, 
+        temperature=0.7, 
+        top_p=0.9, 
+        top_k=50, 
+        truncation=True
+    )[0]['generated_text']
     
     # Step 3: Use the generated English text to create an image
     def query(payload):
@@ -54,8 +63,8 @@ iface = gr.Interface(
              gr.Textbox(label="Generated Descriptive Text"),
              gr.Image(label="Generated Image")],
     title="Tamil to English Translation, Text Generation, and Image Creation",
-    description="Translate Tamil text to English using Facebook's mbart-large-50 model, generate descriptive text using GPT-2, and create an image using the generated text.",
+    description="Translate Tamil text to English using Facebook's mbart-large-50 model, generate descriptive text using GPT-Neo, and create an image using the generated text.",
 )
 
-# Launch Gradio app without `share=True` (Hugging Face Spaces already handles sharing)
+# Launch Gradio app
 iface.launch()