76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
import torch
|
|
from .utils import (
|
|
VIDEO_ASPECT_RATIOS,
|
|
call_api_with_fallback,
|
|
resolve_api_key,
|
|
resolve_generation_url,
|
|
upload_image,
|
|
)
|
|
|
|
|
|
def _get_video_url(data: dict) -> str:
|
|
url = data.get("generation_url", "")
|
|
if not url:
|
|
raise RuntimeError(f"Pruna API returned no generation_url. Response: {data}")
|
|
return resolve_generation_url(url)
|
|
|
|
|
|
class PrunaVideo:
|
|
"""Generate a video from a text prompt and/or an image using the Pruna p-video model.
|
|
|
|
When an image is connected the model animates it; without an image it generates
|
|
from the text prompt alone. Returns the video URL as a STRING.
|
|
|
|
Pricing (p-video):
|
|
720p draft OFF — $0.02/sec | 720p draft ON — $0.005/sec
|
|
1080p draft OFF — $0.04/sec | 1080p draft ON — $0.01/sec
|
|
"""
|
|
|
|
CATEGORY = "pruna ai"
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("video_url",)
|
|
FUNCTION = "generate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"duration": ("INT", {"default": 5, "min": 1, "max": 20, "step": 1}),
|
|
"resolution": (["720p", "1080p"], {"default": "720p"}),
|
|
"aspect_ratio": (VIDEO_ASPECT_RATIOS, {"default": "16:9"}),
|
|
"draft": ("BOOLEAN", {"default": False}),
|
|
"api_key": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
def generate(
|
|
self,
|
|
prompt: str,
|
|
duration: int,
|
|
resolution: str,
|
|
aspect_ratio: str,
|
|
draft: bool,
|
|
api_key: str,
|
|
image: torch.Tensor = None,
|
|
):
|
|
key = resolve_api_key(api_key)
|
|
payload: dict = {
|
|
"input": {
|
|
"prompt": prompt,
|
|
"duration": duration,
|
|
"resolution": resolution,
|
|
"aspect_ratio": aspect_ratio,
|
|
"draft": draft,
|
|
}
|
|
}
|
|
|
|
if image is not None:
|
|
payload["input"]["image"] = upload_image(image, key)
|
|
|
|
data = call_api_with_fallback("p-video", payload, key)
|
|
return (_get_video_url(data),)
|