Update app.py

app.py

@@ -1,82 +1,102 @@
 import os
-from transformers import MarianMTModel, MarianTokenizer, GPTNeoForCausalLM, AutoTokenizer
-import gradio as gr
-from PIL import Image, UnidentifiedImageError
+from huggingface_hub import login
+# Retrieve the actual token from the environment variable
+hf_token = os.getenv("HF_TOKEN")
+
+# Check if the token is retrieved properly
+if hf_token:
+    # Use the retrieved token
+    login(token=hf_token, add_to_git_credential=True)
+else:
+    raise ValueError("Hugging Face token not found in environment variables.")
+
+# Import necessary libraries
+from transformers import MarianMTModel, MarianTokenizer, pipeline
 import requests
 import io
+from PIL import Image
+import matplotlib.pyplot as plt
+import gradio as gr
 
-# Load translation model
+# Load the translation model and tokenizer
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 model = MarianMTModel.from_pretrained(model_name)
 
-# Load GPT-Neo model for creative content generation
-gpt_neo_model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
-gpt_neo_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
-
-# Define language map
-language_map = {
-    "Tamil": "ta",
-    "Russian": "rus"
-}
-
-def translate_text(input_text, selected_languages):
-    if not selected_languages:
-        return "Please select at least one language."
-    
-    selected_language = selected_languages[0]  # Pick the first selected language
-    lang_code = language_map[selected_language]
-    lang_prefix = f">>{lang_code}<< "
-    text_with_lang = lang_prefix + input_text
-    inputs = tokenizer(text_with_lang, return_tensors="pt", padding=True)
-    translated_tokens = model.generate(**inputs)
-    translation = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return translation
+# Create a translation pipeline
+translator = pipeline("translation", model=model, tokenizer=tokenizer)
+
+# Function for translation
+def translate_text(tamil_text):
+    try:
+        translation = translator(tamil_text, max_length=40)
+        translated_text = translation[0]['translation_text']
+        return translated_text
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+# API credentials and endpoint
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+headers = {"Authorization": f"Bearer {hf_token}"}
 
+# Function to send payload and generate image
 def generate_image(prompt):
-    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-    hf_token = os.getenv("HF_TOKEN")
-    headers = {"Authorization": f"Bearer {hf_token}"}
-    
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-
-    if response.status_code == 200:
-        image_bytes = response.content
-        try:
-            image = Image.open(io.BytesIO(image_bytes))
-            return image
-        except UnidentifiedImageError:
+    try:
+        response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
+
+        # Check if the response is successful
+        if response.status_code == 200:
+            print("API call successful, generating image...")
+            image_bytes = response.content
+            # Try opening the image
+            try:
+                image = Image.open(io.BytesIO(image_bytes))
+                return image
+            except Exception as e:
+                print(f"Error opening image: {e}")
+                return None
+        else:
+            print(f"Failed to get image: Status code {response.status_code}")
+            print("Response content:", response.text)  # Print response for debugging
             return None
-    else:
+    except Exception as e:
+        print(f"An error occurred: {e}")
         return None
 
+# Load GPT-Neo model for creative text generation
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+gpt_neo_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
+gpt_neo_model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
+
+# Function to generate creative text based on translated text
 def generate_creative_text(translated_text):
-    prompt = f"Create a creative story based on the following sentence: {translated_text}"
-    inputs = gpt_neo_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=100)
-    output = gpt_neo_model.generate(inputs["input_ids"], max_length=150, do_sample=True, temperature=0.7)
-    creative_text = gpt_neo_tokenizer.decode(output[0], skip_special_tokens=True)
+    input_ids = gpt_neo_tokenizer(translated_text, return_tensors='pt').input_ids
+    generated_text_ids = gpt_neo_model.generate(input_ids, max_length=100)
+    creative_text = gpt_neo_tokenizer.decode(generated_text_ids[0], skip_special_tokens=True)
     return creative_text
 
-def process_input(text_input, selected_languages):
-    translated_output = translate_text(text_input, selected_languages)
-    creative_text = generate_creative_text(translated_output)
-    image = generate_image(translated_output)
-    return translated_output, creative_text, image
+# Function to handle the full workflow
+def translate_generate_image_and_text(tamil_text):
+    # Step 1: Translate Tamil text to English
+    translated_text = translate_text(tamil_text)
+
+    # Step 2: Generate an image based on the translated text
+    image = generate_image(translated_text)
+
+    # Step 3: Generate creative text based on the translated text
+    creative_text = generate_creative_text(translated_text)
+
+    return translated_text, creative_text, image
 
-# Gradio interface
+# Create Gradio interface
 interface = gr.Interface(
-    fn=process_input,
-    inputs=[
-        gr.Textbox(label="Input Text"),
-        gr.CheckboxGroup(choices=["Tamil", "Russian"], label="Select Language")
-    ],
-    outputs=[
-        gr.Textbox(label="Translated Text"),
-        gr.Textbox(label="Creative Text"),
-        gr.Image(label="Generated Image")
-    ],
-    title="Multilingual Translation and Image Generation",
-    description="Translate Tamil or Russian text to English, generate creative content, and create an image."
+    fn=translate_generate_image_and_text,
+    inputs="text",
+    outputs=["text", "text", "image"],
+    title="Tamil to English Translation, Image Generation & Creative Text",
+    description="Enter Tamil text to translate to English, generate an image, and create creative text based on the translation."
 )
 
+# Launch Gradio app
 interface.launch()