Update app.py

app.py

@@ -1,94 +1,116 @@
 import gradio as gr
-from lumaai import Lumaai
-import requests
+import os
+from lumaai import LumaAI
 import time
+import requests
 
-def create_client(api_key):
-    return Lumaai(auth_token=api_key)
-
-def generate_video(api_key, prompt, aspect_ratio, loop):
-    client = create_client(api_key)
-    try:
-        generation = client.generations.create(
-            prompt=prompt,
-            aspect_ratio=aspect_ratio,
-            loop=loop
-        )
-        return generation.id, "Generation started..."
-    except Exception as e:
-        return None, f"Error: {str(e)}"
+def generate_video(api_key, prompt, loop=False, aspect_ratio="16:9"):
+    client = LumaAI(auth_token=api_key)
+    
+    generation = client.generations.create(
+        prompt=prompt,
+        loop=loop,
+        aspect_ratio=aspect_ratio
+    )
+    
+    # Poll for completion
+    while True:
+        status = client.generations.get(id=generation.id)
+        if status.status == "completed":
+            break
+        time.sleep(5)
+    
+    # Download the video
+    video_url = status.assets.video
+    response = requests.get(video_url, stream=True)
+    
+    file_name = f"luma_ai_generated_{generation.id}.mp4"
+    with open(file_name, 'wb') as file:
+        file.write(response.content)
+    
+    return file_name
 
-def poll_generation(api_key, generation_id):
-    if not generation_id:
-        return "No generation in progress", None
+def text_to_video(api_key, prompt, loop, aspect_ratio):
+    if not api_key:
+        return "Please enter your Luma AI API key."
     
-    client = create_client(api_key)
     try:
-        generation = client.generations.get(id=generation_id)
-        return generation.status, generation.assets.thumbnail
+        video_path = generate_video(api_key, prompt, loop, aspect_ratio)
+        return video_path
     except Exception as e:
-        return f"Error: {str(e)}", None
+        return f"An error occurred: {str(e)}"
 
-def download_video(api_key, generation_id):
-    if not generation_id:
-        return None
+def image_to_video(api_key, prompt, image_url, loop, aspect_ratio):
+    if not api_key:
+        return "Please enter your Luma AI API key."
     
-    client = create_client(api_key)
     try:
-        generation = client.generations.get(id=generation_id)
-        video_url = generation.assets.video
+        client = LumaAI(auth_token=api_key)
+        generation = client.generations.create(
+            prompt=prompt,
+            loop=loop,
+            aspect_ratio=aspect_ratio,
+            keyframes={
+                "frame0": {
+                    "type": "image",
+                    "url": image_url
+                }
+            }
+        )
+        
+        # Poll for completion
+        while True:
+            status = client.generations.get(id=generation.id)
+            if status.status == "completed":
+                break
+            time.sleep(5)
+        
+        # Download the video
+        video_url = status.assets.video
         response = requests.get(video_url, stream=True)
         
-        file_name = f"generated_video_{generation_id}.mp4"
+        file_name = f"luma_ai_generated_{generation.id}.mp4"
         with open(file_name, 'wb') as file:
             file.write(response.content)
         
         return file_name
     except Exception as e:
-        return None
+        return f"An error occurred: {str(e)}"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Luma AI Video Generation Demo")
-    gr.Markdown("Generate videos using Luma AI based on text prompts.")
+    gr.Markdown("# Luma AI Text-to-Video Demo")
     
-    api_key = gr.Textbox(label="Enter your Luma AI API Key", type="password")
+    api_key = gr.Textbox(label="Luma AI API Key", type="password")
     
-    with gr.Row():
-        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.Tab("Text to Video"):
+        prompt = gr.Textbox(label="Prompt")
+        generate_btn = gr.Button("Generate Video")
+        video_output = gr.Video(label="Generated Video")
         
-        with gr.Column(scale=3):
-            status = gr.Textbox(label="Status")
-            thumbnail = gr.Image(label="Thumbnail")
-            video = gr.Video(label="Generated Video")
-    
-    generation_id = gr.State()
-    
-    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=[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
+        with gr.Accordion("Advanced Options", open=False):
+            loop = gr.Checkbox(label="Loop", value=False)
+            aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "1:1", "9:16", "4:3", "3:4"], value="16:9")
+        
+        generate_btn.click(
+            text_to_video,
+            inputs=[api_key, prompt, loop, aspect_ratio],
+            outputs=video_output
+        )
     
-    poll_btn = gr.Button("Check Status")
-    poll_btn.click(
-        on_poll,
-        inputs=[api_key, generation_id],
-        outputs=[status, thumbnail, video]
-    )
+    with gr.Tab("Image to Video"):
+        img_prompt = gr.Textbox(label="Prompt")
+        img_url = gr.Textbox(label="Image URL")
+        img_generate_btn = gr.Button("Generate Video from Image")
+        img_video_output = gr.Video(label="Generated Video")
+        
+        with gr.Accordion("Advanced Options", open=False):
+            img_loop = gr.Checkbox(label="Loop", value=False)
+            img_aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "1:1", "9:16", "4:3", "3:4"], value="16:9")
+        
+        img_generate_btn.click(
+            image_to_video,
+            inputs=[api_key, img_prompt, img_url, img_loop, img_aspect_ratio],
+            outputs=img_video_output
+        )
 
 demo.launch()