dfdfdsfgs's picture
Upload project files
d9486d1
raw
history blame
6.05 kB
import gradio as gr
import uuid
import subprocess
import threading
import os
import time
from fastapi import FastAPI
from fastapi.responses import FileResponse
import asyncio
# A simple in-memory dictionary to track task status.
# For a production system, you'd use a database or Redis.
tasks = {}
def run_video_generation(task_id: str, topic: str, context: str):
"""
This function runs the main generation script in a separate process.
"""
tasks[task_id]['status'] = 'running'
# Sanitize topic to create a valid directory name
file_prefix = "".join(c if c.isalnum() else "_" for c in topic.lower())
output_dir = os.path.join("output", file_prefix)
command = [
"python", "generate_video.py",
"--model", "openai/o3-mini", # Or get from request
"--topic", topic,
"--context", context,
"--output_dir", "output",
"--use_langfuse" # Assuming you have secrets set
]
try:
# Using subprocess to run the existing script
process = subprocess.run(command, check=True, capture_output=True, text=True)
# Assume the final video is named based on the topic
# Note: The actual video path might differ. This is an assumption.
# You may need to parse the stdout from generate_video.py to get the exact path.
video_path = None
for file in os.listdir(output_dir):
if file.endswith("_combined.mp4"):
video_path = os.path.join(output_dir, file)
break
if video_path and os.path.exists(video_path):
tasks[task_id]['status'] = 'completed'
tasks[task_id]['video_path'] = video_path
else:
tasks[task_id]['status'] = 'failed'
tasks[task_id]['error'] = "Video file not found after generation."
tasks[task_id]['stdout'] = process.stdout
tasks[task_id]['stderr'] = process.stderr
except subprocess.CalledProcessError as e:
tasks[task_id]['status'] = 'failed'
tasks[task_id]['error'] = str(e)
tasks[task_id]['stdout'] = e.stdout
tasks[task_id]['stderr'] = e.stderr
except Exception as e:
tasks[task_id]['status'] = 'failed'
tasks[task_id]['error'] = str(e)
def start_generation_thread(topic: str, context: str):
if not topic or not context:
return "Topic and Context cannot be empty.", "", None
task_id = str(uuid.uuid4())
tasks[task_id] = {'status': 'queued'}
# Use a background thread to run the time-consuming task
thread = threading.Thread(
target=run_video_generation,
args=(task_id, topic, context)
)
thread.start()
return f"Task started. Your Task ID is: {task_id}", task_id, None
def check_status(task_id: str):
if not task_id:
return "Please provide a Task ID.", None
task = tasks.get(task_id)
if not task:
return "Task not found.", None
status = task.get('status')
if status == 'completed':
video_path = task.get('video_path')
return f"Status: {status}", video_path
elif status == 'failed':
error = task.get('error', 'Unknown error')
stdout = task.get('stdout', '')
stderr = task.get('stderr', '')
return f"Status: {status}\nError: {error}\nOutput: {stdout}\nStderr: {stderr}", None
return f"Status: {status}", None
# We need a lightweight FastAPI app in the background to serve the video files.
# Gradio can't serve files directly from arbitrary paths in a secure way.
fastapi_app = FastAPI()
@fastapi_app.get("/videos/{task_id}")
def get_video(task_id: str):
"""
Serves the final generated video file.
"""
task = tasks.get(task_id)
if not task or task.get('status') != 'completed':
return {"error": "Task not completed or not found"}
video_path = task.get('video_path')
if not os.path.exists(video_path):
return {"error": "Video file not found."}
return FileResponse(video_path, media_type="video/mp4", filename=os.path.basename(video_path))
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("# Theorem-Explain-Agent Video Generation")
gr.Markdown("Start a video generation task and check its status.")
with gr.Tab("Start Generation"):
topic_input = gr.Textbox(label="Topic", placeholder="e.g., The Pythagorean Theorem")
context_input = gr.Textbox(label="Context", placeholder="A short explanation of the theorem.")
start_button = gr.Button("Generate Video")
with gr.Column():
task_id_output = gr.Textbox(label="Task ID", interactive=False)
status_output_start = gr.Textbox(label="Status", interactive=False)
with gr.Tab("Check Status"):
task_id_input = gr.Textbox(label="Task ID", placeholder="Enter the Task ID you received.")
check_button = gr.Button("Check Status")
with gr.Column():
status_output_check = gr.Textbox(label="Status", interactive=False)
video_output = gr.Video(label="Generated Video")
# Actions
start_button.click(
fn=start_generation_thread,
inputs=[topic_input, context_input],
outputs=[status_output_start, task_id_output, video_output] # Clear video on new task
)
check_button.click(
fn=check_status,
inputs=[task_id_input],
outputs=[status_output_check, video_output]
)
gr.Markdown("### How to Use")
gr.Markdown(
"1. Enter a `Topic` and `Context` in the 'Start Generation' tab and click 'Generate Video'.\n"
"2. Copy the `Task ID` that appears.\n"
"3. Go to the 'Check Status' tab, paste the `Task ID`, and click 'Check Status' periodically.\n"
"4. When the generation is complete, the video will appear."
)
# To run both Gradio and FastAPI, we mount the FastAPI app into Gradio's internal FastAPI app.
app = gr.mount_ όπου(demo, fastapi_app, path="/")