Spaces:
Runtime error
Runtime error
File size: 5,110 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import os, sys
import torch
from videogen_hub import MODEL_PATH
class LaVie():
def __init__(self, model_path=os.path.join(MODEL_PATH, "lavie"), device="cuda"):
"""
1. Download all necessary models from huggingface.
2. Initializes the LaVie model with a specific model path and device.
Args:
model_path (str, optional): The path to the model checkpoints. Defaults to "MODEL_PATH/lavie".
device (str, optional): The device to run the model on. Defaults to "cuda".
"""
# Put the source code imports here to avoid dependency version issues
from videogen_hub.pipelines.lavie.lavie_src.base.pipelines.pipeline_videogen import VideoGenPipeline
from videogen_hub.pipelines.lavie.lavie_src.base.download import find_model
from videogen_hub.pipelines.lavie.lavie_src.base.models.unet import UNet3DConditionModel
from diffusers.schedulers import DDPMScheduler
from diffusers.models import AutoencoderKL
from transformers import CLIPTokenizer, CLIPTextModel
from huggingface_hub import snapshot_download
from omegaconf import OmegaConf
snapshot_download(repo_id="Vchitect/LaVie", local_dir=model_path)
snapshot_download(repo_id="CompVis/stable-diffusion-v1-4", local_dir=os.path.join(model_path, "stable-diffusion-v1-4"))
snapshot_download(repo_id="stabilityai/stable-diffusion-x4-upscaler",
local_dir=os.path.join(model_path, "stable-diffusion-x4-upscaler"))
torch.set_grad_enabled(False)
self.device = device
config = {
"model_config": {
"use_compile": False,
"use_fp16": True,
"run_time": 0,
"guidance_scale": 7.5,
"num_sampling_steps": 50
},
"scheduler_config": {
"sample_method": "ddpm",
"beta_start": 0.0001,
"beta_end": 0.02,
"beta_schedule": "linear"
}
}
self.config = OmegaConf.create(config)
sd_path = os.path.join(model_path, "stable-diffusion-v1-4")
unet = UNet3DConditionModel.from_pretrained_2d(sd_path, subfolder="unet").to(device, dtype=torch.float16)
state_dict = find_model(os.path.join(model_path, "lavie_base.pt"))
unet.load_state_dict(state_dict)
vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae", torch_dtype=torch.float16).to(device)
tokenizer_one = CLIPTokenizer.from_pretrained(sd_path, subfolder="tokenizer")
text_encoder_one = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder",
torch_dtype=torch.float16).to(device) # huge
scheduler = DDPMScheduler.from_pretrained(sd_path,
subfolder="scheduler",
beta_start=self.config.scheduler_config.beta_start,
beta_end=self.config.scheduler_config.beta_end,
beta_schedule=self.config.scheduler_config.beta_schedule)
self.videogen_pipeline = VideoGenPipeline(vae=vae,
text_encoder=text_encoder_one,
tokenizer=tokenizer_one,
scheduler=scheduler,
unet=unet).to(device)
self.videogen_pipeline.enable_xformers_memory_efficient_attention()
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.
Args:
prompt (str, optional): The text prompt to generate the video from. Defaults to None.
size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
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.
"""
if seed is not None:
torch.manual_seed(seed)
videos = self.videogen_pipeline(prompt,
video_length=seconds * fps,
height=size[0],
width=size[1],
num_inference_steps=self.config.model_config.num_sampling_steps,
guidance_scale=self.config.model_config.guidance_scale).video
return videos[0]
|