clean

Former-commit-id: a7f2c1c8b2243587cedd6e152e64298842287095

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)

examples/rec_video.py

@@ -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)

examples/rec_video_ae.py

@@ -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)

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):

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"""
 

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; 提示词仅限英文
 ####

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(

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

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 \