Update app.py

app.py

@@ -5,70 +5,71 @@ from PIL import Image
 import torch
 import io
 import os
+import base64
 
-# Load Hugging Face API key securely
-HF_API_KEY = os.getenv("HF_API_KEY")   # give token
+# 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. Add it in Hugging Face 'Variables and Secrets'.")
+    raise ValueError("HF_API_KEY is not set. Please add it in Hugging Face 'Variables and secrets'.")
 
-# API Endpoint for Image Generation
+# 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 Tamil-to-English Translation Model
+# 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
+# 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):
-    """Translates Tamil text to English."""
-    inputs = translator_tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
+    """Translate Tamil text to English."""
+    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)
 
 def generate_text(prompt):
-    """Generates a creative text based on English input."""
-    inputs = generator_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
+    """Generate creative text from English input."""
+    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)
 
 def generate_image(prompt):
-    """Sends request to API for image generation."""
+    """Generate image using Hugging Face inference API."""
     response = requests.post(IMAGE_GEN_URL, headers=HEADERS, json={"inputs": prompt})
     if response.status_code == 200:
-        return Image.open(io.BytesIO(response.content))
-    return Image.new("RGB", (300, 300), "red")  # Placeholder image for errors
+        # Check if raw image or base64 encoded
+        content_type = response.headers.get("content-type", "")
+        if "image" in content_type:
+            return Image.open(io.BytesIO(response.content))
+        else:
+            try:
+                result = response.json()
+                if "image" in result:
+                    image_data = base64.b64decode(result["image"])
+                    return Image.open(io.BytesIO(image_data))
+            except Exception as e:
+                print("Error parsing image:", e)
+    return Image.new("RGB", (300, 300), "red")  # fallback placeholder
 
 def process_input(tamil_text):
-    """Complete pipeline: Translation, Text Generation, and Image Generation."""
+    """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 Interface
-
-# 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, and this app will translate it, generate a creative description, and create an image based on the text."
-# )
+# Create Gradio UI
 interface = gr.Interface(
     fn=process_input,
     inputs=gr.Textbox(label="Enter Tamil Text"),
@@ -78,7 +79,8 @@ interface = gr.Interface(
         gr.Image(label="Generated Image")
     ],
     title="Tamil to English Translator & Image Generator",
-    description="Enter Tamil text, and this app will translate it, generate a creative description, and create an image based on the text."
+    description="Enter Tamil text. This app translates it to English, generates a creative description, and produces an image based on the translated text."
 )
-# Launch the Gradio app
-interface.launch()
+
+# Launch app
+interface.launch()