from fastapi import FastAPI
from pydantic import BaseModel
import torch
from diffusers import StableVideoDiffusionPipeline

# API Init
app = FastAPI()

# Load Model
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid"
)
pipe.to(device)

# Request Model
class VideoRequest(BaseModel):
    prompt: str

# Root Endpoint
@app.get("/")
def home():
    return {"message": "AI Video Generator API is running!"}

# Generate Video Endpoint
@app.post("/generate-video")
def generate_video(request: VideoRequest):
    video_frames = pipe(request.prompt, num_inference_steps=50).frames
    video_path = "output.mp4"
    video_frames[0].save(video_path)
    return {"message": "Video generated successfully!", "video_url": video_path}

# Run API
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7860)