Update app.py

app.py

@@ -1,25 +1,13 @@
 from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
-import requests
-import io
-from PIL import Image
-import os
+import torch
+from diffusers import FluxPipeline
 
 # 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 the Hugging Face API key from environment variables for text-to-image model
-hf_api_key = os.getenv("new_hf_token")
-if hf_api_key is None:
-    raise ValueError("Hugging Face API key not found! Please set 'full_token' environment variable.")
-else:
-    headers = {"Authorization": f"Bearer {hf_api_key}"}
-
-# Define the text-to-image model URL (using a faster text-to-image model)
-API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
-
 # Load a smaller text generation model to reduce generation time
 text_generation_model_name = "EleutherAI/gpt-neo-1.3B"
 text_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
@@ -28,45 +16,69 @@ text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name)
 # Create a pipeline for text generation using the selected model
 text_generator = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer)
 
-# Function to generate an image using Hugging Face's text-to-image model
+# Set up the new FluxPipeline for the text-to-image model
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload()  # Enable CPU offloading to save GPU memory if needed
+
+# Function to generate an image using the new FluxPipeline model
 def generate_image_from_text(translated_text):
     try:
-        # Enhanced prompt to focus on details and clarity
-        enhanced_prompt = f"A high-quality image of a person doing yoga with clear facial features and correct body proportions in a tranquil outdoor setting. " \
-                          f"Include detailed mountains, flowing river, and vibrant greenery, captured in soft sunrise light. Ensure the face and body are realistic and proportional."
-        
-        print(f"Generating image from translated text: {enhanced_prompt}")
-        
-        # Sending the enhanced prompt to the text-to-image model
-        response = requests.post(API_URL, headers=headers, json={"inputs": enhanced_prompt})
-        if response.status_code == 200:
-            image_data = response.content
-            image = Image.open(io.BytesIO(image_data))
-            return image
-        else:
-            raise ValueError(f"Error in image generation: {response.text}")
+        print(f"Generating image from translated text: {translated_text}")
+        # Use the FluxPipeline to generate an image from the text
+        image = pipe(translated_text).images[0]
+        print("Image generation completed.")
+        return image, None
+    except Exception as e:
+        print(f"Error during image generation: {e}")
+        return None, f"Error during image generation: {e}"
+
+# Function to generate a shorter paragraph based on the translated text
+def generate_short_paragraph_from_text(translated_text):
+    try:
+        print(f"Generating a short paragraph from translated text: {translated_text}")
+        # Generate a shorter paragraph from the translated text using smaller settings
+        paragraph = text_generator(translated_text, max_length=150, num_return_sequences=1, temperature=0.2, top_p=0.8)[0]['generated_text']
+        print(f"Paragraph generation completed: {paragraph}")
+        return paragraph
+    except Exception as e:
+        print(f"Error during paragraph generation: {e}")
+        return f"Error during paragraph generation: {e}"
 
+# Define the function to translate Tamil text, generate a short paragraph, and create an image
+def translate_generate_paragraph_and_image(tamil_text):
+    # Step 1: Translate Tamil text to English using mbart-large-50
+    try:
+        print("Translating Tamil text to English...")
+        tokenizer.src_lang = "ta_IN"
+        inputs = tokenizer(tamil_text, return_tensors="pt")
+        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]
+        print(f"Translation completed: {translated_text}")
     except Exception as e:
-        print(f"Error: {e}")
-        return None
+        return f"Error during translation: {e}", "", None, None
+
+    # Step 2: Generate a shorter paragraph based on the translated English text
+    paragraph = generate_short_paragraph_from_text(translated_text)
+    if "Error" in paragraph:
+        return translated_text, paragraph, None, None
+
+    # Step 3: Generate an image using the translated English text
+    image, error_message = generate_image_from_text(translated_text)
+    if error_message:
+        return translated_text, paragraph, None, error_message
 
-# Translation Function
-def translate_text(input_text, src_lang="en_XX", tgt_lang="hi_IN"):
-    tokenizer.src_lang = src_lang
-    encoded_input = tokenizer(input_text, return_tensors="pt")
-    generated_tokens = model.generate(encoded_input["input_ids"], forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang])
-    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    return translated_text, paragraph, image, None
 
-# Gradio Interface for image generation
-def translate_and_generate_image(input_text):
-    translated_text = translate_text(input_text)
-    image = generate_image_from_text(translated_text)
-    return image
+# Gradio interface setup
+iface = gr.Interface(
+    fn=translate_generate_paragraph_and_image,
+    inputs=gr.Textbox(lines=2, placeholder="Enter Tamil text here..."),
+    outputs=[gr.Textbox(label="Translated English Text"), 
+             gr.Textbox(label="Generated Short Paragraph"),
+             gr.Image(label="Generated Image")],
+    title="Tamil to English Translation, Short Paragraph Generation, and Image Creation",
+    description="Translate Tamil text to English using Facebook's mbart-large-50 model, generate a short paragraph, and create an image using the translated text.",
+)
 
-# Create a simple Gradio Interface
-iface = gr.Interface(fn=translate_and_generate_image, 
-                     inputs="text", 
-                     outputs="image", 
-                     title="Yoga Image Generator",
-                     description="Enter a description to translate and generate a high-quality yoga image.")
+# Launch the Gradio app
 iface.launch()