Update app.py

app.py

@@ -11,12 +11,12 @@ model_name = "facebook/mbart-large-50-many-to-one-mmt"
 tokenizer = MBart50Tokenizer.from_pretrained(model_name)
 model = MBartForConditionalGeneration.from_pretrained(model_name)
 
-# 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
+# Use a more powerful text generation model, e.g., GPT-J-6B
+text_gen_model = "EleutherAI/gpt-j-6B"  # Or use 'EleutherAI/gpt-neox-20b' for better results
 pipe = pipeline(
-    "text-generation", 
-    model=text_gen_model, 
-    torch_dtype=torch.float32, 
+    "text-generation",
+    model=text_gen_model,
+    torch_dtype=torch.float32,
     device_map="auto"
 )
 
@@ -32,17 +32,9 @@ 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-Neo
+    # Step 2: Generate high-quality English text using GPT-J
     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']
+    generated_text = pipe(prompt, max_length=150, 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):
@@ -63,7 +55,7 @@ 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-Neo, and create an image using the generated text.",
+    description="Translate Tamil text to English using Facebook's mbart-large-50 model, generate high-quality text using GPT-J, and create an image using the generated text.",
 )
 
 # Launch Gradio app