Spaces:
Runtime error
Runtime error
| import os | |
| from huggingface_hub import snapshot_download, hf_hub_download | |
| import torch | |
| from videogen_hub import MODEL_PATH | |
| class OpenSoraPlan(): | |
| def __init__(self, device="cuda"): | |
| """ | |
| 1. Download the pretrained model and put it inside MODEL_PATH | |
| 2. Create Pipeline | |
| Note: it seems that the model needed from model_dir cannot support cpu | |
| Args: | |
| device: 'cuda' or 'cpu' the device to use the model | |
| """ | |
| from videogen_hub.pipelines.opensora_plan.opensora.sample_t2v import OpenSoraPlanPipeline | |
| model_path = snapshot_download('LanguageBind/Open-Sora-Plan-v1.1.0', local_dir = os.path.join(MODEL_PATH, 'Open-Sora-Plan-v1.1.0')) | |
| arg_list = ['--model_path', model_path, | |
| '--version', '65x512x512', | |
| '--num_frames', '65', | |
| '--height', '512', | |
| '--width', '512', | |
| '--cache_dir', MODEL_PATH, | |
| '--text_encoder_name', 'DeepFloyd/t5-v1_1-xxl', | |
| '--text_prompt', 'prompt_list_0.txt', | |
| '--ae', 'CausalVAEModel_4x8x8', | |
| '--ae_path', "/remote-home1/yeyang/CausalVAEModel_4x8x8", | |
| '--save_img_path', "./sample_video_65x512x512", | |
| '--fps', '24', | |
| '--guidance_scale', '7.5', | |
| '--num_sampling_steps', '150', | |
| '--enable_tiling'] | |
| self.pipeline = OpenSoraPlanPipeline(arg_list, device) | |
| def infer_one_video( | |
| self, | |
| prompt: str = None, | |
| size: list = [320, 512], | |
| seconds: int = 2, | |
| fps: int = 8, | |
| seed: int = 42, | |
| ): | |
| """ | |
| Generates a single video based on the provided prompt and parameters. | |
| Note that there are only 3 available shapes: (1 or 65 or 221)xHxW | |
| The output is of shape [frames, channels, height, width]. | |
| 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. | |
| """ | |
| torch.manual_seed(seed) | |
| self.pipeline.args.text_prompt = prompt | |
| self.pipeline.args.num_frames = fps * seconds | |
| self.pipeline.args.fps = fps | |
| self.pipeline.args.height = size[0] | |
| self.pipeline.args.width = size[1] | |
| samples = self.pipeline.inference(save_output=False) | |
| # samples is torch.Size([B, T, H, W, C]) | |
| output = samples.squeeze(0).permute(0, 3, 1, 2).cpu().float() | |
| return output | |