nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor vae to hf

Browse Source 
Former-commit-id: a938d3f7f962709f7ecff5e9d49c9bf1ce40d7fa
This commit is contained in:
LinB203
2024-04-06 19:26:08 +08:00
parent cea0c041df
commit a4a1b93bbc
5 changed files with 217 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
opensora/models/ae/videobase/causal_vae/__init__.py
+6 -2
View File
@@ -1,3 +1,5 @@
import os.path
import torch
from .configuration_causalvae import CausalVAEConfiguration
@@ -9,9 +11,11 @@ from einops import rearrange
from torch import nn
class CausalVAEModelWrapper(nn.Module):
def __init__(self, ckpt):
def __init__(self, model_path, subfolder=None, cache_dir=None):
super(CausalVAEModelWrapper, self).__init__()
self.vae = CausalVAEModel.load_from_checkpoint(ckpt)
# if os.path.exists(ckpt):
# self.vae = CausalVAEModel.load_from_checkpoint(ckpt)
self.vae = CausalVAEModel.from_pretrained(model_path, subfolder=subfolder, cache_dir=cache_dir)
def encode(self, x): # b c t h w
# x = self.vae.encode(x).sample()
x = self.vae.encode(x).sample().mul_(0.18215)
opensora/models/ae/videobase/modeling_videobase.py
+2 -1
View File
@@ -1,4 +1,5 @@
import torch
from diffusers import ModelMixin, ConfigMixin
from torch import nn
import os
import json
@@ -9,7 +10,7 @@ so there is no content here for the moment.
It will be gradually optimized over time.
Contributions from the open-source community are welcome.
"""
class VideoBaseAE(nn.Module):
class VideoBaseAE(ModelMixin, ConfigMixin):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
opensora/sample/sample_t2v.py
+11 -24
View File
@@ -32,32 +32,23 @@ def main(args):
torch.set_grad_enabled(False)
device = "cuda" if torch.cuda.is_available() else "cpu"
vae = getae(args).to(device, dtype=torch.float16)
vae = CausalVAEModelWrapper(args.model_path, subfolder="vae", cache_dir='cache_dir').to(device, dtype=torch.float16)
if args.enable_tiling:
vae.vae.enable_tiling()
vae.vae.tile_overlap_factor = args.tile_overlap_factor
if getae_wrapper(args) == CausalVQVAEModelWrapper or getae_wrapper(args) == CausalVAEModelWrapper:
video_length = args.num_frames // ae_stride_config[args.ae][0] + 1
else:
video_length = args.num_frames // ae_stride_config[args.ae][0]
video_length, image_size = int(args.version.split('x')[0]), int(args.version.split('x')[1])
if args.force_images:
video_length = 1
ext = '.jpg'
ext = 'jpg'
else:
ext = 'mp4'
# Load model:
latent_size = (args.image_size // ae_stride_config[args.ae][1], args.image_size // ae_stride_config[args.ae][2])
args.latent_size = latent_size
vae.latent_size = latent_size
transformer_model = LatteT2V.from_pretrained(args.ckpt, subfolder="model", torch_dtype=torch.float16).to(device)
transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, torch_dtype=torch.float16).to(device)
transformer_model.force_images = args.force_images
# tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir='./cache_dir')
# text_encoder = get_text_enc(args).to(device).eval()
tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir="cache_dir")
text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir="cache_dir", torch_dtype=torch.float16).to(device)
transformer_model.eval() # important!
# set eval mode
transformer_model.eval()
vae.eval()
@@ -101,8 +92,8 @@ def main(args):
print('Processing the ({}) prompt'.format(prompt))
videos = videogen_pipeline(prompt,
video_length=video_length,
height=args.image_size,
width=args.image_size,
height=image_size,
width=image_size,
num_inference_steps=args.num_sampling_steps,
guidance_scale=args.guidance_scale,
enable_temporal_attentions=not args.force_images,
@@ -113,14 +104,14 @@ def main(args):
if args.force_images:
videos = videos[:, 0].permute(0, 3, 1, 2) # b t h w c -> b c h w
save_image(videos / 255.0, os.path.join(args.save_img_path,
prompt.replace(' ', '_') + f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
prompt.replace(' ', '_')[:100] + f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
nrow=1, normalize=True, value_range=(0, 1)) # t c h w
else:
imageio.mimwrite(
os.path.join(
args.save_img_path,
prompt.replace(' ', '_') + f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
prompt.replace(' ', '_')[:100] + f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
), videos[0],
fps=args.fps, quality=9) # highest quality is 10, lowest is 0
except:
@@ -144,13 +135,12 @@ def main(args):
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--ckpt", type=str, default="")
parser.add_argument("--model", type=str, default='Latte-XL/122')
parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
parser.add_argument("--version", type=str, default='65x512x512', choices=['65x512x512', '65x256x256', '17x256x256'])
parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
parser.add_argument("--beta_start", type=float, default=0.0001)
parser.add_argument("--image_size", type=int, default=512)
parser.add_argument("--beta_end", type=float, default=0.02)
parser.add_argument("--guidance_scale", type=float, default=7.5)
parser.add_argument("--beta_schedule", type=str, default="linear")
@@ -158,12 +148,9 @@ if __name__ == "__main__":
parser.add_argument("--sample_method", type=str, default="PNDM")
parser.add_argument("--num_sampling_steps", type=int, default=50)
parser.add_argument("--fps", type=int, default=24)
parser.add_argument("--num_frames", type=int, default=65)
parser.add_argument("--run_time", type=int, default=0)
parser.add_argument("--text_prompt", nargs='+')
parser.add_argument('--force_images', action='store_true')
parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
parser.add_argument('--enable_tiling', action='store_true')
args = parser.parse_args()
opensora/serve/gradio_utils.py
+82
View File
@@ -0,0 +1,82 @@
import random
import torch
def set_env(seed=0):
torch.manual_seed(seed)
torch.set_grad_enabled(False)
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
if randomize_seed:
seed = random.randint(0, 203279)
return seed
title_markdown = ("""
<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
<a href="https://github.com/PKU-YuanGroup/Open-Sora-Plan" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">
<img src="https://www.pnglog.com/AOuPMh.png" alt="Open-Sora-Plan🚀" style="max-width: 200px; height: auto;">
</a>
<div>
<h1 >Open-Sora-Plan v1.0.0</h1>
<h3 style="margin: 0;">If you like our project, please give us a star ✨ on Github for the latest update.</h3>
</div>
</div>
<br>
<div align="center">
<div style="display:flex; gap: 0.25rem;" align="center">
<a href='https://github.com/PKU-YuanGroup/Open-Sora-Plan'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
<a href='https://github.com/PKU-YuanGroup/Open-Sora-Plan/stargazers'><img src='https://img.shields.io/github/stars/PKU-YuanGroup/Open-Sora-Plan.svg?style=social'></a>
</div>
</div>
""")
block_css = """
#buttons button {
min-width: min(120px,100%);
}
"""
examples = [
[
"A soaring drone footage captures the majestic beauty of a coastal cliff, "
"its red and yellow stratified rock faces rich in color and against the vibrant turquoise of the sea. "
"Seabirds can be seen taking flight around the cliff's precipices. "
"As the drone slowly moves from different angles, "
"the changing sunlight casts shifting shadows that highlight the rugged textures of the cliff and the surrounding calm sea. "
"The water gently laps at the rock base and the greenery that clings to the top of the cliff, "
"and the scene gives a sense of peaceful isolation at the fringes of the ocean. "
"The video captures the essence of pristine natural beauty untouched by human structures.",
50, 7.5,
],
[
"The video captures the majestic beauty of a waterfall cascading down a cliff into a serene lake. "
"The waterfall, with its powerful flow, is the central focus of the video. "
"The surrounding landscape is lush and green, with trees and foliage adding to the natural beauty of the scene. "
"The camera angle provides a bird's eye view of the waterfall, "
"allowing viewers to appreciate the full height and grandeur of the waterfall. "
"The video is a stunning representation of nature's power and beauty.",
50, 7.5,
],
[
"A vibrant scene of a snowy mountain landscape. "
"The sky is filled with a multitude of colorful hot air balloons, "
"each floating at different heights, creating a dynamic and lively atmosphere. "
"The balloons are scattered across the sky, some closer to the viewer, "
"others further away, adding depth to the scene. "
"Below, the mountainous terrain is blanketed in a thick layer of snow, "
"with a few patches of bare earth visible here and there. "
"The snow-covered mountains provide a stark contrast to the colorful balloons, "
"enhancing the visual appeal of the scene. In the foreground, "
"a few cars can be seen driving along a winding road that cuts through the mountains. "
"The cars are small compared to the vastness of the landscape, emphasizing the grandeur of the surroundings. "
"The overall style of the video is a mix of adventure and tranquility, "
"with the hot air balloons adding a touch of whimsy to the otherwise serene mountain landscape. "
"The video is likely shot during the day, as the lighting is bright and even, "
"casting soft shadows on the snow-covered mountains.",
50, 7.5,
]
]
opensora/serve/gradio_web_server.py
+116
View File
@@ -0,0 +1,116 @@
import argparse
import sys
import os
import random
import imageio
import torch
from diffusers import PNDMScheduler
from huggingface_hub import hf_hub_download
from torchvision.utils import save_image
from diffusers.models import AutoencoderKL
from datetime import datetime
from typing import List, Union
import gradio as gr
import numpy as np
from gradio.components import Textbox, Video, Image
from transformers import T5Tokenizer, T5EncoderModel
from opensora.models.ae import ae_stride_config, getae, getae_wrapper
from opensora.models.ae.videobase import CausalVQVAEModelWrapper, CausalVAEModelWrapper
from opensora.models.diffusion.latte.modeling_latte import LatteT2V
from opensora.sample.pipeline_videogen import VideoGenPipeline
from opensora.serve.gradio_utils import block_css, title_markdown, randomize_seed_fn, set_env, examples
@torch.inference_mode()
def generate_img(prompt, sample_steps, scale, seed=0, randomize_seed=False, force_images=False):
seed = int(randomize_seed_fn(seed, randomize_seed))
set_env(seed)
videos = videogen_pipeline(prompt,
video_length=transformer_model.config.video_length,
height=int(args.version.split('x')[1]),
width=int(args.version.split('x')[2]),
num_inference_steps=sample_steps,
guidance_scale=scale,
enable_temporal_attentions=not force_images,
num_images_per_prompt=1,
mask_feature=True,
).video
torch.cuda.empty_cache()
videos = videos[0]
tmp_save_path = 'tmp.mp4'
imageio.mimwrite(tmp_save_path, videos, fps=24, quality=9) # highest quality is 10, lowest is 0
display_model_info = f"Image size: {int(args.version.split('x')[1])}, \nSampling Step: {sample_steps}, \nGuidance Scale: {scale}"
return tmp_save_path, prompt, display_model_info, seed
if __name__ == '__main__':
args = type('args', (), {
'force_images': False,
'model_path': 'LanguageBind/Open-Sora-Plan-65x512x512-v1.0.0',
'text_encoder_name': 'DeepFloyd/t5-v1_1-xxl',
'version': '65x512x512'
})
device = torch.device('cuda:0')
# Load model:
transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, torch_dtype=torch.float16, cache_dir='cache_dir').to(device)
vae = CausalVAEModelWrapper(args.model_path, subfolder="vae", cache_dir='cache_dir').to(device, dtype=torch.float16)
vae.vae.enable_tiling()
transformer_model.force_images = args.force_images
tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir="cache_dir")
text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir="cache_dir",
torch_dtype=torch.float16).to(device)
# set eval mode
transformer_model.eval()
vae.eval()
text_encoder.eval()
scheduler = PNDMScheduler()
videogen_pipeline = VideoGenPipeline(vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
scheduler=scheduler,
transformer=transformer_model).to(device=device)
demo = gr.Interface(
fn=generate_img,
inputs=[Textbox(label="",
placeholder="Please enter your prompt. \n"),
gr.Slider(
label='Sample Steps',
minimum=1,
maximum=500,
value=50,
step=10
),
gr.Slider(
label='Guidance Scale',
minimum=0.1,
maximum=30.0,
value=10.0,
step=0.1
),
gr.Slider(
label="Seed",
minimum=0,
maximum=203279,
step=1,
value=0,
),
gr.Checkbox(label="Randomize seed", value=True),
],
outputs=[Video(label="Vid", width=512, height=512),
Textbox(label="input prompt"),
Textbox(label="model info"),
gr.Slider(label='seed')],
title=title_markdown, theme=gr.themes.Default(), css=block_css,
examples=examples,
)
demo.launch()