Spaces:
Running
Running
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 | |
def home(): | |
return {"message": "AI Video Generator API is running!"} | |
# Generate Video Endpoint | |
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) |