From 1f7edca9fda44f7e1cdf86bb4d2ae112dc9fd72f Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Mon, 8 Apr 2024 18:57:06 +0800 Subject: [PATCH] clean Former-commit-id: a7f2c1c8b2243587cedd6e152e64298842287095 --- examples/rec_imvi_vae.py | 66 +++++----- examples/rec_video.py | 120 ------------------ examples/rec_video_ae.py | 120 ------------------ .../models/diffusion/latte/modeling_latte.py | 6 +- opensora/models/diffusion/latte/modules.py | 91 ++++++++++--- opensora/serve/gradio_utils.py | 2 +- opensora/serve/gradio_web_server.py | 5 +- scripts/causalvae/reconstruction.sh | 11 ++ scripts/text_condition/sample_video.sh | 6 +- 9 files changed, 130 insertions(+), 297 deletions(-) delete mode 100644 examples/rec_video.py delete mode 100644 examples/rec_video_ae.py create mode 100644 scripts/causalvae/reconstruction.sh diff --git a/examples/rec_imvi_vae.py b/examples/rec_imvi_vae.py index a0f1469d9..cb280da08 100644 --- a/examples/rec_imvi_vae.py +++ b/examples/rec_imvi_vae.py @@ -14,6 +14,9 @@ from pytorchvideo.transforms import ShortSideScale from torchvision.transforms import Lambda, Compose import sys + +from opensora.models.ae import getae_wrapper + sys.path.append(".") from opensora.dataset.transform import CenterCropVideo, resize from opensora.models.ae.videobase import CausalVAEModel @@ -30,15 +33,17 @@ def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str video_writer.release() + def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None: x = x.detach().cpu() x = torch.clamp(x, -1, 1) x = (x + 1) / 2 - x = x.permute(1, 2, 3, 0).numpy() - x = (255*x).astype(np.uint8) + x = x.permute(0, 2, 3, 1).numpy() + x = (255 * x).astype(np.uint8) array_to_video(x, fps=fps, output_file=output_file) return + def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor: decord_vr = VideoReader(video_path, ctx=cpu(0)) total_frames = len(decord_vr) @@ -56,7 +61,6 @@ def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tens print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path, total_frames) - frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int) video_data = decord_vr.get_batch(frame_id_list).asnumpy() video_data = torch.from_numpy(video_data) @@ -89,8 +93,8 @@ class ResizeVideo: def __repr__(self) -> str: return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" -def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor: +def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor: transform = Compose( [ Lambda(lambda x: ((x / 255.0) * 2 - 1)), @@ -106,30 +110,26 @@ def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optio def main(args: argparse.Namespace): - video_path = args.video_path - num_frames = args.num_frames - resolution = args.resolution - crop_size = args.crop_size - sample_fps = args.sample_fps - sample_rate = args.sample_rate device = args.device - vqvae = CausalVAEModel.from_pretrained(args.ckpt) + kwarg = {} + vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir', **kwarg).to(device) if args.enable_tiling: - vqvae.enable_tiling() - vqvae.tile_overlap_factor = args.tile_overlap_factor - vqvae.eval() - vqvae = vqvae.to(device) - vqvae = vqvae.to(torch.float16) + vae.vae.enable_tiling() + vae.vae.tile_overlap_factor = args.tile_overlap_factor + vae.eval() + vae = vae.to(device) + vae = vae.half() with torch.no_grad(): - x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) - x_vae = x_vae.to(device) # b c t h w - x_vae = x_vae.to(torch.float16) - latents = vqvae.encode(x_vae).sample().to(torch.float16) - video_recon = vqvae.decode(latents) - + x_vae = preprocess(read_video(args.video_path, args.num_frames, args.sample_rate), args.resolution, + args.crop_size) + x_vae = x_vae.to(device, dtype=torch.float16) # b c t h w + latents = vae.encode(x_vae) + latents = latents.to(torch.float16) + video_recon = vae.decode(latents) # b t c h w + if video_recon.shape[2] == 1: - x = video_recon[0, :, 0, :, :] + x = video_recon[0, 0, :, :, :] x = x.squeeze() x = x.detach().cpu().numpy() x = np.clip(x, -1, 1) @@ -139,21 +139,23 @@ def main(args: argparse.Namespace): image = Image.fromarray(x) image.save(args.rec_path.replace('mp4', 'jpg')) else: - custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path) + custom_to_video(video_recon[0], fps=args.fps, output_file=args.rec_path) + if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--video-path', type=str, default='') - parser.add_argument('--rec-path', type=str, default='') - parser.add_argument('--ckpt', type=str, default='results/pretrained') - parser.add_argument('--sample-fps', type=int, default=30) + parser.add_argument('--video_path', type=str, default='') + parser.add_argument('--rec_path', type=str, default='') + parser.add_argument('--ae', type=str, default='') + parser.add_argument('--model_path', type=str, default='results/pretrained') + parser.add_argument('--fps', type=int, default=30) parser.add_argument('--resolution', type=int, default=336) - parser.add_argument('--crop-size', type=int, default=None) - parser.add_argument('--num-frames', type=int, default=100) - parser.add_argument('--sample-rate', type=int, default=1) + parser.add_argument('--crop_size', type=int, default=None) + parser.add_argument('--num_frames', type=int, default=100) + parser.add_argument('--sample_rate', type=int, default=1) parser.add_argument('--device', type=str, default="cuda") parser.add_argument('--tile_overlap_factor', type=float, default=0.25) parser.add_argument('--enable_tiling', action='store_true') - + args = parser.parse_args() main(args) diff --git a/examples/rec_video.py b/examples/rec_video.py deleted file mode 100644 index 5acbeb6fb..000000000 --- a/examples/rec_video.py +++ /dev/null @@ -1,120 +0,0 @@ -import random -import argparse -from typing import Optional - -import cv2 -import imageio -import numpy as np -import numpy.typing as npt -import torch -from decord import VideoReader, cpu -from torch.nn import functional as F -from pytorchvideo.transforms import ShortSideScale -from torchvision.transforms import Lambda, Compose -from torchvision.transforms._transforms_video import RandomCropVideo - -import sys -sys.path.append(".") -from opensora.models.ae import VQVAEModel - - -def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None: - height, width, channels = image_array[0].shape - fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore - video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height)) - - for image in image_array: - image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - video_writer.write(image_rgb) - - video_writer.release() - -def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None: - x = x.detach().cpu() - x = torch.clamp(x, -0.5, 0.5) - x = (x + 0.5) - x = x.permute(1, 2, 3, 0).numpy() # (C, T, H, W) -> (T, H, W, C) - x = (255*x).astype(np.uint8) - # array_to_video(x, fps=fps, output_file=output_file) - imageio.mimwrite(output_file, x, fps=fps, quality=9) - return - -def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor: - decord_vr = VideoReader(video_path, ctx=cpu(0)) - total_frames = len(decord_vr) - sample_frames_len = sample_rate * num_frames - - if total_frames > sample_frames_len: - s = random.randint(0, total_frames - sample_frames_len - 1) - e = s + sample_frames_len - num_frames = num_frames - else: - s = 0 - e = total_frames - num_frames = int(total_frames / sample_frames_len * num_frames) - print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path, - total_frames) - - - frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int) - video_data = decord_vr.get_batch(frame_id_list).asnumpy() - video_data = torch.from_numpy(video_data) - video_data = video_data.permute(3, 0, 1, 2) # (T, H, W, C) -> (C, T, H, W) - return video_data - -def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor: - - transform = Compose( - [ - # UniformTemporalSubsample(num_frames), - Lambda(lambda x: ((x / 255.0) - 0.5)), - # NormalizeVideo(mean=OPENAI_DATASET_MEAN, std=OPENAI_DATASET_STD), - ShortSideScale(size=short_size), - RandomCropVideo(size=crop_size) if crop_size is not None else Lambda(lambda x: x), - # RandomHorizontalFlipVideo(p=0.5), - ] - ) - - video_outputs = transform(video_data) - video_outputs = torch.unsqueeze(video_outputs, 0) - - return video_outputs - - -def main(args: argparse.Namespace): - video_path = args.video_path - num_frames = args.num_frames - resolution = args.resolution - crop_size = args.crop_size - sample_fps = args.sample_fps - sample_rate = args.sample_rate - device = torch.device('cuda') - if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']: - vqvae = VQVAEModel.download_and_load_model(args.ckpt) - else: - vqvae = VQVAEModel.load_from_checkpoint(args.ckpt) - vqvae.eval() - vqvae = vqvae.to(device) - - with torch.no_grad(): - x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) - x_vae = x_vae.to(device) - encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) - video_recon = vqvae.decode(encodings) - - # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4') - custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--video-path', type=str, default='') - parser.add_argument('--rec-path', type=str, default='') - parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4') - parser.add_argument('--sample-fps', type=int, default=30) - parser.add_argument('--resolution', type=int, default=336) - parser.add_argument('--crop-size', type=int, default=None) - parser.add_argument('--num-frames', type=int, default=100) - parser.add_argument('--sample-rate', type=int, default=1) - args = parser.parse_args() - main(args) diff --git a/examples/rec_video_ae.py b/examples/rec_video_ae.py deleted file mode 100644 index 5acbeb6fb..000000000 --- a/examples/rec_video_ae.py +++ /dev/null @@ -1,120 +0,0 @@ -import random -import argparse -from typing import Optional - -import cv2 -import imageio -import numpy as np -import numpy.typing as npt -import torch -from decord import VideoReader, cpu -from torch.nn import functional as F -from pytorchvideo.transforms import ShortSideScale -from torchvision.transforms import Lambda, Compose -from torchvision.transforms._transforms_video import RandomCropVideo - -import sys -sys.path.append(".") -from opensora.models.ae import VQVAEModel - - -def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None: - height, width, channels = image_array[0].shape - fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore - video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height)) - - for image in image_array: - image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) - video_writer.write(image_rgb) - - video_writer.release() - -def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None: - x = x.detach().cpu() - x = torch.clamp(x, -0.5, 0.5) - x = (x + 0.5) - x = x.permute(1, 2, 3, 0).numpy() # (C, T, H, W) -> (T, H, W, C) - x = (255*x).astype(np.uint8) - # array_to_video(x, fps=fps, output_file=output_file) - imageio.mimwrite(output_file, x, fps=fps, quality=9) - return - -def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor: - decord_vr = VideoReader(video_path, ctx=cpu(0)) - total_frames = len(decord_vr) - sample_frames_len = sample_rate * num_frames - - if total_frames > sample_frames_len: - s = random.randint(0, total_frames - sample_frames_len - 1) - e = s + sample_frames_len - num_frames = num_frames - else: - s = 0 - e = total_frames - num_frames = int(total_frames / sample_frames_len * num_frames) - print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path, - total_frames) - - - frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int) - video_data = decord_vr.get_batch(frame_id_list).asnumpy() - video_data = torch.from_numpy(video_data) - video_data = video_data.permute(3, 0, 1, 2) # (T, H, W, C) -> (C, T, H, W) - return video_data - -def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor: - - transform = Compose( - [ - # UniformTemporalSubsample(num_frames), - Lambda(lambda x: ((x / 255.0) - 0.5)), - # NormalizeVideo(mean=OPENAI_DATASET_MEAN, std=OPENAI_DATASET_STD), - ShortSideScale(size=short_size), - RandomCropVideo(size=crop_size) if crop_size is not None else Lambda(lambda x: x), - # RandomHorizontalFlipVideo(p=0.5), - ] - ) - - video_outputs = transform(video_data) - video_outputs = torch.unsqueeze(video_outputs, 0) - - return video_outputs - - -def main(args: argparse.Namespace): - video_path = args.video_path - num_frames = args.num_frames - resolution = args.resolution - crop_size = args.crop_size - sample_fps = args.sample_fps - sample_rate = args.sample_rate - device = torch.device('cuda') - if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']: - vqvae = VQVAEModel.download_and_load_model(args.ckpt) - else: - vqvae = VQVAEModel.load_from_checkpoint(args.ckpt) - vqvae.eval() - vqvae = vqvae.to(device) - - with torch.no_grad(): - x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) - x_vae = x_vae.to(device) - encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) - video_recon = vqvae.decode(encodings) - - # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4') - custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--video-path', type=str, default='') - parser.add_argument('--rec-path', type=str, default='') - parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4') - parser.add_argument('--sample-fps', type=int, default=30) - parser.add_argument('--resolution', type=int, default=336) - parser.add_argument('--crop-size', type=int, default=None) - parser.add_argument('--num-frames', type=int, default=100) - parser.add_argument('--sample-rate', type=int, default=1) - args = parser.parse_args() - main(args) diff --git a/opensora/models/diffusion/latte/modeling_latte.py b/opensora/models/diffusion/latte/modeling_latte.py index 3993dae81..996762afd 100644 --- a/opensora/models/diffusion/latte/modeling_latte.py +++ b/opensora/models/diffusion/latte/modeling_latte.py @@ -8,8 +8,7 @@ from einops import rearrange, repeat from typing import Any, Dict, Optional, Tuple from diffusers.models import Transformer2DModel from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate -from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid, ImagePositionalEmbeddings, CaptionProjection -# from diffusers.models.embeddings import PatchEmbed, CombinedTimestepSizeEmbeddings +from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid, ImagePositionalEmbeddings from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear @@ -19,7 +18,8 @@ import torch.nn.functional as F from torch import nn from opensora.models.diffusion.utils.pos_embed import get_1d_sincos_pos_embed -from .modules import PatchEmbed, BasicTransformerBlock, BasicTransformerBlock_, AdaLayerNormSingle, Transformer3DModelOutput +from .modules import PatchEmbed, BasicTransformerBlock, BasicTransformerBlock_, AdaLayerNormSingle, \ + Transformer3DModelOutput, CaptionProjection class Latte(ModelMixin, ConfigMixin): diff --git a/opensora/models/diffusion/latte/modules.py b/opensora/models/diffusion/latte/modules.py index 786045dc7..81d8e5fde 100644 --- a/opensora/models/diffusion/latte/modules.py +++ b/opensora/models/diffusion/latte/modules.py @@ -1,28 +1,15 @@ from importlib import import_module import numpy as np -import torch - -import os -import json - -from dataclasses import dataclass -from einops import rearrange, repeat from typing import Any, Dict, Optional, Tuple, Callable -from diffusers.models import Transformer2DModel from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, is_xformers_available -from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid, ImagePositionalEmbeddings, CaptionProjection, \ - PatchEmbed, CombinedTimestepSizeEmbeddings -from diffusers.configuration_utils import ConfigMixin, register_to_config -from diffusers.models.modeling_utils import ModelMixin -from diffusers.models.attention import BasicTransformerBlock from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear import torch import torch.nn.functional as F from torch import nn from diffusers.utils.torch_utils import maybe_allow_in_graph -from diffusers.models.embeddings import SinusoidalPositionalEmbedding +from diffusers.models.embeddings import SinusoidalPositionalEmbedding, TimestepEmbedding, Timesteps from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormZero from diffusers.models.attention_processor import SpatialNorm, LORA_ATTENTION_PROCESSORS, \ CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0, \ @@ -33,8 +20,6 @@ from diffusers.models.activations import GEGLU, GELU, ApproximateGELU from dataclasses import dataclass -from torch import nn - from opensora.models.diffusion.utils.pos_embed import get_2d_sincos_pos_embed if is_xformers_available(): @@ -44,6 +29,80 @@ else: xformers = None +class CombinedTimestepSizeEmbeddings(nn.Module): + """ + For PixArt-Alpha. + + Reference: + https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 + """ + + def __init__(self, embedding_dim, size_emb_dim, use_additional_conditions: bool = False): + super().__init__() + + self.outdim = size_emb_dim + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + self.use_additional_conditions = use_additional_conditions + if use_additional_conditions: + self.use_additional_conditions = True + self.additional_condition_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.resolution_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) + self.aspect_ratio_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=size_emb_dim) + + def apply_condition(self, size: torch.Tensor, batch_size: int, embedder: nn.Module): + if size.ndim == 1: + size = size[:, None] + + if size.shape[0] != batch_size: + size = size.repeat(batch_size // size.shape[0], 1) + if size.shape[0] != batch_size: + raise ValueError(f"`batch_size` should be {size.shape[0]} but found {batch_size}.") + + current_batch_size, dims = size.shape[0], size.shape[1] + size = size.reshape(-1) + size_freq = self.additional_condition_proj(size).to(size.dtype) + + size_emb = embedder(size_freq) + size_emb = size_emb.reshape(current_batch_size, dims * self.outdim) + return size_emb + + def forward(self, timestep, resolution, aspect_ratio, batch_size, hidden_dtype): + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) + + if self.use_additional_conditions: + resolution = self.apply_condition(resolution, batch_size=batch_size, embedder=self.resolution_embedder) + aspect_ratio = self.apply_condition( + aspect_ratio, batch_size=batch_size, embedder=self.aspect_ratio_embedder + ) + conditioning = timesteps_emb + torch.cat([resolution, aspect_ratio], dim=1) + else: + conditioning = timesteps_emb + + return conditioning + +class CaptionProjection(nn.Module): + """ + Projects caption embeddings. Also handles dropout for classifier-free guidance. + + Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py + """ + + def __init__(self, in_features, hidden_size, num_tokens=120): + super().__init__() + self.linear_1 = nn.Linear(in_features=in_features, out_features=hidden_size, bias=True) + self.act_1 = nn.GELU(approximate="tanh") + self.linear_2 = nn.Linear(in_features=hidden_size, out_features=hidden_size, bias=True) + self.register_buffer("y_embedding", nn.Parameter(torch.randn(num_tokens, in_features) / in_features**0.5)) + + def forward(self, caption, force_drop_ids=None): + hidden_states = self.linear_1(caption) + hidden_states = self.act_1(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + class PatchEmbed(nn.Module): """2D Image to Patch Embedding""" diff --git a/opensora/serve/gradio_utils.py b/opensora/serve/gradio_utils.py index a7a7ab4c8..1935fe9e9 100644 --- a/opensora/serve/gradio_utils.py +++ b/opensora/serve/gradio_utils.py @@ -23,7 +23,7 @@ DESCRIPTION = """ ## If Open-Sora-Plan is helpful, please help to ✨ the [Github Repo](https://github.com/PKU-YuanGroup/Open-Sora-Plan) and recommend it to your friends 😊' #### [Open-Sora-Plan v1.0.0](https://github.com/PKU-YuanGroup/Open-Sora-Plan) is a transformer-based text-to-video diffusion system trained on text embeddings from T5. #### This demo is only trained on 40k videos, when creating videos, please be aware that it has the potential to generate harmful videos. For more details read our [report](). -#### Image generation is typically 50 steps, video generation maybe 250 steps will yield good results, but this may take 2-3 minutes. +#### Image generation is typically 50 steps, video generation maybe 150 steps will yield good results, but this may take 3-4 minutes. #### Feel free to enjoy the examples. #### English prompts ONLY; 提瀺词仅限英文 #### diff --git a/opensora/serve/gradio_web_server.py b/opensora/serve/gradio_web_server.py index 479b1140b..99dbeff0b 100644 --- a/opensora/serve/gradio_web_server.py +++ b/opensora/serve/gradio_web_server.py @@ -63,7 +63,8 @@ if __name__ == '__main__': # Load model: transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, torch_dtype=torch.float16, cache_dir='cache_dir').to(device) - vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir').to(device, dtype=torch.float16) + vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir').to(device) + vae = vae.half() vae.vae.enable_tiling() image_size = int(args.version.split('x')[1]) latent_size = (image_size // ae_stride_config[args.ae][1], image_size // ae_stride_config[args.ae][2]) @@ -82,7 +83,7 @@ if __name__ == '__main__': text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler, - transformer=transformer_model).to(device=device) + transformer=transformer_model).to(device) demo = gr.Interface( diff --git a/scripts/causalvae/reconstruction.sh b/scripts/causalvae/reconstruction.sh new file mode 100644 index 000000000..6be8f66b7 --- /dev/null +++ b/scripts/causalvae/reconstruction.sh @@ -0,0 +1,11 @@ +CUDA_VISIBLE_DEVICES=0 python examples/rec_imvi_vae.py \ + --model_path LanguageBind/Open-Sora-Plan-v1.0.0 \ + --video_path test.mp4 \ + --rec_path rec.mp4 \ + --device cuda \ + --sample_rate 1 \ + --num_frames 65 \ + --resolution 512 \ + --crop_size 512 \ + --ae CausalVAEModel_4x8x8 \ + --enable_tiling \ No newline at end of file diff --git a/scripts/text_condition/sample_video.sh b/scripts/text_condition/sample_video.sh index 81276d45c..9eb46ab51 100644 --- a/scripts/text_condition/sample_video.sh +++ b/scripts/text_condition/sample_video.sh @@ -1,10 +1,10 @@ -CUDA_VISIBLE_DEVICES=7 python opensora/sample/sample_t2v.py \ +CUDA_VISIBLE_DEVICES=0 python opensora/sample/sample_t2v.py \ --model_path LanguageBind/Open-Sora-Plan-v1.0.0 \ --text_encoder_name DeepFloyd/t5-v1_1-xxl \ - --text_prompt examples/prompt_list_7.txt \ + --text_prompt examples/prompt_list_0.txt \ --ae CausalVAEModel_4x8x8 \ --version 65x512x512 \ - --save_img_path "./sample_videos/prompt_list_7" \ + --save_img_path "./sample_videos/prompt_list_0" \ --fps 24 \ --guidance_scale 7.5 \ --num_sampling_steps 250 \