import streamlit as st
from diffusers import StableDiffusionImg2ImgPipeline
from moviepy.editor import ImageSequenceClip
from PIL import Image
import torch
import os

# Title and instructions
st.title("Image-to-Video Conversion")
st.write("Upload an image, provide a prompt, and generate a video using AI.")

# Sidebar for user input
st.sidebar.title("Settings")
num_frames = st.sidebar.slider("Number of Frames", 5, 50, 10)
fps = st.sidebar.slider("Frames Per Second (FPS)", 1, 30, 12)
guidance_scale = st.sidebar.slider("Guidance Scale", 5.0, 15.0, 7.5)
strength_base = st.sidebar.slider("Base Strength (Image Influence)", 0.1, 1.0, 0.5)

# File uploader
uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "png", "jpeg"])
prompt = st.text_input("Enter a Prompt", value="A cinematic animation of a sunset over mountains")

# Load the pre-trained model
@st.cache_resource
def load_model():
    return StableDiffusionImg2ImgPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4",
        torch_dtype=torch.float16,
        revision="fp16"
    ).to("cuda")

pipe = load_model()

# Process the uploaded image and generate video frames
if uploaded_image and st.button("Generate Video"):
    # Load the input image
    input_image = Image.open(uploaded_image).convert("RGB")
    st.image(input_image, caption="Uploaded Image", use_column_width=True)

    st.write("Generating video frames... This might take a few minutes.")
    progress = st.progress(0)

    frames = []
    for i in range(num_frames):
        progress.progress((i + 1) / num_frames)
        result = pipe(
            prompt=prompt,
            image=input_image,
            strength=strength_base + (i * 0.05),  # Incremental strength
            guidance_scale=guidance_scale
        )
        frames.append(result.images[0])

    # Save video frames as a video file
    video_path = "./output_video.mp4"
    video_clip = ImageSequenceClip([frame for frame in frames], fps=fps)
    video_clip.write_videofile(video_path, codec="libx264")
    st.success("Video generated successfully!")

    # Display video
    st.video(video_path)

    # Download link for the video
    with open(video_path, "rb") as file:
        btn = st.download_button(
            label="Download Video",
            data=file,
            file_name="output_video.mp4",
            mime="video/mp4"
        )

# Footer
st.write("Powered by Hugging Face Diffusers and Streamlit")