Spaces:
Runtime error
Runtime error
File size: 2,231 Bytes
413d4d0 |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import os
import torch
from huggingface_hub import snapshot_download
from videogen_hub import MODEL_PATH
class ModelScope:
def __init__(self, device="gpu"):
"""
1. Download the pretrained model and put it inside checkpoints/modelscope
2. Create Pipeline
Note: it seems that the model needed from model_dir cannot support cpu
Args:
device: 'gpu' or 'cpu' the device to use the model
"""
from modelscope.pipelines import pipeline
from modelscope.models import Model
model_dir = snapshot_download(
repo_id="ali-vilab/modelscope-damo-text-to-video-synthesis",
local_dir=os.path.join(MODEL_PATH, "modelscope"),
)
model = Model.from_pretrained(model_dir)
self.pipeline = pipeline("text-to-video-synthesis", model=model, device=device)
def infer_one_video(
self, prompt: str = None, seconds: int = 2, fps: int = 8, seed: int = 42
):
"""
Generates a single video based on the provided prompt and parameters.
The generated video always has resolution 256x256
Args:
prompt (str, optional): The text prompt to generate the video from. Defaults to None.
seconds (int, optional): The duration of the video in seconds. Defaults to 2.
fps (int, optional): The frames per second of the video. Defaults to 8.
seed (int, optional): The seed for random number generation. Defaults to 42.
Returns:
torch.Tensor: The generated video as a tensor.
"""
from modelscope.outputs import OutputKeys
from decord import VideoReader
from decord import cpu, gpu
import io
torch.manual_seed(seed)
self.pipeline.model.config.model.model_args.max_frames = fps * seconds
test_text = {
"text": prompt,
}
output_video_path = self.pipeline(
test_text,
)[OutputKeys.OUTPUT_VIDEO]
result = io.BytesIO(output_video_path)
result = VideoReader(result, ctx=cpu(0))
result = torch.from_numpy(result.get_batch(range(len(result))).asnumpy())
return result
|