Update app.py

app.py

@@ -1,18 +1,21 @@
 import gradio as gr
+import os
 from lumaai import LumaAI
+import requests
 
 def generate_content(api_key, prompt, aspect_ratio, loop):
     """
-    Generates content using LumaAI's API based on user inputs.
+    Generates video content using LumaAI's API based on user inputs.
+    Polls for the video once the generation is completed.
 
     Parameters:
     - api_key (str): User's LumaAI API key.
-    - prompt (str): The prompt for content generation.
+    - prompt (str): The prompt for video generation.
     - aspect_ratio (str): Desired aspect ratio (e.g., "16:9").
     - loop (bool): Whether the generation should loop.
 
     Returns:
-    - str: The ID of the generated content or an error message.
+    - str: URL to the generated video or an error message.
     """
     try:
         # Initialize the LumaAI client with the provided API key
@@ -24,15 +27,43 @@ def generate_content(api_key, prompt, aspect_ratio, loop):
             loop=loop,
             prompt=prompt,
         )
+
+        # Polling to check if the generation is complete
+        generation_id = generation.id
+        video_url = None
         
-        return f"Generation ID: {generation.id}"
+        # Poll until the video is ready
+        while video_url is None:
+            generation_status = client.generations.get(id=generation_id)
+            video_url = generation_status.assets.get('video', None)
+
+        return video_url
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
+
+def download_video(video_url):
+    """
+    Downloads the video from the provided URL and saves it locally.
+
+    Parameters:
+    - video_url (str): URL of the video to download.
+
+    Returns:
+    - str: File path of the downloaded video.
+    """
+    response = requests.get(video_url, stream=True)
+    file_name = 'generated_video.mp4'
+    
+    with open(file_name, 'wb') as file:
+        file.write(response.content)
+    
+    return file_name
+
 # Create the Gradio Blocks interface
 with gr.Blocks() as demo:
     with gr.Row():
-        gr.Markdown("# LumaAI Content Generator")
+        gr.Markdown("# LumaAI Video Generator")
 
     with gr.Row():
         api_key = gr.Textbox(label="LumaAI API Key", placeholder="Enter your LumaAI API Key", type="password")
@@ -47,14 +78,23 @@ with gr.Blocks() as demo:
         loop = gr.Checkbox(label="Loop", value=False)
 
     with gr.Row():
-        generate_button = gr.Button("Generate")
+        generate_button = gr.Button("Generate Video")
     
     with gr.Row():
-        output = gr.Textbox(label="Generation Result")
-    
+        output = gr.Video(label="Generated Video")
+
+    # Define the button click behavior
+    def handle_generation(api_key, prompt, aspect_ratio, loop):
+        video_url = generate_content(api_key, prompt, aspect_ratio, loop)
+        if video_url.startswith("http"):
+            file_name = download_video(video_url)
+            return file_name
+        else:
+            return video_url
+
     # Link the button click to the function
     generate_button.click(
-        generate_content,
+        handle_generation,
         inputs=[api_key, prompt, aspect_ratio, loop],
         outputs=output
     )