Faizkhanvideo / app.py
Faizbulbul's picture
Update app.py
43c9d4d verified
raw
history blame contribute delete
932 Bytes
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)