import os
import streamlit as st
from diffusers import CogVideoXImageToVideoPipeline
from diffusers.utils import export_to_video, load_image
import torch

# Streamlit interface for uploading an image and inputting a prompt
st.title("Image to Video with Hugging Face")
st.write("Upload an image and provide a prompt to generate a video.")

# File uploader for the input image
uploaded_file = st.file_uploader("Upload an image (JPG or PNG):", type=["jpg", "jpeg", "png"])
prompt = st.text_input("Enter your prompt:", "A little girl is riding a bicycle at high speed. Focused, detailed, realistic.")

if uploaded_file and prompt:
    try:
        # Save the uploaded file to a temporary location
        with open("uploaded_image.jpg", "wb") as f:
            f.write(uploaded_file.read())

        # Load the image
        image = load_image("uploaded_image.jpg")

        # Initialize the CogVideoX pipeline
        st.write("Initializing the pipeline...")
        pipe = CogVideoXImageToVideoPipeline.from_pretrained(
            "THUDM/CogVideoX1.5-5B-I2V",
            torch_dtype=torch.bfloat16
        )

        pipe.enable_sequential_cpu_offload()
        pipe.vae.enable_tiling()
        pipe.vae.enable_slicing()

        # Generate the video
        st.write("Generating video... this may take a while.")
        video_frames = pipe(
            prompt=prompt,
            image=image,
            num_videos_per_prompt=1,
            num_inference_steps=50,
            num_frames=81,
            guidance_scale=6,
            generator=torch.Generator(device="cuda").manual_seed(42),
        ).frames[0]

        # Export the video
        video_path = "output.mp4"
        export_to_video(video_frames, video_path, fps=8)

        # Display the video in Streamlit
        st.video(video_path)

    except Exception as e:
        st.error(f"An error occurred: {e}")
else:
    st.write("Please upload an image and provide a prompt to get started.")