Update app.py

app.py

@@ -1,91 +1,140 @@
-import gradio as gr
+# Install the required libraries
+pip install transformers gradio Pillow requests
+import os
 import requests
 from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
-from PIL import Image
-import torch
+from PIL import Image, ImageDraw
 import io
-import os
-import base64
-
-# Load Hugging Face API key from environment
-HF_API_KEY = os.getenv("HF_API_KEY")  # Add this in 'Variables and secrets' on HF Spaces
-if not HF_API_KEY:
-    raise ValueError("HF_API_KEY is not set. Please add it in Hugging Face 'Variables and secrets'.")
-
-# API endpoint for image generation
-IMAGE_GEN_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
-HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
-
-# Check if GPU is available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load translation model: Tamil to English
-translator_model = "Helsinki-NLP/opus-mt-mul-en"
-translator = MarianMTModel.from_pretrained(translator_model).to(device)
-translator_tokenizer = MarianTokenizer.from_pretrained(translator_model)
-
-# Load text generation model
-generator_model = "EleutherAI/gpt-neo-1.3B"
-generator = AutoModelForCausalLM.from_pretrained(generator_model).to(device)
-generator_tokenizer = AutoTokenizer.from_pretrained(generator_model)
-if generator_tokenizer.pad_token is None:
-    generator_tokenizer.pad_token = generator_tokenizer.eos_token
-generator.config.pad_token_id = generator_tokenizer.pad_token_id
-
-def translate_tamil_to_english(text):
-    """Translate Tamil text to English."""
-    try:
-        inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
-        output = translator.generate(**inputs)
-        return translator_tokenizer.decode(output[0], skip_special_tokens=True)
-    except Exception as e:
-        return f"Translation error: {str(e)}"
+import gradio as gr
+import torch
 
-def generate_text(prompt):
-    """Generate creative text from English input."""
-    try:
-        inputs = generator_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
-        output = generator.generate(**inputs, max_length=100)
-        return generator_tokenizer.decode(output[0], skip_special_tokens=True)
-    except Exception as e:
-        return f"Text generation error: {str(e)}"
+# Detect if GPU is available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Load the MarianMT model and tokenizer for translation (Tamil to English)
+model_name = "Helsinki-NLP/opus-mt-mul-en"
+translation_model = MarianMTModel.from_pretrained(model_name).to(device)
+translation_tokenizer = MarianTokenizer.from_pretrained(model_name)
+
+# Load GPT-Neo for creative text generation
+text_generation_model_name = "EleutherAI/gpt-neo-1.3B"
+text_generation_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name).to(device)
+text_generation_tokenizer = AutoTokenizer.from_pretrained(text_generation_model_name)
+
+# Add padding token to GPT-Neo tokenizer if not present
+if text_generation_tokenizer.pad_token is None:
+    text_generation_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+
+# Set your Hugging Face API key
+os.environ['HF_API_KEY'] = 'Your_HF_TOKEN'  # Replace with your actual API key
+api_key = os.getenv('HF_API_KEY')
+if api_key is None:
+    raise ValueError("Hugging Face API key is not set. Please set it in your environment.")
+
+headers = {"Authorization": f"Bearer {api_key}"}
 
+# Define the API URL for image generation (replace with actual model URL)
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"  # Replace with a valid image generation model
+
+# Query Hugging Face API to generate image with error handling
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    if response.status_code != 200:
+        print(f"Error: Received status code {response.status_code}")
+        print(f"Response: {response.text}")
+        return None
+    return response.content
+
+# Translate Tamil text to English
+def translate_text(tamil_text):
+    inputs = translation_tokenizer(tamil_text, return_tensors="pt", padding=True, truncation=True).to(device)
+    translated_tokens = translation_model.generate(**inputs)
+    translation = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    return translation
+
+# Generate an image based on the translated text with error handling
 def generate_image(prompt):
-    """Generate image using Hugging Face inference API."""
+    image_bytes = query({"inputs": prompt})
+
+    if image_bytes is None:
+        # Return a blank image with error message
+        error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
+        d = ImageDraw.Draw(error_img)
+        d.text((10, 150), "Image Generation Failed", fill=(255, 255, 255))
+        return error_img
+
     try:
