Update app.py

app.py

@@ -1,112 +1,94 @@
 import gradio as gr
-import os
-from lumaai import LumaAI
+from lumaai import Lumaai
 import requests
 import time
 
-def generate_content(api_key, prompt, aspect_ratio, loop):
-    """
-    Generates video content using LumaAI's API based on user inputs.
-    Polls for the video once the generation is completed.
+def create_client(api_key):
+    return Lumaai(auth_token=api_key)
 
-    Parameters:
-    - api_key (str): User's LumaAI API key.
-    - 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: URL to the generated video or an error message.
-    """
+def generate_video(api_key, prompt, aspect_ratio, loop):
+    client = create_client(api_key)
     try:
-        # Initialize the LumaAI client with the provided API key
-        client = LumaAI(auth_token=api_key)
-        
-        # Create a generation request
         generation = client.generations.create(
-            aspect_ratio=aspect_ratio,
-            loop=loop,
             prompt=prompt,
+            aspect_ratio=aspect_ratio,
+            loop=loop
         )
-
-        # Polling to check if the generation is complete
-        generation_id = generation.id
-        video_url = None
-        
-        # Poll until the video is ready (adding a timeout of 60 seconds max)
-        timeout = 60  # seconds
-        start_time = time.time()
-        
-        while video_url is None and time.time() - start_time < timeout:
-            generation_status = client.generations.get(id=generation_id)
-            video_url = generation_status.assets.get('video', None)
-            if video_url is None:
-                time.sleep(5)  # Wait for 5 seconds before polling again
-
-        if video_url is None:
-            return "The video generation timed out. Please try again later."
-        return video_url
+        return generation.id, "Generation started..."
     except Exception as e:
-        return f"An error occurred: {str(e)}"
-
+        return None, f"Error: {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'
+def poll_generation(api_key, generation_id):
+    if not generation_id:
+        return "No generation in progress", None
     
-    with open(file_name, 'wb') as file:
-        file.write(response.content)
+    client = create_client(api_key)
+    try:
+        generation = client.generations.get(id=generation_id)
+        return generation.status, generation.assets.thumbnail
+    except Exception as e:
+        return f"Error: {str(e)}", None
+
+def download_video(api_key, generation_id):
+    if not generation_id:
+        return None
     
-    return file_name
+    client = create_client(api_key)
+    try:
+        generation = client.generations.get(id=generation_id)
+        video_url = generation.assets.video
+        response = requests.get(video_url, stream=True)
+        
+        file_name = f"generated_video_{generation_id}.mp4"
+        with open(file_name, 'wb') as file:
+            file.write(response.content)
+        
+        return file_name
+    except Exception as e:
+        return None
 
-# Create the Gradio Blocks interface
 with gr.Blocks() as demo:
-    with gr.Row():
-        gr.Markdown("# LumaAI Video Generator")
-
-    with gr.Row():
-        api_key = gr.Textbox(label="LumaAI API Key", placeholder="Enter your LumaAI API Key", type="password")
+    gr.Markdown("# Luma AI Video Generation Demo")
+    gr.Markdown("Generate videos using Luma AI based on text prompts.")
     
-    with gr.Row():
-        prompt = gr.Textbox(label="Prompt", placeholder="Describe what you want to generate...", lines=2)
+    api_key = gr.Textbox(label="Enter your Luma AI API Key", type="password")
     
     with gr.Row():
-        aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "4:3", "1:1", "21:9"], value="16:9")
+        with gr.Column(scale=2):
+            prompt = gr.Textbox(label="Prompt")
+            aspect_ratio = gr.Dropdown(["16:9", "9:16", "1:1", "4:3", "3:4"], label="Aspect Ratio")
+            loop = gr.Checkbox(label="Loop")
+            generate_btn = gr.Button("Generate Video")
+        
+        with gr.Column(scale=3):
+            status = gr.Textbox(label="Status")
+            thumbnail = gr.Image(label="Thumbnail")
+            video = gr.Video(label="Generated Video")
     
-    with gr.Row():
-        loop = gr.Checkbox(label="Loop", value=False)
-
-    with gr.Row():
-        generate_button = gr.Button("Generate Video")
+    generation_id = gr.State()
     
-    with gr.Row():
-        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(
-        handle_generation,
+    def on_generate(api_key, prompt, aspect_ratio, loop):
+        gen_id, message = generate_video(api_key, prompt, aspect_ratio, loop)
+        return gen_id, message
+    
+    generate_btn.click(
+        on_generate,
         inputs=[api_key, prompt, aspect_ratio, loop],
-        outputs=output
+        outputs=[generation_id, status]
+    )
+    
+    def on_poll(api_key, gen_id):
+        status, thumb = poll_generation(api_key, gen_id)
+        if status == "completed":
+            video_path = download_video(api_key, gen_id)
+            return status, thumb, video_path
+        return status, thumb, None
+    
+    poll_btn = gr.Button("Check Status")
+    poll_btn.click(
+        on_poll,
+        inputs=[api_key, generation_id],
+        outputs=[status, thumbnail, video]
     )
 
-# Launch the demo
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()