File size: 932 Bytes
43c9d4d
 
 
 
79d3231
43c9d4d
 
3f3c56c
43c9d4d
 
 
 
 
 
afcab5d
43c9d4d
 
 
afcab5d
43c9d4d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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)