Update app.py

app.py

@@ -7,22 +7,56 @@ import os
 import torch
 import openai
 from huggingface_hub import InferenceApi
-
-
+from PIL import Image
+import requests
+import io
+import time
 
 # Set up Groq API key
 api_key = os.getenv("GROQ_API_KEY")
 client = Groq(api_key=api_key)
 
+# Hugging Face API details for image generation
+H_key = os.getenv('API_KEY')
+API_URL = "https://api-inference.huggingface.co/models/Artples/LAI-ImageGeneration-vSDXL-2"
+headers = {"Authorization": f"Bearer {H_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)
 
-# Set device: CUDA if available, else CPU
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-# torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        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
+
+        return response.content
+
+    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
 
-model_id1 = os.getenv("API_KEY")
-pipe = StableDiffusionPipeline.from_pretrained(model_id1, low_cpu_mem_usage=True, use_safetensors=True)
-pipe = pipe.to('cpu')
 
 # Updated function for text generation using the new API structure
 def generate_creative_text(prompt):
@@ -40,8 +74,6 @@ def process_audio(audio_path, image_option, creative_text_option):
     if audio_path is None:
         return "Please upload an audio file.", None, None, None
 
-      
-
     # Step 1: Transcribe audio
     try:
         with open(audio_path, "rb") as file:
@@ -67,20 +99,15 @@ def process_audio(audio_path, image_option, creative_text_option):
     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:
-            model_id1 = "dreamlike-art/dreamlike-diffusion-1.0"
-            pipe = StableDiffusionPipeline.from_pretrained(model_id1, torch_dtype=torch.float16, use_safetensors=True)
-            pipe = pipe.to('cpu')
-            image = pipe(translation).images[0]
-        except Exception as e:
-            return tamil_text, translation, creative_text, f"An error occurred during image generation: {str(e)}"
+        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      
-    
+
 
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Base()) as iface:
@@ -104,4 +131,3 @@ with gr.Blocks(theme=gr.themes.Base()) as iface:
 
 # Launch the interface
 iface.launch()
-