Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import json
 import urllib.request
@@ -11,6 +10,7 @@ from hercai import Hercai
 import uuid
 import time
 import gradio as gr
+import requests
 
 # Configure detailed logging
 log_dir = os.getenv('LOG_DIRECTORY', './')  # Get log directory from environment variable, default to current directory
@@ -42,7 +42,7 @@ LOGGER.setLevel(log_level)  # Set the log level for the logger instance
 
 class Text2Video:
     """
-    A class to generate videos from text prompts, with detailed logging and a user-friendly interface.
+    A class to generate videos from text prompts, with detailed logging, model selection, and a user-friendly interface.
     """
 
     def __init__(self) -> None:
@@ -53,12 +53,14 @@ class Text2Video:
         self.herc = Hercai("")  # Replace "" with your actual Hercai API key if you have one
         LOGGER.info("Hercai initialized successfully")
 
-    def get_image(self, img_prompt: str) -> str:
+    def get_image(self, img_prompt: str, image_generator: str, image_model: str) -> str:
         """
-        Generate an image from a text prompt using Hercai, with detailed logging and comic book styling.
+        Generate an image from a text prompt using the selected AI model, with detailed logging and comic book styling.
 
         Args:
             img_prompt (str): The text prompt to generate the image from.
+            image_generator (str): The name of the AI image generation service (Hercai, Prodia, or Pollinations).
+            image_model (str): The specific model to use within the selected AI image generation service.
 
         Returns:
             str: The URL of the generated image. Returns an empty string if an error occurred.
@@ -69,12 +71,46 @@ class Text2Video:
             modified_prompt = f"Generate a comic book style image with speech bubbles containing the following text: '{img_prompt}'. " \
                               f"Include elements like vibrant colors, onomatopoeia, and exaggerated expressions to enhance the comic book aesthetic."
             # Log the modified prompt
-            LOGGER.info(f"Modified prompt for Hercai: {modified_prompt}")
-
-            # Generate the image using Hercai
-            image_result = self.herc.draw_image(model="v3", prompt=modified_prompt, negative_prompt="Dark and gloomy")
-            # Extract the image URL from the result
-            image_url = image_result["url"]
+            LOGGER.info(f"Modified prompt for {image_generator}: {modified_prompt}")
+            image_url = ""
+
+            if image_generator == "Hercai":
+                # Log the selected Hercai model
+                LOGGER.info(f"Using Hercai model: {image_model}")
+
+                # Generate the image using Hercai
+                image_result = self.herc.draw_image(model=image_model, prompt=modified_prompt, negative_prompt="Dark and gloomy")
+                # Extract the image URL from the result
+                image_url = image_result["url"]
+
+            elif image_generator == "Prodia":
+                # Log the selected Prodia model
+                LOGGER.info(f"Using Prodia model: {image_model}")
+                # Create the Prodia API call
+                api_url = "https://api.prodia.com/v1/generate"
+                payload = {
+                    "model": image_model,
+                    "prompt": modified_prompt,
+                    "negative_prompt": "Dark and gloomy"
+                }
+                headers = {
+                    "Authorization": "Bearer YOUR_PRODIA_API_KEY"  # Replace YOUR_PRODIA_API_KEY with your actual Prodia API key
+                }
+                response = requests.post(api_url, json=payload, headers=headers)
+                if response.status_code == 200:
+                    image_url = response.json()["url"]
+                    # Log the generated image URL
+                    LOGGER.info(f"Image generated successfully using Prodia: {image_url}")
+                else:
+                    # Log an error if the Prodia API call failed
+                    LOGGER.error(f"Error generating image using Prodia: {response.text}")
+
+            elif image_generator == "Pollinations":
+                # Log the selected Pollinations model
+                LOGGER.info(f"Using Pollinations model: {image_model}")
+                # Implement Pollinations API call here, similar to Prodia
+                # Replace the following placeholder with your Pollinations API call
+                # ...
 
             # Log the generated image URL
             LOGGER.info(f"Image generated successfully: {image_url}")
@@ -82,7 +118,7 @@ class Text2Video:
 
         except Exception as e:
             # Log any errors encountered during image generation
-            LOGGER.error(f"Error generating image for prompt '{img_prompt}': {e}")
+            LOGGER.error(f"Error generating image for prompt '{img_prompt}' using {image_generator}: {e}")
             return ""
 
     def download_img_from_url(self, image_url: str, image_path: str) -> str:
@@ -138,12 +174,14 @@ class Text2Video:
             LOGGER.error(f"Error converting text '{img_prompt}' to audio: {e}")
             return ""
 
-    def get_images_and_audio(self, list_prompts: list) -> tuple:
+    def get_images_and_audio(self, list_prompts: list, image_generator: str, image_model: str) -> tuple:
         """
