From a4a1b93bbcfe5eae3071b77c3012f0566ecd9ee6 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sat, 6 Apr 2024 19:26:08 +0800 Subject: [PATCH] refactor vae to hf Former-commit-id: a938d3f7f962709f7ecff5e9d49c9bf1ce40d7fa --- .../ae/videobase/causal_vae/__init__.py | 8 +- .../models/ae/videobase/modeling_videobase.py | 3 +- opensora/sample/sample_t2v.py | 35 ++---- opensora/serve/gradio_utils.py | 82 +++++++++++++ opensora/serve/gradio_web_server.py | 116 ++++++++++++++++++ 5 files changed, 217 insertions(+), 27 deletions(-) create mode 100644 opensora/serve/gradio_utils.py create mode 100644 opensora/serve/gradio_web_server.py diff --git a/opensora/models/ae/videobase/causal_vae/__init__.py b/opensora/models/ae/videobase/causal_vae/__init__.py index 85e75f141..b32009597 100644 --- a/opensora/models/ae/videobase/causal_vae/__init__.py +++ b/opensora/models/ae/videobase/causal_vae/__init__.py @@ -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) diff --git a/opensora/models/ae/videobase/modeling_videobase.py b/opensora/models/ae/videobase/modeling_videobase.py index d78d17e57..efa689d1b 100644 --- a/opensora/models/ae/videobase/modeling_videobase.py +++ b/opensora/models/ae/videobase/modeling_videobase.py @@ -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) diff --git a/opensora/sample/sample_t2v.py b/opensora/sample/sample_t2v.py index cb70d11ff..8b3d6ff0e 100644 --- a/opensora/sample/sample_t2v.py +++ b/opensora/sample/sample_t2v.py @@ -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() diff --git a/opensora/serve/gradio_utils.py b/opensora/serve/gradio_utils.py new file mode 100644 index 000000000..a77c7f824 --- /dev/null +++ b/opensora/serve/gradio_utils.py @@ -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 = (""" +
+ + + Open-Sora-Plan🚀 + +
+

Open-Sora-Plan v1.0.0

+

If you like our project, please give us a star ✨ on Github for the latest update.

+
+
+
+
+
+ + +
+
+""") + +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, + ] +] \ No newline at end of file diff --git a/opensora/serve/gradio_web_server.py b/opensora/serve/gradio_web_server.py new file mode 100644 index 000000000..38e266623 --- /dev/null +++ b/opensora/serve/gradio_web_server.py @@ -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() \ No newline at end of file