-        response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
-        if response.status_code == 200:
-            content_type = response.headers.get("content-type", "")
-            if "image" in content_type:
-                return Image.open(io.BytesIO(response.content))
-            else:
-                result = response.json()
-                if "image" in result:
-                    image_data = base64.b64decode(result["image"])
-                    return Image.open(io.BytesIO(image_data))
+        image = Image.open(io.BytesIO(image_bytes))
+        return image
     except Exception as e:
-        print("Error generating image:", e)
-    return Image.new("RGB", (300, 300), "red")  # fallback placeholder
-
-def process_input(tamil_text):
-    """Pipeline: Translate → Generate Text → Generate Image"""
-    english_text = translate_tamil_to_english(tamil_text)
-    creative_text = generate_text(english_text)
-    image = generate_image(english_text)
-    return english_text, creative_text, image
-
-# Create Gradio UI
-interface = gr.Interface(
-    fn=process_input,
-    inputs=gr.Textbox(label="Enter Tamil Text"),
-    outputs=[
-        gr.Textbox(label="Translated English Text"),
-        gr.Textbox(label="Creative Text"),
-        gr.Image(label="Generated Image")
-    ],
-    title="Tamil to English Translator & Image Generator",
-    description="Enter Tamil text. This app translates it to English, generates a creative description, and produces an image based on the translated text."
-)
-
-# Launch app
-interface.launch()
+        print(f"Error: {e}")
+        # Return an error image in case of failure
+        error_img = Image.new('RGB', (300, 300), color=(255, 0, 0))
+        d = ImageDraw.Draw(error_img)
+        d.text((10, 150), "Invalid Image Data", fill=(255, 255, 255))
+        return error_img
+
+# Generate creative text based on the translated English text
+def generate_creative_text(translated_text):
+    inputs = text_generation_tokenizer(translated_text, return_tensors="pt", padding=True, truncation=True).to(device)
+    generated_tokens = text_generation_model.generate(**inputs, max_length=100)
+    creative_text = text_generation_tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+    return creative_text
+
+# Function to handle the full workflow
+def translate_generate_image_and_text(tamil_text):
+    # Step 1: Translate Tamil to English
+    translated_text = translate_text(tamil_text)
+
+    # Step 2: Generate an image from the translated text
+    image = generate_image(translated_text)
+
+    # Step 3: Generate creative text from the translated text
+    creative_text = generate_creative_text(translated_text)
+
+    return translated_text, creative_text, image
+
+# Create a visually appealing Gradio interface
+css = """
+#transart-title {
+    font-size: 2.5em;
+    font-weight: bold;
+    color: #4CAF50;
+    text-align: center;
+    margin-bottom: 10px;
+}
+#transart-subtitle {
+    font-size: 1.25em;
+    text-align: center;
+    color: #555555;
+    margin-bottom: 20px;
+}
+body {
+    background-color: #f0f0f5;
+}
+.gradio-container {
+    font-family: 'Arial', sans-serif;
+}
+"""
+
+# Custom HTML for title and subtitle (can be displayed in Markdown)
+title_markdown = """
+# <div id="transart-title">TransArt</div>
+### <div id="transart-subtitle">Tamil to English Translation, Creative Text & Image Generation</div>
+"""
+
+# Gradio interface with customized layout and aesthetics
+with gr.Blocks(css=css) as interface:
+    gr.Markdown(title_markdown)  # Title and subtitle in Markdown
+    with gr.Row():
+        with gr.Column():
+            tamil_input = gr.Textbox(label="Enter Tamil Text", placeholder="Type Tamil text here...", lines=3)  # Input for Tamil text
+        with gr.Column():
+            translated_output = gr.Textbox(label="Translated Text", interactive=False)        # Output for translated text
+            creative_text_output = gr.Textbox(label="Creative Generated Text", interactive=False)  # Output for creative text
+            generated_image_output = gr.Image(label="Generated Image")  # Output for generated image
+
+    gr.Button("Generate").click(fn=translate_generate_image_and_text, inputs=tamil_input, outputs=[translated_output, creative_text_output, generated_image_output])
+
+# Launch the Gradio app
+interface.launch(debug=True, server_name="0.0.0.0")