Update app.py

app.py

@@ -6,115 +6,61 @@ Original file is located at
 """
 
 import os
-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
+from transformers import MarianMTModel, MarianTokenizer
+import gradio as gr
+from PIL import Image, UnidentifiedImageError
 import requests
 import io
-from PIL import Image
-import matplotlib.pyplot as plt
-import gradio as gr
 
-# Load the translation model and tokenizer
+# Load translation models
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 model = MarianMTModel.from_pretrained(model_name)
 
-# Create a translation pipeline
-translator = pipeline("translation", model=model, tokenizer=tokenizer)
+# Define language map
+language_map = {
+    "Tamil": "ta",
+    "Russian": "rus"
+}
+
+def translate_text(input_text, 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
 
-# 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):
-    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
+    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:
             return None
-
-    except Exception as e:
-        print(f"An error occurred: {e}")
-        return None
-
-# Display image
-def show_image(image):
-    if image:
-        plt.imshow(image)
-        plt.axis('off')  # Hide axes
-        plt.show()
     else:
-        print("No image to display")
-
-# 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):
-    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
-
-# 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 None
 
-    return translated_text, creative_text, image
+def process_input(text_input, selected_language):
+    translated_output = translate_text(text_input, selected_language)
+    image = generate_image(translated_output)
+    return translated_output, image
 
-# Create Gradio interface
+# Gradio interface
 interface = gr.Interface(
-    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."
+    fn=process_input,
+    inputs=[gr.Textbox(label="Input Text"), gr.CheckboxGroup(choices=["Tamil", "Russian"], label="Select Language")],
+    outputs=[gr.Textbox(label="Translated Text"), gr.Image(label="Generated Image")],
+    title="Multilingual Translation and Image Generation",
+    description="Translate Tamil or Russian text to English and generate an image."
 )
 
-# Launch Gradio app
-interface.launch()
+interface.launch()