import streamlit as st
import subprocess
import os

# Title
st.title("🎥 WAN 2.1 - ")

# Input fields
prompt = st.text_area("Enter your text prompt:", "A tiger walking.")
frame_num = st.slider("Number of frames:", min_value=30, max_value=120, value=60, step=10)
resolution = st.selectbox("Select resolution:", ["200*200", "1280*720"])
sample_steps = st.slider("Sampling steps:", min_value=10, max_value=50, value=20, step=5)

# Logging output placeholder
log_area = st.empty()

# Button to generate video
if st.button("🚀 Generate Video"):
    st.info("Generating video... This may take a few minutes.")
    
    # Define the command
    command = (
        f"python generate.py --task t2v-14B --size {resolution} "
        f"--frame_num {frame_num} --sample_steps {sample_steps} "
        f"--ckpt_dir ./Wan2.1-T2V-14B --offload_model True --prompt \"{prompt}\""
    )
    
    # Log command execution
    log_area.text_area("🔹 Running Command:", command)

    # Run process and stream output
    with subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=1, text=True) as process:
        logs = []
        for line in process.stdout:
            logs.append(line)
            log_area.text_area("📜 Logs (Real-time)", "".join(logs), height=250)
        
        # Capture stderr separately
        stderr_output = process.stderr.read()
        if stderr_output:
            logs.append("\n🔺 Error Logs:\n" + stderr_output)
            log_area.text_area("📜 Logs (Real-time)", "".join(logs), height=250)
    
    # Check if video was created
    if os.path.exists("output.mp4"):
        st.video("output.mp4")
        st.success("✅ Video generated successfully!")
    else:
        st.error("❌ Video generation failed! Check logs above.")