app.py

import os
import sys
import streamlit as st

# Ensure moviepy and dependencies are installed
def install_packages():
    import pip
    try:
        import moviepy
    except ModuleNotFoundError:
        pip.main(["install", "moviepy", "ffmpeg", "imageio[ffmpeg]"])

install_packages()

import torch
import numpy as np
from PIL import Image
import moviepy.editor as mp  # Now this should work
from diffusers import StableDiffusionPipeline

# Load Stable Diffusion model (optimized for CPU)
@st.cache_resource
def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-1-base",
        torch_dtype=torch.float16
    )
    pipe.enable_attention_slicing()
    pipe = pipe.to("cpu")  # Use CPU
    return pipe

pipe = load_model()

# Function to generate an image from a text prompt
def generate_image(prompt):
    image = pipe(prompt).images[0]
    image_path = "generated_image.png"
    image.save(image_path)
    return image_path

# Function to overlay image on video
def overlay_image_on_video(video_path, image_path, output_path="output_with_overlay.mp4"):
    video = mp.VideoFileClip(video_path)
    image = Image.open(image_path).resize((video.size[0]//3, video.size[1]//3)) 

    overlay = mp.ImageClip(np.array(image), duration=video.duration)
    overlay = overlay.set_position(("center", "bottom")).set_duration(video.duration)
    
    final_video = mp.CompositeVideoClip([video, overlay])
    final_video.write_videofile(output_path, codec="libx264", fps=24)
    return output_path

# Streamlit UI
st.title("AI Video Generator (CPU Optimized)")

prompt = st.text_input("Enter your prompt:")
generate_button = st.button("Generate Image")

if generate_button and prompt:
    st.write("Generating image... This may take a few minutes on CPU.")
    image_path = generate_image(prompt)
    st.image(image_path, caption="Generated Image", use_column_width=True)

# Video Upload
st.write("Upload a video to overlay the image")
video_file = st.file_uploader("Choose a video", type=["mp4"])

if video_file is not None:
    video_path = "uploaded_video.mp4"
    with open(video_path, "wb") as f:
        f.write(video_file.getvalue())

    st.video(video_path)

    if st.button("Overlay Image on Video"):
        output_video_path = overlay_image_on_video(video_path, image_path)
        st.video(output_video_path)
        st.download_button("Download Final Video", data=open(output_video_path, "rb"), file_name="final_video.mp4", mime="video/mp4")