Merge pull request #91 from RuslanPeresy/typecheck

[refactor] add static type checking

Former-commit-id: 969dfa6dedaeb64cf776d977b466b14568cc4b1c

.gitignore

@@ -9,3 +9,4 @@ vae
 build/
 opensora.egg-info/
 wandb/
+.idea

README.md

@@ -164,6 +164,12 @@ conda activate opensora
 pip install -e .
 ```
 
+* Install optional requirements such as static type checking:
+
+```
+pip install -e '.[dev]'
+```
+
 ## 🗝️ Usage
 
 ### Datasets

examples/rec_video.py

@@ -4,19 +4,20 @@ sys.path.append(".")
 import random
 import cv2
 import numpy as np
+import numpy.typing as npt
 import torch
 from decord import VideoReader, cpu
 from pytorchvideo.transforms import ShortSideScale
-from torchvision.io import read_video
 from torchvision.transforms import Lambda, Compose
 from torchvision.transforms._transforms_video import RandomCropVideo
 from torch.nn import functional as F
+from typing import Optional
 from opensora.models.ae import VQVAEModel
 import argparse
 
-def array_to_video(image_array, fps=30.0, output_file='output_video.mp4'):
+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')
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')    # type: ignore
     video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
 
     for image in image_array:
@@ -25,7 +26,7 @@ def array_to_video(image_array, fps=30.0, output_file='output_video.mp4'):
 
     video_writer.release()
 
-def custom_to_video(x, fps=2.0, output_file='output_video.mp4'):
+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)
@@ -34,7 +35,7 @@ def custom_to_video(x, fps=2.0, output_file='output_video.mp4'):
     array_to_video(x, fps=fps, output_file=output_file)
     return
 
-def read_video(video_path, num_frames, sample_rate):
+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
@@ -57,7 +58,7 @@ def read_video(video_path, num_frames, sample_rate):
     video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
     return video_data
 
-def preprocess(video_data, short_size=128, crop_size=None):
+def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor:
 
     transform = Compose(
         [
@@ -76,7 +77,7 @@ def preprocess(video_data, short_size=128, crop_size=None):
     return video_outputs
 
 
-def main(args):
+def main(args: argparse.Namespace):
     video_path = args.video_path
     num_frames = args.num_frames
     resolution = args.resolution

opensora/models/ae/videobase/modeling_videobase.py

@@ -1,3 +1,4 @@
+import torch
 from torch import nn
 
 """
@@ -17,8 +18,8 @@ class VideoBaseAE(nn.Module):
     def download_and_load_model(cls, model_name, cache_dir=None):
         pass
     
-    def encode(self, x, *args, **kwargs):
+    def encode(self, x: torch.Tensor, *args, **kwargs):
         pass
 
-    def decode(self, encoding, *args, **kwargs):
+    def decode(self, encoding: torch.Tensor, *args, **kwargs):
         pass

opensora/models/ae/videobase/vqvae/modeling_vqvae.py

@@ -1,12 +1,13 @@
 from ..modeling_videobase import VideoBaseAE
 import torch
-from torch import nn
+from torch import nn, Tensor
 import numpy as np
 import torch.distributed as dist
 import torch.nn.functional as F
 import math
 import os
 import json
+from typing import Tuple, Dict, Union
 from .configuration_vqvae import VQVAEConfiguration
 
 
@@ -730,15 +731,15 @@ class VQVAEModel(VideoBaseAE):
         recon_loss = F.mse_loss(x_recon, x) / 0.06
         return recon_loss, x_recon, vq_output
 
-    def encode(self, x, include_embeddings=False):
+    def encode(self, x: Tensor, include_embeddings: bool = False) -> Union[Tuple[Tensor, Tensor], Tensor]:
         h = self.pre_vq_conv(self.encoder(x))
-        vq_output = self.codebook(h)
+        vq_output: Dict[str, Tensor] = self.codebook(h)
         if include_embeddings:
             return vq_output["encodings"], vq_output["embeddings"]
         else:
             return vq_output["encodings"]
 
-    def decode(self, encodings):
+    def decode(self, encodings: Tensor) -> Tensor:
         h = F.embedding(encodings, self.codebook.embeddings)
         h = self.post_vq_conv(shift_dim(h, -1, 1))
         return self.decoder(h)

pyproject.toml

@@ -23,6 +23,11 @@ dependencies = [
     "diffusers==0.24.0", "scikit-video==1.1.11", "imageio-ffmpeg==0.4.9"
 ]
 
+[project.optional-dependencies]
+dev = [
+    "mypy==1.8.0"
+]
+
 
 [project.urls]
 "Homepage" = "https://github.com/PKU-YuanGroup/Open-Sora-Plan"
@@ -33,3 +38,11 @@ exclude = ["assets*", "docker*", "docs", "scripts*"]
 
 [tool.wheel]
 exclude = ["assets*", "docker*", "docs", "scripts*"]
+
+[tool.mypy]
+warn_return_any = true
+warn_unused_configs = true
+ignore_missing_imports = true
+disallow_untyped_calls = true
+check_untyped_defs = true
+no_implicit_optional = true