import streamlit as st
import replicate
import requests
from PIL import Image
from io import BytesIO

def run_replicate_model(api_key, input_params):
    output = replicate.run(
        "stability-ai/stable-video-diffusion:3f0457e4619daac51203dedb472816fd4af51f3149fa7a9e0b5ffcf1b8172438",
        api_key=api_key,
        input=input_params
    )
    return output

st.title("Stable Video Diffusion")
st.write("Fill in the input parameters and click on the button to run the Stable Video Diffusion model.")

# Input field for Replicate API token
api_key = st.text_input("Replicate API Token")

# Allow users to upload an image file or provide a URL
uploaded_file = st.file_uploader("Upload Image File", type=["jpg", "png", "jpeg"])
input_image = None
if uploaded_file is not None:
    input_image = uploaded_file
else:
    input_image = st.text_input("Input Image URL", "https://example.com/s-widua-icjKmXvsO7U-unsplash.jpg")

# Additional parameters
with st.expander("Additional Parameters"):
    cond_aug = st.number_input("Conditioning Augmentation", value=0.02, step=0.01)
    decoding_t = st.number_input("Decoding Time", value=7)
    sizing_strategy = st.selectbox("Sizing Strategy", ["maintain_aspect_ratio", "other_options"])
    motion_bucket_id = st.number_input("Motion Bucket ID", value=127)

frames_per_second = st.number_input("Frames Per Second", value=6)

input_params = {
    "cond_aug": cond_aug,
    "decoding_t": decoding_t,
    "input_image": input_image,
    "video_length": "14_frames_with_svd",
    "sizing_strategy": sizing_strategy,
    "motion_bucket_id": motion_bucket_id,
    "frames_per_second": frames_per_second
}

if st.button("Run Model"):
    if api_key:
        output_result = run_replicate_model(api_key, input_params)
        st.write("Output Video:")
        
        try:
            response = requests.get(output_result)
            video_bytes = BytesIO(response.content)
            st.video(video_bytes)
        except Exception as e:
            st.write("Unable to display the video.")
    else:
        st.write("Please provide your Replicate API Token.")