Update app.py

app.py

@@ -1,28 +1,76 @@
-import os
-import io
-import requests
+import whisper
 import gradio as gr
 from groq import Groq
-from transformers import MarianMTModel, MarianTokenizer, AutoModelForCausalLM, AutoTokenizer
 from deep_translator import GoogleTranslator
-from PIL import Image, ImageDraw
-import joblib
-import time
+from diffusers import StableDiffusionPipeline
+import os
 import torch
-import warnings
+import openai
 from huggingface_hub import InferenceApi
-from diffusers import StableDiffusionPipeline
-
-# Load text generation model and tokenizer
-device = "cuda" if torch.cuda.is_available() else "cpu"
-text_generation_model = AutoModelForCausalLM.from_pretrained("gpt2").to(device)  # Move model to the correct device
-text_generation_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+from PIL import Image
+import requests
+import io
+import time
+ 
+# Set up Groq API key
+api_key = os.getenv("g_key")
+client = Groq(api_key=api_key)
+ 
+# Hugging Face API details for image generation
+key = os.getenv("h_key")
+API_URL = "https://api-inference.huggingface.co/models/Artples/LAI-ImageGeneration-vSDXL-2"
+headers = {"Authorization": f"Bearer {key}"}
+ 
+ 
+# Function for querying image generation with retries
+def query_image_generation(payload, max_retries=5):
+    for attempt in range(max_retries):
+        response = requests.post(API_URL, headers=headers, json=payload)
+ 
+        if response.status_code == 503:
+            print(f"Model is still loading, retrying... Attempt {attempt + 1}/{max_retries}")
+            estimated_time = min(response.json().get("estimated_time", 60), 60)
+            time.sleep(estimated_time)
+            continue
+ 
+        if response.status_code != 200:
+            print(f"Error: Received status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return None
  
-# Set the padding token
-text_generation_tokenizer.pad_token = text_generation_tokenizer.eos_token  # Use EOS token as padding token
+        return response.content
  
-# Function to transcribe, translate, analyze sentiment, and generate image
-def process_audio(audio_path, image_option):
+    print(f"Failed to generate image after {max_retries} attempts.")
+    return None
+ 
+# Function for generating an image from text
+def generate_image(prompt):
+    image_bytes = query_image_generation({"inputs": prompt})
+ 
+    if image_bytes is None:
+        return None
+ 
+    try:
+        image = Image.open(io.BytesIO(image_bytes))  # Opening the image from bytes
+        return image
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+ 
+ 
+# Updated function for text generation using the new API structure
+def generate_creative_text(prompt):
+    chat_completion = client.chat.completions.create(
+                messages=[
+                    {"role": "user", "content":prompt}
+                ],
+                model="llama-3.2-90b-text-preview"
+            )
+    chatbot_response = chat_completion.choices[0].message.content
+    return chatbot_response
+ 
+ 
+def process_audio(audio_path, image_option, creative_text_option):
     if audio_path is None:
         return "Please upload an audio file.", None, None, None
  
@@ -38,7 +86,6 @@ def process_audio(audio_path, image_option):
         tamil_text = transcription.text
     except Exception as e:
         return f"An error occurred during transcription: {str(e)}", None, None, None
- 
     # Step 2: Translate Tamil to English
     try:
         translator = GoogleTranslator(source='ta', target='en')
@@ -46,58 +93,39 @@ def process_audio(audio_path, image_option):
     except Exception as e:
         return tamil_text, f"An error occurred during translation: {str(e)}", None, None
  
-    # Step 3: Generate creative text
-    def generate_creative_text(english_text):
-        if not english_text:
-            return "Please provide text to generate creative content."
- 
-        try:
-            inputs = text_generation_tokenizer(english_text, return_tensors="pt", padding=True, truncation=True).to(device)  # Move inputs to the same device
-            generated_tokens = text_generation_model.generate(
-                **inputs,
-                max_length=60,
-                num_return_sequences=1,
-                no_repeat_ngram_size=3,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True,
-                early_stopping=True
-            )
-            creative_text = text_generation_tokenizer.decode(generated_tokens[0], skip_special_tokens=True).strip()
-            return creative_text
-        except Exception as e:
-            return f"An error occurred during text generation: {str(e)}"
- 
-    creative_text = generate_creative_text(translation)
+    # Step 3: Generate creative text (if selected)
+    creative_text = None
+    if creative_text_option == "Generate Creative Text":
+        creative_text = generate_creative_text(translation)
  
     # Step 4: Generate image (if selected)
     image = None
     if image_option == "Generate Image":
-        try:
-            image = pipe(translation).images[0]
-        except Exception as e:
-            return tamil_text, translation, f"An error occurred during image generation: {str(e)}", None
+        image = generate_image(translation)
+        if image is None:
+            return tamil_text, translation, creative_text, f"An error occurred during image generation"
+ 
+    return tamil_text, translation, creative_text, image      
  
-    return tamil_text, translation, image, creative_text
  
 # Create Gradio interface
-with gr.Blocks() as iface:
-    gr.Markdown("# Audio Transcription, Translation, and Image Generation")
+with gr.Blocks(theme=gr.themes.Base()) as iface:
+    gr.Markdown("# Audio Transcription, Translation, Image & Creative Text Generation")
     with gr.Row():
         with gr.Column():
             audio_input = gr.Audio(type="filepath", label="Upload Audio File")
             image_option = gr.Dropdown(["Generate Image", "Skip Image"], label="Image Generation", value="Generate Image")
+            creative_text_option = gr.Dropdown(["Generate Creative Text", "Skip Creative Text"], label="Creative Text Generation", value="Generate Creative Text")
             submit_button = gr.Button("Process Audio")
         with gr.Column():
-            tamil_text_output = gr.Textbox(label="Tamil Transcription", interactive=False)
-            translation_output = gr.Textbox(label="English Translation", interactive=False)
+            tamil_text_output = gr.Textbox(label="Tamil Transcription")
+            translation_output = gr.Textbox(label="English Translation")
+            creative_text_output = gr.Textbox(label="Creative Text")
             image_output = gr.Image(label="Generated Image")
-            creative_text_output = gr.Textbox(label="Creative Text", interactive=False)
- 
     submit_button.click(
         fn=process_audio,
-        inputs=[audio_input, image_option],
-        outputs=[tamil_text_output, translation_output, image_output, creative_text_output]
+        inputs=[audio_input, image_option, creative_text_option],
+        outputs=[tamil_text_output, translation_output, creative_text_output, image_output]
     )
  
 # Launch the interface