import streamlit as st
import torch
import subprocess
import os

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Title
st.title("WAN 2.1 - 1.3B Text-to-Video Generator 🎥")

# Model selection
model_path = "./Wan2.1-T2V-1.3B"

# Input fields
prompt = st.text_area("Enter your text prompt:", "A cat in military dress wearing headphones, laughing and walking.")
frame_num = st.slider("Number of frames:", min_value=30, max_value=180, value=60, step=10)
resolution = st.selectbox("Select resolution:", ["832*480", "1280*720"])
sample_steps = st.slider("Sampling steps:", min_value=10, max_value=50, value=30, step=5)

# Button to generate video
if st.button("Generate Video"):
    st.info("Generating video... Please wait.")

    # Run WAN 2.1 with user settings
    command = f"python generate.py --task t2v-1.3B --size {resolution} --frame_num {frame_num} --sample_steps {sample_steps} --ckpt_dir {model_path} --prompt \"{prompt}\""
    
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    process.wait()

    # Display video if generated
    if os.path.exists("output.mp4"):
        st.video("output.mp4")
        st.success("✅ Video generated successfully!")
    else:
        st.error("❌ Video generation failed. Check logs for details.")