Update app.py

app.py

@@ -1,92 +1,72 @@
-import torch
-from transformers import MBartForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import MBartForConditionalGeneration, MBart50Tokenizer, AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
 import requests
 import io
 from PIL import Image
 import os
 
-# Set up the Hugging Face API key from environment variables
-hf_api_key = os.getenv("new_hf_token")
-if not hf_api_key:
-    raise ValueError("Hugging Face API key not found! Please set the 'HF_API_KEY' environment variable.")
-headers = {"Authorization": f"Bearer {hf_api_key}"}
+# 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)
 
-# Define the text-to-image model URLs
-model_urls = {
-    "stable_diffusion_v1_4": "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4",
-    "stable_diffusion_v1_5": "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5",
-}
-API_URL = model_urls["stable_diffusion_v1_4"]
+# Use the Hugging Face API key from environment variables for text-to-image model
+hf_api_key = os.getenv("full_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 translation model for multilingual text inputs
-translation_model_name = "facebook/mbart-large-50-many-to-one-mmt"
-tokenizer = AutoTokenizer.from_pretrained(translation_model_name)
-translation_model = MBartForConditionalGeneration.from_pretrained(translation_model_name)
+# 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 text generation model from Hugging Face
-text_generation_model_name = "EleutherAI/gpt-neo-2.7B"
+# 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)
-text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name, device_map="auto", torch_dtype=torch.float32)
+text_model = AutoModelForCausalLM.from_pretrained(text_generation_model_name)
 
-# Create a pipeline for text generation
+# 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
 def generate_image_from_text(translated_text):
-    payload = {"inputs": translated_text, "options": {"wait_for_model": True}}
-    response = requests.post(API_URL, headers=headers, json=payload)
-    if response.status_code == 200:
-        image_data = response.content
-        image = Image.open(io.BytesIO(image_data))
-        return image
-    else:
-        # If the model is loading, check the estimated wait time
-        if response.status_code == 503:
-            error_message = response.json()
-            estimated_time = error_message.get("estimated_time", "Unknown")
-            return f"Model is currently loading. Estimated wait time: {estimated_time} seconds. Try again later."
+    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:
-            return f"Failed to generate image. Error: {response.status_code}, Message: {response.text}"
+            raise ValueError(f"Error in image generation: {response.text}")
 
-# Function to translate text using the MBart model
-def translate_text(input_text, src_lang="en"):
-    # Tokenize and translate
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+# 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")
-    translated_tokens = translation_model.generate(**encoded_input)
-    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return translated_text
-
-# Function to generate text using the GPT-Neo model
-def generate_text(prompt, max_length=50):
-    generated_texts = text_generator(prompt, max_length=max_length, num_return_sequences=1)
-    return generated_texts[0]["generated_text"]
-
-# Define the Gradio Interface
-def app_interface(input_text, src_language="en"):
-    translated_text = translate_text(input_text, src_lang=src_language)
-    generated_image = generate_image_from_text(translated_text)
-    generated_text = generate_text(translated_text)
-    return generated_text, generated_image
-
-# Launch the Gradio App using the new Gradio components
-with gr.Blocks() as demo:
-    gr.Markdown("# Multilingual Text-to-Image & Text Generation")
-
-    # Define Gradio components
-    input_text = gr.Textbox(lines=2, placeholder="Enter text here...")
-    src_language = gr.Dropdown(["en", "fr", "de", "es"], value="en", label="Source Language")
-    
-    # Display outputs for text and image generation
-    generated_text_output = gr.Textbox(label="Generated Text")
-    generated_image_output = gr.Image(label="Generated Image")
-
-    # Button to trigger the processing
-    generate_button = gr.Button("Generate")
+    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]
 
-    # Link the button to the function call
-    generate_button.click(fn=app_interface, inputs=[input_text, src_language], outputs=[generated_text_output, generated_image_output])
+# 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
 
-# Run the app
-demo.launch()
+# 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.")
+iface.launch()