-        Generate images and corresponding audio files for a list of text prompts.
+        Generate images and corresponding audio files for a list of text prompts using the selected AI model.
 
         Args:
             list_prompts (list): A list of text prompts.
+            image_generator (str): The name of the AI image generation service (Hercai, Prodia, or Pollinations).
+            image_model (str): The specific model to use within the selected AI image generation service.
 
         Returns:
             tuple: A tuple containing two lists: image paths and audio paths.
@@ -160,8 +198,8 @@ class Text2Video:
                 # Construct the image path using the unique identifier
                 image_path = f"{img_prompt[:9]}_{unique_id}.png"
 
-                # Generate the image URL using Hercai
-                img_url = self.get_image(img_prompt)
+                # Generate the image URL using the selected AI model
+                img_url = self.get_image(img_prompt, image_generator, image_model)
 
                 # Download the image from the generated URL
                 image = self.download_img_from_url(img_url, image_path)
@@ -237,13 +275,14 @@ class Text2Video:
             # Log any errors encountered during video creation
             LOGGER.error(f"Error creating video: {e}")
 
-
-    def generate_video(self, text: str) -> str:
+    def generate_video(self, text: str, image_generator: str, image_model: str) -> str:
         """
-        Generate a video from a comma-separated string of text prompts
+        Generate a video from a comma-separated string of text prompts using the selected AI model.
 
         Args:
             text (str): A comma-separated string of text prompts, where each prompt represents a scene or frame in the video.
+            image_generator (str): The name of the AI image generation service (Hercai, Prodia, or Pollinations).
+            image_model (str): The specific model to use within the selected AI image generation service.
 
         Returns:
             str: The file path of the generated video file. Returns an empty string if an error occurred.
@@ -257,8 +296,8 @@ class Text2Video:
             # Define the output path for the generated video
             output_path = "output_video.mp4"
 
-            # Generate images and corresponding audio files for each prompt
-            img_list, audio_paths = self.get_images_and_audio(list_prompts)
+            # Generate images and corresponding audio files for each prompt using the selected AI model
+            img_list, audio_paths = self.get_images_and_audio(list_prompts, image_generator, image_model)
 
             # Create the video from the generated images and audio files
             self.create_video_from_images_and_audio(img_list, audio_paths, output_path)
@@ -271,7 +310,6 @@ class Text2Video:
             LOGGER.error(f"Error generating video from text '{text}': {e}")
             return ""
 
-
     def gradio_interface(self):
         """
         Creates a user-friendly Gradio interface for the video generation application.
@@ -287,6 +325,20 @@ class Text2Video:
                 input_text = gr.Textbox(label="Comics Text",
                                        placeholder="Enter the comics text, separating scenes with double commas (,,)")
 
+            # Create a dropdown menu for selecting the AI image generation service
+            with gr.Row(elem_id="col-container"):
+                image_generator = gr.Dropdown(label="Image Generator",
+                                             choices=["Hercai", "Prodia", "Pollinations"],
+                                             value="Hercai",
+                                             interactive=True)
+
+            # Create a dropdown menu for selecting the specific model within the chosen service
+            with gr.Row(elem_id="col-container"):
+                image_model = gr.Dropdown(label="Image Model",
+                                         choices=["v1", "v2", "v3", "simurg", "animefy", "raava", "shonin"],
+                                         value="v3",
+                                         interactive=True)
+
             # Create a button that triggers the video generation process
             with gr.Row(elem_id="col-container"):
                 button = gr.Button("Generate Video")
@@ -302,8 +354,8 @@ class Text2Video:
                 'That's okay', said the visitor. 'I will make stone soup for everyone'. Then he took a stone and dropped it into a giant pot,,"""
                 example = gr.Examples([example_txt], input_text)
 
-            # Define the button's click event to call the generate_video function with the user's input
-            button.click(self.generate_video, [input_text], output)
+            # Define the button's click event to call the generate_video function with the user's input and model selection
+            button.click(self.generate_video, [input_text, image_generator, image_model], output)
 
             LOGGER.info("Gradio interface launched successfully")
             # Launch the Gradio interface