Upload 49 files

src/videogen_hub/__init__.py

@@ -0,0 +1,12 @@
+import os
+
+from ._version import __version__
+MODEL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "checkpoints"))
+if os.environ.get("VIDEO_MODEL_PATH"):
+    MODEL_PATH = os.environ.get("VIDEO_MODEL_PATH")
+
+# (cogVideo) Set the SAT_HOME env variable to MODEL_PATH if not set
+if not os.environ.get("SAT_HOME"):
+    os.environ["SAT_HOME"] = MODEL_PATH
+
+from .infermodels import load, get_model, load_model

src/videogen_hub/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.4a0"

src/videogen_hub/depend/icetk/__init__.py

@@ -0,0 +1,5 @@
+from .ice_tokenizer import IceTokenizer
+
+icetk = IceTokenizer()
+
+__all__ = ['icetk']

src/videogen_hub/depend/icetk/ice_tokenizer.py

@@ -0,0 +1,116 @@
+# -*- encoding: utf-8 -*-
+
+import os
+import sys
+import math
+import random
+from typing import List, Tuple, Union
+
+import torch
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms.functional import pil_to_tensor
+
+from .text_tokenizer import TextTokenizer
+from .image_tokenizer import ImageTokenizer
+from .utils import auto_create
+
+class IceTokenizer:
+    def __init__(self, path='~/.icetk_models', device='cuda', fp16=True):
+        self.configure(path, device, fp16)
+        
+    def configure(self, path=None, device=None, fp16=None):
+        if path is not None:
+            self.path = os.path.expanduser(path)
+        if device is not None:
+            self.device = device
+        if fp16 is not None:
+            self.fp16 = fp16
+            
+    @property
+    def text_tokenizer(self):
+        if not hasattr(self, '_text_tokenizer'):
+            fp = os.path.join(self.path, 'ice_text.model')
+            auto_create(fp)
+            self._text_tokenizer = TextTokenizer(fp)
+        return self._text_tokenizer
+    
+    @property
+    def image_tokenizer(self):
+        if not hasattr(self, '_image_tokenizer'):
+            fp = os.path.join(self.path, 'ice_image.pt')
+            auto_create(fp)
+            self._image_tokenizer = ImageTokenizer(fp, device=self.device, fp16=self.fp16)
+        return self._image_tokenizer
+    
+    @property
+    def num_image_tokens(self):
+        return 20000 # self.image_tokenizer.num_tokens # allow not load
+    
+    @property
+    def num_text_tokens(self):
+        return self.text_tokenizer.num_tokens
+    @property
+    def num_tokens(self):
+        return self.num_image_tokens + self.num_text_tokens
+        
+    def add_special_tokens(self, special_tokens: List[str]):
+        self.text_tokenizer.add_special_tokens(special_tokens)
+    
+    def encode(self, text=None, 
+               image_path=None, image_pil=None, image_torch=None, 
+               image_size: int=None, compress_rate=8, ignore_linebreak=True):
+        assert (text is None) + (image_path is None) + (image_pil is None) + (image_torch is None) == 3
+        assert int(compress_rate) in [4, 8, 16]
+        if text is not None:
+            if not ignore_linebreak:
+                text = text.replace('\n', '<n>')
+            tmp = self.text_tokenizer.encode(text)
+            return [x + self.num_image_tokens for x in tmp]
+        else:
+            need_norm_to_1 = False
+            if image_path is not None:
+                image_pil = Image.open(image_path)
+            if image_torch is None:
+                image_torch = pil_to_tensor(image_pil)
+                need_norm_to_1 = True
+            if image_size is not None:
+                # for speed in large-scale preprocessing, set this to None and transform in Dataloader.
+                # TODO: test speed
+                tr = transforms.Compose([
+                    transforms.Resize(image_size),
+                    transforms.CenterCrop(image_size),
+                ])
+                image_torch = tr(image_torch)
+            image_torch = image_torch.to(self.image_tokenizer.device).float()
+            if need_norm_to_1:
+                image_torch /= 255.
+            return self.image_tokenizer.encode(image_torch, l=int(math.log2(compress_rate))-2)
+            
+
+    def decode(self, text_ids: List[int]=None, image_ids: Union[List[int], torch.LongTensor]=None, compress_rate=8):
+        assert (text_ids is None) + (image_ids is None) == 1
+        if text_ids is not None:
+            ids = [int(_id) - self.num_image_tokens for _id in text_ids]
+            return self.text_tokenizer.decode(ids).replace('<n>', '\n')
+        else:
+            return self.image_tokenizer.decode(image_ids, l=int(math.log2(compress_rate))-2)
+            
+    def tokenize(self, text):
+        return self.text_tokenizer.tokenize(text)
+
+    def __getitem__(self, x):
+        if isinstance(x, int):
+            if x < self.num_image_tokens:
+                return '<image_{}>'.format(x)
+            else:
+                return self.text_tokenizer.convert_id_to_token(x - self.num_image_tokens)
+        elif isinstance(x, str):
+            if x.startswith('<image_') and x.endswith('>') and x[7:-1].isdigit():
+                return int(x[7:-1])
+            else:
+                return self.text_tokenizer.convert_token_to_id(x) + self.num_image_tokens
+        else:
+            raise ValueError('The key should be str or int.')
+        
+        

src/videogen_hub/depend/icetk/image_tokenizer.py

@@ -0,0 +1,77 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   image_tokenizer.py
+@Time    :   2021/12/20 14:19:49
+@Author  :   Ming Ding 
+@Contact :   dm18@mails.tsinghua.edu.cn
+'''
+
+# here put the import lib
+import os
+import sys
+import math
+import random
+
+import torch
+import torch.nn.functional as F
+from torchvision import transforms 
+
+from .vqvae import load_default_HVQVAE, load_ckpt
+
+class ImageTokenizer(object):
+    def __init__(self,
+                model_path,
+                device='cuda',
+                fp16=True):
+        model = load_default_HVQVAE()
+        model = load_ckpt(model, model_path)
+        model = model.to(device)
+        model.eval()
+        
+        self.tr_normalize = transforms.Normalize(
+            [0.79093, 0.76271, 0.75340], 
+            [0.30379, 0.32279, 0.32800]
+            )
+
+        self.model = model
+        self.device = device
+        self.fp16 = fp16
+        self.num_tokens = model.quantize.n_embed
+        
+        if fp16:
+            model = model.half()
+
+    def __len__(self):
+        return self.num_tokens
+
+    def encode(self, image_torch, l=1):
+        '''Convert a batch of img to code
+        Args:
+            model: The tokenizer model.
+            img: [b, c, h, w]
+        '''
+        if len(image_torch.shape) == 3:
+            image_torch = image_torch.unsqueeze(0)
+        img = self.tr_normalize(image_torch).to(self.device)
+        if self.fp16:
+            img = img.half()
+        with torch.no_grad():
+            quant, diff, id = self.model.single_encode(img, l)
+        return id.view(img.shape[0], -1)
+
+    def decode(self, codes, l=1):
+        '''Convert a batch of code to imgs
+        Args:
+            codes : [b, h, w] or [b, h*w] or [h*w] LongTensor / list
+        '''
+        if isinstance(codes, list):
+            codes = torch.tensor(codes, dtype=torch.long, device=self.device)
+        if len(codes.shape) == 1:   
+            codes = codes.unsqueeze(0)
+        if len(codes.shape) == 2:
+            s = int(math.sqrt(len(codes.view(-1))) + 1e-5)
+            codes = codes.view(codes.shape[0], s, s)
+        with torch.no_grad():
+            out = self.model.single_decode_code(codes, l)
+            out = out * torch.tensor([0.30379, 0.32279, 0.32800], device=out.device).view(1, -1, 1, 1) + torch.tensor([0.79093, 0.76271, 0.75340], device=out.device).view(1, -1, 1, 1)
+        return out

src/videogen_hub/depend/icetk/sentencepiece_model_pb2.py

@@ -0,0 +1,722 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: sentencepiece_model.proto
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor.FileDescriptor(
+  name='sentencepiece_model.proto',
+  package='sentencepiece',
+  syntax='proto2',
+  serialized_options=b'H\003',
+  create_key=_descriptor._internal_create_key,
+  serialized_pb=b'\n\x19sentencepiece_model.proto\x12\rsentencepiece\"\xa1\n\n\x0bTrainerSpec\x12\r\n\x05input\x18\x01 \x03(\t\x12\x14\n\x0cinput_format\x18\x07 \x01(\t\x12\x14\n\x0cmodel_prefix\x18\x02 \x01(\t\x12\x41\n\nmodel_type\x18\x03 \x01(\x0e\x32$.sentencepiece.TrainerSpec.ModelType:\x07UNIGRAM\x12\x18\n\nvocab_size\x18\x04 \x01(\x05:\x04\x38\x30\x30\x30\x12\x17\n\x0f\x61\x63\x63\x65pt_language\x18\x05 \x03(\t\x12 \n\x15self_test_sample_size\x18\x06 \x01(\x05:\x01\x30\x12\"\n\x12\x63haracter_coverage\x18\n \x01(\x02:\x06\x30.9995\x12\x1e\n\x13input_sentence_size\x18\x0b \x01(\x04:\x01\x30\x12$\n\x16shuffle_input_sentence\x18\x13 \x01(\x08:\x04true\x12 \n\x14mining_sentence_size\x18\x0c \x01(\x05\x42\x02\x18\x01\x12\"\n\x16training_sentence_size\x18\r \x01(\x05\x42\x02\x18\x01\x12(\n\x17seed_sentencepiece_size\x18\x0e \x01(\x05:\x07\x31\x30\x30\x30\x30\x30\x30\x12\x1e\n\x10shrinking_factor\x18\x0f \x01(\x02:\x04\x30.75\x12!\n\x13max_sentence_length\x18\x12 \x01(\x05:\x04\x34\x31\x39\x32\x12\x17\n\x0bnum_threads\x18\x10 \x01(\x05:\x02\x31\x36\x12\x1d\n\x12num_sub_iterations\x18\x11 \x01(\x05:\x01\x32\x12$\n\x18max_sentencepiece_length\x18\x14 \x01(\x05:\x02\x31\x36\x12%\n\x17split_by_unicode_script\x18\x15 \x01(\x08:\x04true\x12\x1d\n\x0fsplit_by_number\x18\x17 \x01(\x08:\x04true\x12!\n\x13split_by_whitespace\x18\x16 \x01(\x08:\x04true\x12)\n\x1atreat_whitespace_as_suffix\x18\x18 \x01(\x08:\x05\x66\x61lse\x12\x1b\n\x0csplit_digits\x18\x19 \x01(\x08:\x05\x66\x61lse\x12\x17\n\x0f\x63ontrol_symbols\x18\x1e \x03(\t\x12\x1c\n\x14user_defined_symbols\x18\x1f \x03(\t\x12\x16\n\x0erequired_chars\x18$ \x01(\t\x12\x1c\n\rbyte_fallback\x18# \x01(\x08:\x05\x66\x61lse\x12+\n\x1dvocabulary_output_piece_score\x18  \x01(\x08:\x04true\x12\x1e\n\x10hard_vocab_limit\x18! \x01(\x08:\x04true\x12\x1c\n\ruse_all_vocab\x18\" \x01(\x08:\x05\x66\x61lse\x12\x11\n\x06unk_id\x18( \x01(\x05:\x01\x30\x12\x11\n\x06\x62os_id\x18) \x01(\x05:\x01\x31\x12\x11\n\x06\x65os_id\x18* \x01(\x05:\x01\x32\x12\x12\n\x06pad_id\x18+ \x01(\x05:\x02-1\x12\x18\n\tunk_piece\x18- \x01(\t:\x05<unk>\x12\x16\n\tbos_piece\x18. \x01(\t:\x03<s>\x12\x17\n\teos_piece\x18/ \x01(\t:\x04</s>\x12\x18\n\tpad_piece\x18\x30 \x01(\t:\x05<pad>\x12\x1a\n\x0bunk_surface\x18, \x01(\t:\x05 \xe2\x81\x87 \x12+\n\x1ctrain_extremely_large_corpus\x18\x31 \x01(\x08:\x05\x66\x61lse\"5\n\tModelType\x12\x0b\n\x07UNIGRAM\x10\x01\x12\x07\n\x03\x42PE\x10\x02\x12\x08\n\x04WORD\x10\x03\x12\x08\n\x04\x43HAR\x10\x04*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\"\xd1\x01\n\x0eNormalizerSpec\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x1c\n\x14precompiled_charsmap\x18\x02 \x01(\x0c\x12\x1e\n\x10\x61\x64\x64_dummy_prefix\x18\x03 \x01(\x08:\x04true\x12&\n\x18remove_extra_whitespaces\x18\x04 \x01(\x08:\x04true\x12 \n\x12\x65scape_whitespaces\x18\x05 \x01(\x08:\x04true\x12\x1e\n\x16normalization_rule_tsv\x18\x06 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\"y\n\x0cSelfTestData\x12\x33\n\x07samples\x18\x01 \x03(\x0b\x32\".sentencepiece.SelfTestData.Sample\x1a)\n\x06Sample\x12\r\n\x05input\x18\x01 \x01(\t\x12\x10\n\x08\x65xpected\x18\x02 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\"\xfe\x03\n\nModelProto\x12\x37\n\x06pieces\x18\x01 \x03(\x0b\x32\'.sentencepiece.ModelProto.SentencePiece\x12\x30\n\x0ctrainer_spec\x18\x02 \x01(\x0b\x32\x1a.sentencepiece.TrainerSpec\x12\x36\n\x0fnormalizer_spec\x18\x03 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x12\x33\n\x0eself_test_data\x18\x04 \x01(\x0b\x32\x1b.sentencepiece.SelfTestData\x12\x38\n\x11\x64\x65normalizer_spec\x18\x05 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x1a\xd2\x01\n\rSentencePiece\x12\r\n\x05piece\x18\x01 \x01(\t\x12\r\n\x05score\x18\x02 \x01(\x02\x12\x42\n\x04type\x18\x03 \x01(\x0e\x32,.sentencepiece.ModelProto.SentencePiece.Type:\x06NORMAL\"T\n\x04Type\x12\n\n\x06NORMAL\x10\x01\x12\x0b\n\x07UNKNOWN\x10\x02\x12\x0b\n\x07\x43ONTROL\x10\x03\x12\x10\n\x0cUSER_DEFINED\x10\x04\x12\x08\n\x04\x42YTE\x10\x06\x12\n\n\x06UNUSED\x10\x05*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\x42\x02H\x03'
+)
+
+
+
+_TRAINERSPEC_MODELTYPE = _descriptor.EnumDescriptor(
+  name='ModelType',
+  full_name='sentencepiece.TrainerSpec.ModelType',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='UNIGRAM', index=0, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='BPE', index=1, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='WORD', index=2, number=3,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='CHAR', index=3, number=4,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=1294,
+  serialized_end=1347,
+)
+_sym_db.RegisterEnumDescriptor(_TRAINERSPEC_MODELTYPE)
+
+_MODELPROTO_SENTENCEPIECE_TYPE = _descriptor.EnumDescriptor(
+  name='Type',
+  full_name='sentencepiece.ModelProto.SentencePiece.Type',
+  filename=None,
+  file=DESCRIPTOR,
+  create_key=_descriptor._internal_create_key,
+  values=[
+    _descriptor.EnumValueDescriptor(
+      name='NORMAL', index=0, number=1,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='UNKNOWN', index=1, number=2,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='CONTROL', index=2, number=3,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='USER_DEFINED', index=3, number=4,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='BYTE', index=4, number=6,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='UNUSED', index=5, number=5,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+  ],
+  containing_type=None,
+  serialized_options=None,
+  serialized_start=2100,
+  serialized_end=2184,
+)
+_sym_db.RegisterEnumDescriptor(_MODELPROTO_SENTENCEPIECE_TYPE)
+
+
+_TRAINERSPEC = _descriptor.Descriptor(
+  name='TrainerSpec',
+  full_name='sentencepiece.TrainerSpec',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='input', full_name='sentencepiece.TrainerSpec.input', index=0,
+      number=1, type=9, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='input_format', full_name='sentencepiece.TrainerSpec.input_format', index=1,
+      number=7, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='model_prefix', full_name='sentencepiece.TrainerSpec.model_prefix', index=2,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='model_type', full_name='sentencepiece.TrainerSpec.model_type', index=3,
+      number=3, type=14, cpp_type=8, label=1,
+      has_default_value=True, default_value=1,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='vocab_size', full_name='sentencepiece.TrainerSpec.vocab_size', index=4,
+      number=4, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=8000,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='accept_language', full_name='sentencepiece.TrainerSpec.accept_language', index=5,
+      number=5, type=9, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='self_test_sample_size', full_name='sentencepiece.TrainerSpec.self_test_sample_size', index=6,
+      number=6, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='character_coverage', full_name='sentencepiece.TrainerSpec.character_coverage', index=7,
+      number=10, type=2, cpp_type=6, label=1,
+      has_default_value=True, default_value=float(0.9995),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='input_sentence_size', full_name='sentencepiece.TrainerSpec.input_sentence_size', index=8,
+      number=11, type=4, cpp_type=4, label=1,
+      has_default_value=True, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='shuffle_input_sentence', full_name='sentencepiece.TrainerSpec.shuffle_input_sentence', index=9,
+      number=19, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='mining_sentence_size', full_name='sentencepiece.TrainerSpec.mining_sentence_size', index=10,
+      number=12, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=b'\030\001', file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='training_sentence_size', full_name='sentencepiece.TrainerSpec.training_sentence_size', index=11,
+      number=13, type=5, cpp_type=1, label=1,
+      has_default_value=False, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=b'\030\001', file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='seed_sentencepiece_size', full_name='sentencepiece.TrainerSpec.seed_sentencepiece_size', index=12,
+      number=14, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=1000000,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='shrinking_factor', full_name='sentencepiece.TrainerSpec.shrinking_factor', index=13,
+      number=15, type=2, cpp_type=6, label=1,
+      has_default_value=True, default_value=float(0.75),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='max_sentence_length', full_name='sentencepiece.TrainerSpec.max_sentence_length', index=14,
+      number=18, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=4192,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_threads', full_name='sentencepiece.TrainerSpec.num_threads', index=15,
+      number=16, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=16,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='num_sub_iterations', full_name='sentencepiece.TrainerSpec.num_sub_iterations', index=16,
+      number=17, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=2,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='max_sentencepiece_length', full_name='sentencepiece.TrainerSpec.max_sentencepiece_length', index=17,
+      number=20, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=16,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='split_by_unicode_script', full_name='sentencepiece.TrainerSpec.split_by_unicode_script', index=18,
+      number=21, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='split_by_number', full_name='sentencepiece.TrainerSpec.split_by_number', index=19,
+      number=23, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='split_by_whitespace', full_name='sentencepiece.TrainerSpec.split_by_whitespace', index=20,
+      number=22, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='treat_whitespace_as_suffix', full_name='sentencepiece.TrainerSpec.treat_whitespace_as_suffix', index=21,
+      number=24, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='split_digits', full_name='sentencepiece.TrainerSpec.split_digits', index=22,
+      number=25, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='control_symbols', full_name='sentencepiece.TrainerSpec.control_symbols', index=23,
+      number=30, type=9, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='user_defined_symbols', full_name='sentencepiece.TrainerSpec.user_defined_symbols', index=24,
+      number=31, type=9, cpp_type=9, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='required_chars', full_name='sentencepiece.TrainerSpec.required_chars', index=25,
+      number=36, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='byte_fallback', full_name='sentencepiece.TrainerSpec.byte_fallback', index=26,
+      number=35, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='vocabulary_output_piece_score', full_name='sentencepiece.TrainerSpec.vocabulary_output_piece_score', index=27,
+      number=32, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='hard_vocab_limit', full_name='sentencepiece.TrainerSpec.hard_vocab_limit', index=28,
+      number=33, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='use_all_vocab', full_name='sentencepiece.TrainerSpec.use_all_vocab', index=29,
+      number=34, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='unk_id', full_name='sentencepiece.TrainerSpec.unk_id', index=30,
+      number=40, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=0,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='bos_id', full_name='sentencepiece.TrainerSpec.bos_id', index=31,
+      number=41, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=1,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='eos_id', full_name='sentencepiece.TrainerSpec.eos_id', index=32,
+      number=42, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=2,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='pad_id', full_name='sentencepiece.TrainerSpec.pad_id', index=33,
+      number=43, type=5, cpp_type=1, label=1,
+      has_default_value=True, default_value=-1,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='unk_piece', full_name='sentencepiece.TrainerSpec.unk_piece', index=34,
+      number=45, type=9, cpp_type=9, label=1,
+      has_default_value=True, default_value=b"<unk>".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='bos_piece', full_name='sentencepiece.TrainerSpec.bos_piece', index=35,
+      number=46, type=9, cpp_type=9, label=1,
+      has_default_value=True, default_value=b"<s>".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='eos_piece', full_name='sentencepiece.TrainerSpec.eos_piece', index=36,
+      number=47, type=9, cpp_type=9, label=1,
+      has_default_value=True, default_value=b"</s>".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='pad_piece', full_name='sentencepiece.TrainerSpec.pad_piece', index=37,
+      number=48, type=9, cpp_type=9, label=1,
+      has_default_value=True, default_value=b"<pad>".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='unk_surface', full_name='sentencepiece.TrainerSpec.unk_surface', index=38,
+      number=44, type=9, cpp_type=9, label=1,
+      has_default_value=True, default_value=b" \342\201\207 ".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='train_extremely_large_corpus', full_name='sentencepiece.TrainerSpec.train_extremely_large_corpus', index=39,
+      number=49, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=False,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _TRAINERSPEC_MODELTYPE,
+  ],
+  serialized_options=None,
+  is_extendable=True,
+  syntax='proto2',
+  extension_ranges=[(200, 536870912), ],
+  oneofs=[
+  ],
+  serialized_start=45,
+  serialized_end=1358,
+)
+
+
+_NORMALIZERSPEC = _descriptor.Descriptor(
+  name='NormalizerSpec',
+  full_name='sentencepiece.NormalizerSpec',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='name', full_name='sentencepiece.NormalizerSpec.name', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='precompiled_charsmap', full_name='sentencepiece.NormalizerSpec.precompiled_charsmap', index=1,
+      number=2, type=12, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"",
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='add_dummy_prefix', full_name='sentencepiece.NormalizerSpec.add_dummy_prefix', index=2,
+      number=3, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='remove_extra_whitespaces', full_name='sentencepiece.NormalizerSpec.remove_extra_whitespaces', index=3,
+      number=4, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='escape_whitespaces', full_name='sentencepiece.NormalizerSpec.escape_whitespaces', index=4,
+      number=5, type=8, cpp_type=7, label=1,
+      has_default_value=True, default_value=True,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='normalization_rule_tsv', full_name='sentencepiece.NormalizerSpec.normalization_rule_tsv', index=5,
+      number=6, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=True,
+  syntax='proto2',
+  extension_ranges=[(200, 536870912), ],
+  oneofs=[
+  ],
+  serialized_start=1361,
+  serialized_end=1570,
+)
+
+
+_SELFTESTDATA_SAMPLE = _descriptor.Descriptor(
+  name='Sample',
+  full_name='sentencepiece.SelfTestData.Sample',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='input', full_name='sentencepiece.SelfTestData.Sample.input', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='expected', full_name='sentencepiece.SelfTestData.Sample.expected', index=1,
+      number=2, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=False,
+  syntax='proto2',
+  extension_ranges=[],
+  oneofs=[
+  ],
+  serialized_start=1641,
+  serialized_end=1682,
+)
+
+_SELFTESTDATA = _descriptor.Descriptor(
+  name='SelfTestData',
+  full_name='sentencepiece.SelfTestData',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='samples', full_name='sentencepiece.SelfTestData.samples', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[_SELFTESTDATA_SAMPLE, ],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=True,
+  syntax='proto2',
+  extension_ranges=[(200, 536870912), ],
+  oneofs=[
+  ],
+  serialized_start=1572,
+  serialized_end=1693,
+)
+
+
+_MODELPROTO_SENTENCEPIECE = _descriptor.Descriptor(
+  name='SentencePiece',
+  full_name='sentencepiece.ModelProto.SentencePiece',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='piece', full_name='sentencepiece.ModelProto.SentencePiece.piece', index=0,
+      number=1, type=9, cpp_type=9, label=1,
+      has_default_value=False, default_value=b"".decode('utf-8'),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='score', full_name='sentencepiece.ModelProto.SentencePiece.score', index=1,
+      number=2, type=2, cpp_type=6, label=1,
+      has_default_value=False, default_value=float(0),
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='type', full_name='sentencepiece.ModelProto.SentencePiece.type', index=2,
+      number=3, type=14, cpp_type=8, label=1,
+      has_default_value=True, default_value=1,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[],
+  enum_types=[
+    _MODELPROTO_SENTENCEPIECE_TYPE,
+  ],
+  serialized_options=None,
+  is_extendable=True,
+  syntax='proto2',
+  extension_ranges=[(200, 536870912), ],
+  oneofs=[
+  ],
+  serialized_start=1985,
+  serialized_end=2195,
+)
+
+_MODELPROTO = _descriptor.Descriptor(
+  name='ModelProto',
+  full_name='sentencepiece.ModelProto',
+  filename=None,
+  file=DESCRIPTOR,
+  containing_type=None,
+  create_key=_descriptor._internal_create_key,
+  fields=[
+    _descriptor.FieldDescriptor(
+      name='pieces', full_name='sentencepiece.ModelProto.pieces', index=0,
+      number=1, type=11, cpp_type=10, label=3,
+      has_default_value=False, default_value=[],
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='trainer_spec', full_name='sentencepiece.ModelProto.trainer_spec', index=1,
+      number=2, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='normalizer_spec', full_name='sentencepiece.ModelProto.normalizer_spec', index=2,
+      number=3, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='self_test_data', full_name='sentencepiece.ModelProto.self_test_data', index=3,
+      number=4, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+    _descriptor.FieldDescriptor(
+      name='denormalizer_spec', full_name='sentencepiece.ModelProto.denormalizer_spec', index=4,
+      number=5, type=11, cpp_type=10, label=1,
+      has_default_value=False, default_value=None,
+      message_type=None, enum_type=None, containing_type=None,
+      is_extension=False, extension_scope=None,
+      serialized_options=None, file=DESCRIPTOR,  create_key=_descriptor._internal_create_key),
+  ],
+  extensions=[
+  ],
+  nested_types=[_MODELPROTO_SENTENCEPIECE, ],
+  enum_types=[
+  ],
+  serialized_options=None,
+  is_extendable=True,
+  syntax='proto2',
+  extension_ranges=[(200, 536870912), ],
+  oneofs=[
+  ],
+  serialized_start=1696,
+  serialized_end=2206,
+)
+
+_TRAINERSPEC.fields_by_name['model_type'].enum_type = _TRAINERSPEC_MODELTYPE
+_TRAINERSPEC_MODELTYPE.containing_type = _TRAINERSPEC
+_SELFTESTDATA_SAMPLE.containing_type = _SELFTESTDATA
+_SELFTESTDATA.fields_by_name['samples'].message_type = _SELFTESTDATA_SAMPLE
+_MODELPROTO_SENTENCEPIECE.fields_by_name['type'].enum_type = _MODELPROTO_SENTENCEPIECE_TYPE
+_MODELPROTO_SENTENCEPIECE.containing_type = _MODELPROTO
+_MODELPROTO_SENTENCEPIECE_TYPE.containing_type = _MODELPROTO_SENTENCEPIECE
+_MODELPROTO.fields_by_name['pieces'].message_type = _MODELPROTO_SENTENCEPIECE
+_MODELPROTO.fields_by_name['trainer_spec'].message_type = _TRAINERSPEC
+_MODELPROTO.fields_by_name['normalizer_spec'].message_type = _NORMALIZERSPEC
+_MODELPROTO.fields_by_name['self_test_data'].message_type = _SELFTESTDATA
+_MODELPROTO.fields_by_name['denormalizer_spec'].message_type = _NORMALIZERSPEC
+DESCRIPTOR.message_types_by_name['TrainerSpec'] = _TRAINERSPEC
+DESCRIPTOR.message_types_by_name['NormalizerSpec'] = _NORMALIZERSPEC
+DESCRIPTOR.message_types_by_name['SelfTestData'] = _SELFTESTDATA
+DESCRIPTOR.message_types_by_name['ModelProto'] = _MODELPROTO
+_sym_db.RegisterFileDescriptor(DESCRIPTOR)
+
+TrainerSpec = _reflection.GeneratedProtocolMessageType('TrainerSpec', (_message.Message,), {
+  'DESCRIPTOR' : _TRAINERSPEC,
+  '__module__' : 'sentencepiece_model_pb2'
+  # @@protoc_insertion_point(class_scope:sentencepiece.TrainerSpec)
+  })
+_sym_db.RegisterMessage(TrainerSpec)
+
+NormalizerSpec = _reflection.GeneratedProtocolMessageType('NormalizerSpec', (_message.Message,), {
+  'DESCRIPTOR' : _NORMALIZERSPEC,
+  '__module__' : 'sentencepiece_model_pb2'
+  # @@protoc_insertion_point(class_scope:sentencepiece.NormalizerSpec)
+  })
+_sym_db.RegisterMessage(NormalizerSpec)
+
+SelfTestData = _reflection.GeneratedProtocolMessageType('SelfTestData', (_message.Message,), {
+
+  'Sample' : _reflection.GeneratedProtocolMessageType('Sample', (_message.Message,), {
+    'DESCRIPTOR' : _SELFTESTDATA_SAMPLE,
+    '__module__' : 'sentencepiece_model_pb2'
+    # @@protoc_insertion_point(class_scope:sentencepiece.SelfTestData.Sample)
+    })
+  ,
+  'DESCRIPTOR' : _SELFTESTDATA,
+  '__module__' : 'sentencepiece_model_pb2'
+  # @@protoc_insertion_point(class_scope:sentencepiece.SelfTestData)
+  })
+_sym_db.RegisterMessage(SelfTestData)
+_sym_db.RegisterMessage(SelfTestData.Sample)
+
+ModelProto = _reflection.GeneratedProtocolMessageType('ModelProto', (_message.Message,), {
+
+  'SentencePiece' : _reflection.GeneratedProtocolMessageType('SentencePiece', (_message.Message,), {
+    'DESCRIPTOR' : _MODELPROTO_SENTENCEPIECE,
+    '__module__' : 'sentencepiece_model_pb2'
+    # @@protoc_insertion_point(class_scope:sentencepiece.ModelProto.SentencePiece)
+    })
+  ,
+  'DESCRIPTOR' : _MODELPROTO,
+  '__module__' : 'sentencepiece_model_pb2'
+  # @@protoc_insertion_point(class_scope:sentencepiece.ModelProto)
+  })
+_sym_db.RegisterMessage(ModelProto)
+_sym_db.RegisterMessage(ModelProto.SentencePiece)
+
+
+DESCRIPTOR._options = None
+_TRAINERSPEC.fields_by_name['mining_sentence_size']._options = None
+_TRAINERSPEC.fields_by_name['training_sentence_size']._options = None
+# @@protoc_insertion_point(module_scope)

src/videogen_hub/depend/icetk/text_tokenizer.py

@@ -0,0 +1,77 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   text_tokenizer.py
+@Time    :   2021/12/20 01:26:12
+@Author  :   Ming Ding 
+@Contact :   dm18@mails.tsinghua.edu.cn
+'''
+
+# here put the import lib
+import os
+import sys
+import math
+import random
+from copy import copy
+from typing import List
+
+import sentencepiece as spm
+from . import sentencepiece_model_pb2 as model
+
+
+class TextTokenizer:
+    def __init__(self, model_path):
+        self.proto = model.ModelProto()
+        with open(model_path, 'rb') as fin:
+            proto_str = fin.read()
+            self.proto.ParseFromString(proto_str)
+        self.refresh()
+        
+    def refresh(self):
+        self.sp = spm.SentencePieceProcessor()
+        self.sp.Load(model_proto=self.proto.SerializeToString())
+        self.num_tokens = self.sp.vocab_size()
+
+    def add_special_tokens(self, tokens):
+        for token in tokens:
+            new_token = model.ModelProto().SentencePiece()
+            new_token.piece = token
+            new_token.score = 0
+            self.proto.pieces.append(new_token)
+        self.refresh()
+    
+    def discourage_tokens(self, tokens):
+        if isinstance(tokens, str): # single token
+            tokens = [tokens]
+        for token in tokens:
+            for piece in self.proto.pieces:    
+                if piece.piece == token:
+                    piece.score = -100
+        self.refresh()
+    
+    def discourage_ids(self, ids):
+        if isinstance(ids, int):
+            ids = [ids]
+        for idx in ids:
+            self.proto.pieces[idx].score = -100
+        self.refresh()
+
+    def encode(self, text):
+        return self.sp.EncodeAsIds(text)
+
+    def decode(self, ids: List[int]):
+        return self.sp.DecodeIds(ids)
+
+    def tokenize(self, text):
+        return self.sp.EncodeAsPieces(text)
+
+    def convert_tokens_to_ids(self, tokens):
+        return [self.sp.PieceToId(token) for token in tokens]
+
+    def convert_token_to_id(self, token):
+        return self.sp.PieceToId(token)
+
+    def convert_id_to_token(self, idx):
+        return self.sp.IdToPiece(idx)
+    
+    def __len__(self):
+        return self.num_tokens

src/videogen_hub/depend/icetk/utils.py

@@ -0,0 +1,46 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   utils.py
+@Time    :   2021/12/22 23:00:33
+@Author  :   Ming Ding 
+@Contact :   dm18@mails.tsinghua.edu.cn
+'''
+
+# here put the import lib
+import os
+import sys
+import math
+import random
+import requests
+
+from tqdm import tqdm
+import requests
+from filelock import FileLock
+
+def download_with_progress_bar(save_path, url):
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        with open(save_path, 'wb') as f:
+            pbar = tqdm(total=int(r.headers['Content-Length']), unit_scale=True)
+            for chunk in r.iter_content(chunk_size=32 * 1024):
+                if chunk:  # filter out keep-alive new chunks
+                    f.write(chunk)
+                    pbar.update(len(chunk))
+
+MODEL_ULRS = {
+    'ice_text.model': 'https://cloud.tsinghua.edu.cn/f/2c73ea6d3e7f4aed82ec/?dl=1',
+    'ice_image.pt': 'https://cloud.tsinghua.edu.cn/f/ae2cd37af814429d875d/?dl=1'
+}
+
+def auto_create(file_path):
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+    lock = FileLock(file_path + '.lock')
+    with lock:
+        if os.path.exists(file_path):
+            return False 
+        else:
+            url = MODEL_ULRS[os.path.basename(file_path)]
+            print(f'Downloading tokenizer models {url} into {file_path} ...')
+            download_with_progress_bar(file_path, url)
+            return True

src/videogen_hub/depend/icetk/vqvae/__init__.py

@@ -0,0 +1,5 @@
+from .vqvae_hierarchical import HVQVAE
+from .enc_dec import Encoder, Decoder, ResidualDownSample
+from .quantize import VectorQuantizeEMA
+
+from .api import load_default_HVQVAE, load_ckpt

src/videogen_hub/depend/icetk/vqvae/api.py

@@ -0,0 +1,93 @@
+import importlib
+import torch
+import json
+import math
+import os
+import numpy as np
+
+import torch.nn.functional as F
+
+def new_module(config):
+    '''in config:
+            "target": module type
+            "params": dict of params'''
+    if type(config) == str:
+        with open(config, 'r') as file:
+            config = json.load(file)
+    assert type(config) == dict
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    module, cls = config.get('target').rsplit(".", 1)
+    model = getattr(importlib.import_module(module, package=__package__), cls)(**config.get("params", dict()))
+    
+    return model
+
+def load_ckpt(model, path):
+    sd = torch.load(path, map_location="cpu")['module']
+    model.load_state_dict(sd, strict=False)
+    return model
+
+def load_default_HVQVAE():
+    config = {
+        "target": "..vqvae.HVQVAE",
+        "params": {
+            "levels": 3,
+            "embedding_dim": 256,
+            "codebook_scale": 1,
+            "down_sampler_configs": [
+                {
+                    "target": "..vqvae.ResidualDownSample",
+                    "params": {
+                        "in_channels": 256
+                    }
+                },
+                {
+                    "target": "..vqvae.ResidualDownSample",
+                    "params": {
+                        "in_channels": 256
+                    }
+                }
+            ],
+            "enc_config": {
+                    "target": "..vqvae.Encoder",
+                    "params": {
+                        "num_res_blocks": 2,
+                        "channels_mult": [1,2,4]
+                    }
+            },
+            "quantize_config": {
+                    "target": "..vqvae.VectorQuantizeEMA",
+                    "params": {
+                        "hidden_dim": 256,
+                        "embedding_dim": 256,
+                        "n_embed": 20000,
+                        "training_loc": False
+                    }
+            },
+            "dec_configs": [
+                {
+                    "target": "..vqvae.Decoder",
+                    "params": {
+                        "channels_mult": [1,1,1,2,4]
+                    }
+                },
+                {
+                    "target": "..vqvae.Decoder",
+                    "params": {
+                        "channels_mult": [1,1,2,4]
+                    }
+                },
+                {
+                    "target": "..vqvae.Decoder",
+                    "params": {
+                        "channels_mult": [1,2,4]
+                    }
+                }
+            ]
+        }
+    }
+    return new_module(config)
+
+
+if __name__ == '__main__':
+    pass

src/videogen_hub/depend/icetk/vqvae/enc_dec.py

@@ -0,0 +1,386 @@
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+import numpy as np
+
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+def Normalize(in_channels):
+    return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+class Upsample(nn.Module):
+    def __init__(self, 
+                    in_channels,
+                    with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if with_conv:
+            self.conv = nn.Conv2d(in_channels,
+                                    in_channels,
+                                    kernel_size=3,
+                                    stride=1,
+                                    padding=1)
+    
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2., mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+class DownSample(nn.Module):
+    def __init__(self,
+                    in_channels,
+                    with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if with_conv:
+            self.conv = nn.Conv2d(in_channels,
+                                    in_channels,
+                                    kernel_size=3,
+                                    stride=2,
+                                    padding=0)
+    
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = F.pad(x, pad, mode='constant', value=0)
+            x = self.conv(x)
+        else:
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+class ResidualDownSample(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+        self.pooling_down_sampler = DownSample(in_channels, with_conv=False)
+        self.conv_down_sampler = DownSample(in_channels, with_conv=True)
+
+    def forward(self, x):
+        return self.pooling_down_sampler(x) + self.conv_down_sampler(x)
+
+class ResnetBlock(nn.Module):
+    def __init__(self,
+                    in_channels,
+                    dropout,
+                    out_channels=None,
+                    conv_shortcut=False):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels,
+                                out_channels,
+                                kernel_size=3,
+                                stride=1,
+                                padding=1)
+
+        self.norm2 = Normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels,
+                                out_channels,
+                                kernel_size=3,
+                                stride=1,
+                                padding=1)
+        if in_channels != out_channels:
+            if conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(in_channels,
+                                                out_channels,
+                                                kernel_size=3,
+                                                stride=1,
+                                                padding=1)
+            else:
+                self.nin_shortcut = nn.Conv2d(in_channels,
+                                                out_channels,
+                                                kernel_size=1,
+                                                stride=1,
+                                                padding=0)
+                                            
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = nn.Conv2d(in_channels,
+                            in_channels,
+                            kernel_size=1,
+                            stride=1,
+                            padding=0)
+        self.k = nn.Conv2d(in_channels,
+                            in_channels,
+                            kernel_size=1,
+                            stride=1,
+                            padding=0)
+        self.v = nn.Conv2d(in_channels,
+                            in_channels,
+                            kernel_size=1,
+                            stride=1,
+                            padding=0)
+        self.proj_out = nn.Conv2d(in_channels,
+                                    in_channels,
+                                    kernel_size=1, 
+                                    stride=1,
+                                    padding=0)
+    
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        B, C, H, W = q.shape
+        q = q.reshape(B, C, -1)
+        q = q.permute(0, 2, 1) # (B, H*W, C)
+        k = k.reshape(B, C, -1) # (B, C, H*W)
+        w_ = torch.bmm(q, k) # (B, H*W, H*W)
+        w_ = w_ * C**(-0.5)
+        w_ = F.softmax(w_, dim=2)
+
+        v = v.reshape(B, C, -1) # (B, C, H*W)
+        w_ = w_.permute(0, 2, 1)
+        h_ = torch.bmm(v, w_)
+        h_ = h_.reshape(B, C, H, W)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_
+
+class Encoder(nn.Module):
+    def __init__(self,
+                    in_channels=3,
+                    out_channels=3,
+                    z_channels=256,
+                    channels=128,
+                    num_res_blocks=0,
+                    resolution=256,
+                    attn_resolutions=[16],
+                    resample_with_conv=True,
+                    channels_mult=(1,2,4,8),
+                    dropout=0.
+                    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.z_channels = z_channels
+        self.channels = channels
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+
+        self.conv_in = nn.Conv2d(in_channels,
+                                    channels,
+                                    kernel_size=3,
+                                    stride=1,
+                                    padding=1)
+
+        current_resolution = resolution
+        in_channels_mult = (1,) + tuple(channels_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = channels * in_channels_mult[i_level]
+            block_out = channels * channels_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                            out_channels=block_out,
+                                            dropout=dropout))
+                block_in = block_out
+                if current_resolution in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = DownSample(block_in,
+                                                resample_with_conv)
+                current_resolution = current_resolution // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                        out_channels=block_in,
+                                        dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                        out_channels=block_in,
+                                        dropout=dropout)
+        
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in,
+                                    z_channels,
+                                    kernel_size=3,
+                                    stride=1,
+                                    padding=1)
+
+    def test_forward(self, x):
+        # downsample
+        import pdb
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+            
+        return hs
+
+    def forward(self, x):
+        # downsample
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+
+        return h
+
+class Decoder(nn.Module):
+    def __init__(self,
+                    in_channels=3,
+                    out_channels=3,
+                    z_channels=256,
+                    channels=128,
+                    num_res_blocks=0,
+                    resolution=256,
+                    attn_resolutions=[16],
+                    channels_mult=(1,2,4,8),
+                    resample_with_conv=True,
+                    dropout=0.
+                    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.z_channels = z_channels
+        self.channels = channels
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        
+        in_channels_mult = (1,) + tuple(channels_mult)
+        block_in = channels * channels_mult[self.num_resolutions - 1]
+        current_resolution = resolution // 2**(self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, current_resolution, current_resolution)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels,
+                                    block_in,
+                                    kernel_size=3,
+                                    stride=1,
+                                    padding=1)
+        
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                        out_channels=block_in,
+                                        dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                        out_channels=block_in,
+                                        dropout=dropout)
+        
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = channels * channels_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                            out_channels=block_out,
+                                            dropout=dropout))
+                block_in = block_out
+                if current_resolution in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in,
+                                        resample_with_conv)
+                current_resolution = current_resolution * 2
+            self.up.insert(0, up)
+        
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in,
+                                    out_channels,
+                                    kernel_size=3,
+                                    stride=1,
+                                    padding=1)
+
+    def forward(self, z):
+        self.last_z_shape = z.shape
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+        
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+    def get_last_layer(self):
+        return self.conv_out.weight

src/videogen_hub/depend/icetk/vqvae/quantize.py

@@ -0,0 +1,156 @@
+import torch
+from torch import nn
+from torch import einsum
+from torch.nn import functional as F
+
+class VectorQuantize(nn.Module):
+    def __init__(self, 
+                    hidden_dim,
+                    embedding_dim,
+                    n_embed,
+                    commitment_cost=1):
+        super().__init__()
+        
+        self.hidden_dim = hidden_dim
+        self.embedding_dim = embedding_dim
+        self.n_embed = n_embed
+        self.commitment_cost = commitment_cost
+
+        self.proj = nn.Conv2d(hidden_dim, embedding_dim, 1)
+        self.embed = nn.Embedding(n_embed, embedding_dim)
+        self.embed.weight.data.uniform_(-1. / n_embed, 1. / n_embed)
+
+    def forward(self, z):
+        B, C, H, W = z.shape
+
+        z_e = self.proj(z)
+        z_e = z_e.permute(0, 2, 3, 1) # (B, H, W, C)
+        flatten = z_e.reshape(-1, self.embedding_dim)
+
+        dist = (
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ self.embed.weight.t()
+            + self.embed.weight.pow(2).sum(1, keepdim=True).t()
+        )
+        _, embed_ind = (-dist).max(1)
+        embed_ind = embed_ind.view(B, H, W)
+
+        z_q = self.embed_code(embed_ind)
+        diff = self.commitment_cost * (z_q.detach() - z_e).pow(2).mean() \
+                + (z_q - z_e.detach()).pow(2).mean()
+
+        z_q = z_e + (z_q - z_e).detach()
+        return z_q, diff, embed_ind
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.embed.weight)
+    
+
+class VectorQuantizeEMA(nn.Module):
+    def __init__(self,
+                    hidden_dim,
+                    embedding_dim,
+                    n_embed,
+                    commitment_cost=1,
+                    decay=0.99,
+                    eps=1e-5,
+                    pre_proj=True,
+                    training_loc=True):
+        super().__init__()
+        
+        self.hidden_dim = hidden_dim
+        self.embedding_dim = embedding_dim
+        self.n_embed = n_embed
+        self.commitment_cost = commitment_cost
+        self.training_loc = training_loc
+        
+        self.pre_proj = pre_proj
+        if self.pre_proj:
+            self.proj = nn.Conv2d(hidden_dim, embedding_dim, 1)
+        self.embed = nn.Embedding(n_embed, embedding_dim)
+        self.embed.weight.data.uniform_(-1. / n_embed, 1. / n_embed)
+        
+        self.register_buffer("cluster_size", torch.zeros(n_embed))
+        self.register_buffer("embed_avg", self.embed.weight.data.clone())
+        
+        self.decay = decay
+        self.eps = eps
+        
+    def forward(self, z):
+        B, C, H, W = z.shape
+        
+        if self.pre_proj:
+            z_e = self.proj(z)
+        else:
+            z_e = z
+        z_e = z_e.permute(0, 2, 3, 1) # (B, H, W, C)
+        flatten = z_e.reshape(-1, self.embedding_dim)
+
+        dist = (
+            flatten.pow(2).sum(1, keepdim=True)
+            - 2 * flatten @ self.embed.weight.t()
+            + self.embed.weight.pow(2).sum(1, keepdim=True).t()
+        )
+        _, embed_ind = (-dist).max(1)
+        embed_onehot = F.one_hot(embed_ind, self.n_embed).type(flatten.dtype)
+        embed_ind = embed_ind.view(B, H, W)
+
+        z_q = self.embed_code(embed_ind)
+        
+        diff = self.commitment_cost * (z_q.detach() - z_e).pow(2).mean()
+
+        z_q = z_e + (z_q - z_e).detach()
+        return z_q, diff, embed_ind
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.embed.weight)
+ 
+
+class GumbelQuantize(nn.Module):
+    def __init__(self,
+                    hidden_dim,
+                    embedding_dim,
+                    n_embed,
+                    commitment_cost=1,
+                    straight_through=True,
+                    kl_weight=5e-4,
+                    temp_init=1.,
+                    eps=1e-5):
+        super().__init__()
+        
+        self.hidden_dim = hidden_dim
+        self.embedding_dim = embedding_dim
+        self.n_embed = n_embed
+        self.commitment_cost = commitment_cost
+        
+        self.kl_weight = kl_weight
+        self.temperature = temp_init
+        self.eps = eps
+        
+        self.proj = nn.Conv2d(hidden_dim, n_embed, 1)
+        self.embed = nn.Embedding(n_embed, embedding_dim)
+        self.embed.weight.data.uniform_(-1. / n_embed, 1. / n_embed)
+        
+        self.straight_through = straight_through
+        
+    def forward(self, z, temp=None):
+        hard = self.straight_through if self.training else True
+        temp = self.temperature if temp is None else temp
+        
+        B, C, H, W = z.shape
+        
+        z_e = self.proj(z)
+        
+        soft_one_hot = F.gumbel_softmax(z_e, tau=temp, dim=1, hard=hard)
+        z_q = einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight)
+        
+        qy = F.softmax(z_e, dim=1)
+        diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.n_embed + self.eps), dim=1).mean()
+        
+        embed_ind = soft_one_hot.argmax(dim=1)
+        z_q = z_q.permute(0, 2, 3, 1)
+        return z_q, diff, embed_ind
+        
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.embed.weight)
+ 

src/videogen_hub/depend/icetk/vqvae/vqvae_hierarchical.py

@@ -0,0 +1,97 @@
+import torch
+from torch import nn
+import json
+import os
+
+from .api import new_module
+
+class HVQVAE(nn.Module):
+    def __init__(
+        self,
+        levels,
+        embedding_dim,
+        enc_config,
+        quantize_config,
+        down_sampler_configs,
+        dec_configs,
+        codebook_scale=1.
+    ):
+        super().__init__()
+        
+        self.levels = levels
+
+        self.enc = new_module(enc_config)
+            
+        self.decs = nn.ModuleList()
+        for i in range(levels):
+            self.decs.append(new_module(dec_configs[i]))
+            
+        self.quantize = new_module(quantize_config)
+        self.down_samplers = nn.ModuleList()
+        for i in range(levels-1):
+            self.down_samplers.append(new_module(down_sampler_configs[i]))
+        self.codebook_scale = codebook_scale
+            
+    def forward(self, input):        
+        quants, diffs, ids = self.encode(input)
+        dec_outputs = self.decode(quants[::-1])
+        
+        total_diff = diffs[0]
+        scale = 1.
+        for diff in diffs[1:]:
+            scale *= self.codebook_scale
+            total_diff = total_diff + diff * scale
+        return dec_outputs, total_diff
+
+    def encode(self, input):
+        enc_output = self.enc(input)
+        enc_outputs = [enc_output]
+        for l in range(self.levels-1):
+            enc_outputs.append(self.down_samplers[l](enc_outputs[-1]))
+       
+        quants, diffs, ids = [], [], []
+        for enc_output in enc_outputs:
+            quant, diff, id = self.quantize(enc_output)
+            quants.append(quant.permute(0, 3, 1, 2))
+            diffs.append(diff)
+            ids.append(id)
+            
+        return quants, diffs, ids
+        
+    def decode(self, quants):
+        dec_outputs = []
+        for l in range(self.levels-1, -1, -1):
+            dec_outputs.append(self.decs[l](quants[l]))
+            
+        return dec_outputs
+
+    def decode_code(self, codes):
+        quants = []
+        for l in range(self.levels):
+            quants.append(self.quantize.embed_code(codes[l]).permute(0, 3, 1, 2))
+        dec_outputs = self.decode(quants)
+
+        return dec_outputs
+
+    def single_encode(self, input, l):
+        assert l >= 0 and l <= 2
+        enc_output = self.enc(input)
+        for i in range(l):
+            enc_output = self.down_samplers[i](enc_output)
+        
+        quant, diff, id = self.quantize(enc_output)
+
+        return quant, diff, id
+    
+    def single_decode(self, quant, l):
+        assert l >= 0 and l <= 2
+        return self.decs[l](quant)
+    
+    def single_decode_code(self, code, l):
+        assert l >= 0 and l <= 2
+        quant = self.quantize.embed_code(code).permute(0, 3, 1, 2)
+        return self.decs[2-l](quant)
+    
+    def get_last_layer(self):
+        return self.decs[-1].get_last_layer()
+

src/videogen_hub/infermodels/__init__.py

@@ -0,0 +1,59 @@
+# ==========================================================
+# Text-to-Video Generation
+from .lavie import LaVie
+from .videocrafter import VideoCrafter2
+from .modelscope import ModelScope
+from .streamingt2v import StreamingT2V
+from .show_one import ShowOne
+from .opensora import OpenSora
+from .opensora_plan import OpenSoraPlan
+from .t2v_turbo import T2VTurbo
+from .opensora_12 import OpenSora12
+from .cogvideox import CogVideoX
+
+# from .cogvideo import CogVideo # Not supporting CogVideo ATM
+
+# ==========================================================
+# Image-to-Video Generation
+from .seine import SEINE
+from .consisti2v import ConsistI2V
+from .dynamicrafter import DynamiCrafter
+from .i2vgen_xl import I2VGenXL
+
+# ==========================================================
+
+import sys
+from functools import partial
+
+
+def get_model(model_name: str = None, init_with_default_params: bool = True):
+    """
+    Retrieves a model class or instance by its name.
+
+    Args:
+        model_name (str): Name of the model class. Triggers an error if the module name does not exist.
+        init_with_default_params (bool, optional): If True, returns an initialized model instance; otherwise, returns
+            the model class. Default is True. If set to True, be cautious of potential ``OutOfMemoryError`` with insufficient CUDA memory.
+
+    Returns:
+        model_class or model_instance: Depending on ``init_with_default_params``, either the model class or an instance of the model.
+
+    Examples::
+        initialized_model = infermodels.get_model(model_name='<Model>', init_with_default_params=True)
+
+        uninitialized_model = infermodels.get_model(model_name='<Model>', init_with_default_params=False)
+        initialized_model = uninitialized_model(device="cuda", <...>)
+    """
+
+    if not hasattr(sys.modules[__name__], model_name):
+        raise ValueError(f"No model named {model_name} found in infermodels.")
+
+    model_class = getattr(sys.modules[__name__], model_name)
+    if init_with_default_params:
+        model_instance = model_class()
+        return model_instance
+    return model_class
+
+
+load_model = partial(get_model, init_with_default_params=True)
+load = partial(get_model)

src/videogen_hub/infermodels/cogvideo.py

@@ -0,0 +1,54 @@
+class CogVideo:
+    def __init__(self, device="cuda"):
+        """
+        Initializes the CogVideo model with a specific device.
+
+        Args:
+            device (str, optional): The device to run the model on. Defaults to "cuda".
+        """
+
+        import argparse
+
+        # Manually creating an args object
+        self.args = argparse.Namespace(
+            generate_frame_num=5,
+            coglm_temperature2=0.89,
+            use_guidance_stage1=True,
+            use_guidance_stage2=False,  # Assuming this is not set
+            guidance_alpha=3.0,
+            stage_1=False,  # Assuming this is not set
+            stage_2=False,  # Assuming this is not set
+            both_stages=True,
+            parallel_size=1,
+            stage1_max_inference_batch_size=-1,
+            multi_gpu=False,  # Assuming this is not set
+            device=3,
+        )
+
+    def infer_one_video(
+        self,
+        prompt: str = None,
+        size: list = [320, 512],
+        seconds: int = 2,
+        fps: int = 8,
+        seed: int = 42,
+    ):
+        """
+        Generates a single video based on the provided prompt and parameters.
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+
+        from videogen_hub.pipelines.cogvideo.cogvideo_pipeline import pipeline
+
+        return pipeline(
+            self.args, raw_text=prompt, height=size[0], width=size[1], duration=seconds
+        )

src/videogen_hub/infermodels/cogvideox.py

@@ -0,0 +1,48 @@
+import torch
+
+class CogVideoX:
+    def __init__(self, weight="THUDM/CogVideoX-2b", device="cuda"):
+        """
+        Initializes the CogVideo model with a specific device.
+
+        Args:
+            device (str, optional): The device to run the model on. Defaults to "cuda".
+        """
+        from diffusers import CogVideoXPipeline
+
+        self.pipe = CogVideoXPipeline.from_pretrained(weight).to("cuda")
+
+    def infer_one_video(
+        self,
+        prompt: str = None,
+        size: list = [320, 512],
+        seconds: int = 2,
+        fps: int = 8,
+        seed: int = 42,
+    ):
+        """
+        Generates a single video based on the provided prompt and parameters.
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+
+        video = self.pipe(prompt=prompt, 
+                        guidance_scale=6,
+                        num_frames=seconds * fps, 
+                        #height=size[0],
+                        #width=size[1],
+                        num_inference_steps=50,
+                        generator=torch.manual_seed(seed)).frames[0]
+        from videogen_hub.utils import images_to_tensor
+        video = video[:-1] # drop the last frame
+        video = images_to_tensor(video) # parse it back to tensor (T, C, H, W)
+        
+        return video

src/videogen_hub/infermodels/consisti2v.py

@@ -0,0 +1,116 @@
+import os
+
+from PIL import Image
+from huggingface_hub import snapshot_download
+
+from videogen_hub import MODEL_PATH
+
+
+class ConsistI2V:
+    def __init__(self, device="cuda"):
+
+        class Args:
+            def __init__(self):
+                self.inference_config = "configs/inference/inference.yaml"
+                self.prompt = None
+                self.n_prompt = ""
+                self.seed = "random"
+                self.path_to_first_frame = None
+                self.prompt_config = "configs/prompts/default.yaml"
+                self.format = "mp4"
+                self.save_model = False
+                self.optional_args = []
+
+        self.args = Args()
+        model_path = os.path.join(MODEL_PATH, "TIGER-Lab", "ConsistI2V").replace("\\", "\\\\")
+        yaml_config = f"""
+            output_dir: "samples/inference"
+            output_name: "i2v"
+            pretrained_model_path: "{model_path}"
+            unet_path: null
+            unet_ckpt_prefix: "module."
+            pipeline_pretrained_path: null
+
+            sampling_kwargs:
+            height: 256
+            width: 256
+            n_frames: 16
+            steps: 50
+            ddim_eta: 0.0
+            guidance_scale_txt: 7.5
+            guidance_scale_img: 1.0
+            guidance_rescale: 0.0
+            num_videos_per_prompt: 1
+            frame_stride: 3
+
+            unet_additional_kwargs:
+            variant: null
+            n_temp_heads: 8
+            augment_temporal_attention: true
+            temp_pos_embedding: "rotary" # "rotary" or "sinusoidal"
+            first_frame_condition_mode: "concat"
+            use_frame_stride_condition: true
+            noise_sampling_method: "pyoco_mixed" # "vanilla" or "pyoco_mixed" or "pyoco_progressive"
+            noise_alpha: 1.0
+
+            noise_scheduler_kwargs:
+            beta_start: 0.00085
+            beta_end: 0.012
+            beta_schedule: "linear"
+            steps_offset: 1
+            clip_sample: false
+            rescale_betas_zero_snr: false     # true if using zero terminal snr
+            timestep_spacing:       "leading" # "trailing" if using zero terminal snr
+            prediction_type:        "epsilon" # "v_prediction" if using zero terminal snr
+
+            frameinit_kwargs:
+            enable: true
+            camera_motion: null
+            noise_level: 850
+            filter_params:
+                method: 'gaussian'
+                d_s: 0.25
+                d_t: 0.25
+        """
+
+        from omegaconf import OmegaConf
+
+        self.config = OmegaConf.create(yaml_config)
+        model_path = os.path.join(MODEL_PATH, "ConsistI2V").replace("\\", "\\\\")
+        snapshot_download("TIGER-Lab/ConsistI2V", local_dir=model_path)
+        from videogen_hub.pipelines.consisti2v.scripts.animate import main
+
+        self.pipeline = main
+
+    def infer_one_video(
+        self,
+        input_image: Image.Image,
+        prompt: str = None,
+        size: list = [320, 512],
+        seconds: int = 2,
+        fps: int = 8,
+        seed: int = 42,
+    ):
+        """
+        Generates a single video based on a textual prompt and first frame image, using either a provided image or an image path as the starting point. The output is a tensor representing the video.
+
+        Args:
+            input_image (PIL.Image.Image): The input image to use as the basis for video generation.
+            prompt (str, optional): The text prompt that guides the video generation. If not specified, the video generation will rely solely on the input image. Defaults to None.
+            size (list, optional): Specifies the resolution of the output video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The number of frames per second in the generated video. This determines how smooth the video appears. Defaults to 8.
+            seed (int, optional): A seed value for random number generation, ensuring reproducibility of the video generation process. Defaults to 42.
+
+        Returns:
+            torch.Tensor: A tensor representing the generated video, structured as (time, channel, height, width).
+        """
+
+        self.args.prompt = prompt
+        self.args.path_to_first_frame = input_image
+        self.args.seed = str(seed)
+        self.config.sampling_kwargs.height = size[0]
+        self.config.sampling_kwargs.width = size[1]
+        self.config.sampling_kwargs.n_frames = seconds * fps
+
+        return self.pipeline(self.args, self.config)

src/videogen_hub/infermodels/dynamicrafter.py

@@ -0,0 +1,104 @@
+import os
+
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from videogen_hub import MODEL_PATH
+
+
+class DynamiCrafter:
+    def __init__(self, version: str = "256"):
+        """
+        Initializes the DynamiCrafter model using the Doubiiu/DynamiCrafter_{version} checkpoint from the Hugging Face Hub.
+        and load them to "MODEL_DIR/dynamicrafter_{version}_v1"
+
+        Args:
+            version (str, optional): The resolution of the video to generate. Choose from '256', '512', or '1024'. Defaults to '256'.
+        """
+        from videogen_hub.pipelines.dynamicrafter.inference import DynamiCrafterPipeline
+
+        if version == "256":
+            (self.height, self.width) = 256, 256
+            self.fs = 3
+            self.model_path = hf_hub_download(
+                repo_id="Doubiiu/DynamiCrafter",
+                filename="model.ckpt",
+                local_dir=os.path.join(MODEL_PATH, "dynamicrafter_256_v1"),
+            )
+
+        elif version == "512":
+            (self.height, self.width) = 320, 512
+            self.fs = 24
+            self.model_path = hf_hub_download(
+                repo_id="Doubiiu/DynamiCrafter_512",
+                filename="model.ckpt",
+                local_dir=os.path.join(MODEL_PATH, "dynamicrafter_512_v1"),
+            )
+
+        elif version == "1024":
+            (self.height, self.width) = 576, 1024
+            self.fs = 10
+            self.model_path = hf_hub_download(
+                repo_id="Doubiiu/DynamiCrafter_1024",
+                filename="model.ckpt",
+                local_dir=os.path.join(MODEL_PATH, "dynamicrafter_1024_v1"),
+            )
+        else:
+            raise ValueError("Invalid input. Please enter 256, 512, or 1024.")
+
+        self.arg_list = [
+            "--ckpt_path",
+            self.model_path,
+            "--config",
+            f"src/videogen_hub/pipelines/dynamicrafter/configs/inference_{version}_v1.0.yaml",
+            "--n_samples",
+            "1",
+            "--bs",
+            "1",
+            "--height",
+            str(self.height),
+            "--width",
+            str(self.width),
+            "--text_input",
+            "--unconditional_guidance_scale",
+            "7.5",
+            "--ddim_steps",
+            "50",
+            "--ddim_eta",
+            "1.0",
+            "--video_length",
+            "16",
+            "--frame_stride",
+            str(self.fs),
+        ]
+
+        self.pipeline = DynamiCrafterPipeline(self.arg_list)
+
+    def infer_one_video(
+        self,
+        input_image: Image.Image,
+        prompt: str = None,
+        seconds: int = 2,
+        fps: int = 8,
+        seed: int = 42,
+    ):
+        """
+        Generates a single video based on a textual prompt and first frame image, using either a provided image or an image path as the starting point. The output is a tensor representing the video.
+
+        Args:
+            input_image (PIL.Image.Image): The input image to use as the basis for video generation.
+            prompt (str, optional): The text prompt that guides the video generation. If not specified, the video generation will rely solely on the input image. Defaults to None.
+            size (list, optional): Specifies the resolution of the output video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The number of frames per second in the generated video. This determines how smooth the video appears. Defaults to 8.
+            seed (int, optional): A seed value for random number generation, ensuring reproducibility of the video generation process. Defaults to 42.
+
+        Returns:
+            torch.Tensor: A tensor representing the generated video, structured as (time, channel, height, width).
+        """
+        self.pipeline.args.seed = seed
+        self.pipeline.args.text_input = prompt
+        self.pipeline.args.video_length = fps * seconds
+        video = self.pipeline.run_inference(input_image)
+
+        return video

src/videogen_hub/infermodels/i2vgen_xl.py

@@ -0,0 +1,57 @@
+import os
+from typing import Union
+import torch
+from huggingface_hub import snapshot_download, hf_hub_download
+from PIL import Image
+
+from videogen_hub import MODEL_PATH
+
+
+class I2VGenXL:
+    def __init__(self):
+        """
+        Initializes the I2VGenXL model using the ali-vilab/i2vgen-xl checkpoint from the Hugging Face Hub.
+
+        Args:
+            None
+        """
+
+        from diffusers import I2VGenXLPipeline
+        model_path = os.path.join(MODEL_PATH, "i2vgen-xl")
+        model_path = snapshot_download("ali-vilab/i2vgen-xl", local_dir=model_path, ignore_patterns=["*fp16*", "*png"])
+        self.pipeline = I2VGenXLPipeline.from_pretrained(
+            model_path, torch_dtype=torch.float16, variant="fp16"
+        )
+
+    def infer_one_video(
+            self,
+            input_image: Image.Image,
+            prompt: str = None,
+            size: list = [320, 512],
+            seconds: int = 2,
+            fps: int = 8,
+            seed: int = 42,
+    ):
+        """
+        Generates a single video based on a textual prompt and first frame image, using either a provided image or an image path as the starting point. The output is a tensor representing the video.
+
+        Args:
+            input_image (Image.Image): The input image path or tensor to use as the basis for video generation.
+            prompt (str, optional): The text prompt that guides the video generation. If not specified, the video generation will rely solely on the input image. Defaults to None.
+            size (list, optional): Specifies the resolution of the output video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The number of frames per second in the generated video. This determines how smooth the video appears. Defaults to 8.
+            seed (int, optional): A seed value for random number generation, ensuring reproducibility of the video generation process. Defaults to 42.
+
+        Returns:
+            torch.Tensor: A tensor representing the generated video, structured as (time, channel, height, width).
+        """
+        return self.pipeline(
+            prompt=prompt,
+            image=input_image,
+            height=size[0],
+            width=size[1],
+            target_fps=fps,
+            num_frames=seconds * fps,
+            generator=torch.manual_seed(seed),
+        )

src/videogen_hub/infermodels/lavie.py

@@ -0,0 +1,103 @@
+import os, sys
+import torch
+
+from videogen_hub import MODEL_PATH
+
+
+class LaVie():
+    def __init__(self, model_path=os.path.join(MODEL_PATH, "lavie"), device="cuda"):
+        """
+        1. Download all necessary models from huggingface.
+        2. Initializes the LaVie model with a specific model path and device.
+
+        Args:
+            model_path (str, optional): The path to the model checkpoints. Defaults to "MODEL_PATH/lavie".
+            device (str, optional): The device to run the model on. Defaults to "cuda".
+        """
+
+        # Put the source code imports here to avoid dependency version issues
+        from videogen_hub.pipelines.lavie.lavie_src.base.pipelines.pipeline_videogen import VideoGenPipeline
+        from videogen_hub.pipelines.lavie.lavie_src.base.download import find_model
+        from videogen_hub.pipelines.lavie.lavie_src.base.models.unet import UNet3DConditionModel
+        from diffusers.schedulers import DDPMScheduler
+        from diffusers.models import AutoencoderKL
+        from transformers import CLIPTokenizer, CLIPTextModel
+        from huggingface_hub import snapshot_download
+        from omegaconf import OmegaConf
+
+        snapshot_download(repo_id="Vchitect/LaVie", local_dir=model_path)
+        snapshot_download(repo_id="CompVis/stable-diffusion-v1-4", local_dir=os.path.join(model_path, "stable-diffusion-v1-4"))
+        snapshot_download(repo_id="stabilityai/stable-diffusion-x4-upscaler",
+                          local_dir=os.path.join(model_path, "stable-diffusion-x4-upscaler"))
+
+        torch.set_grad_enabled(False)
+        self.device = device
+
+        config = {
+            "model_config": {
+                "use_compile": False,
+                "use_fp16": True,
+                "run_time": 0,
+                "guidance_scale": 7.5,
+                "num_sampling_steps": 50
+            },
+            "scheduler_config": {
+                "sample_method": "ddpm",
+                "beta_start": 0.0001,
+                "beta_end": 0.02,
+                "beta_schedule": "linear"
+            }
+        }
+        self.config = OmegaConf.create(config)
+
+        sd_path = os.path.join(model_path, "stable-diffusion-v1-4")
+        unet = UNet3DConditionModel.from_pretrained_2d(sd_path, subfolder="unet").to(device, dtype=torch.float16)
+        state_dict = find_model(os.path.join(model_path, "lavie_base.pt"))
+        unet.load_state_dict(state_dict)
+
+        vae = AutoencoderKL.from_pretrained(sd_path, subfolder="vae", torch_dtype=torch.float16).to(device)
+        tokenizer_one = CLIPTokenizer.from_pretrained(sd_path, subfolder="tokenizer")
+        text_encoder_one = CLIPTextModel.from_pretrained(sd_path, subfolder="text_encoder",
+                                                         torch_dtype=torch.float16).to(device)  # huge
+
+        scheduler = DDPMScheduler.from_pretrained(sd_path,
+                                                  subfolder="scheduler",
+                                                  beta_start=self.config.scheduler_config.beta_start,
+                                                  beta_end=self.config.scheduler_config.beta_end,
+                                                  beta_schedule=self.config.scheduler_config.beta_schedule)
+
+        self.videogen_pipeline = VideoGenPipeline(vae=vae,
+                                                  text_encoder=text_encoder_one,
+                                                  tokenizer=tokenizer_one,
+                                                  scheduler=scheduler,
+                                                  unet=unet).to(device)
+        self.videogen_pipeline.enable_xformers_memory_efficient_attention()
+
+    def infer_one_video(self,
+                        prompt: str = None,
+                        size: list = [320, 512],
+                        seconds: int = 2,
+                        fps: int = 8,
+                        seed: int = 42):
+        """
+        Generates a single video based on the provided prompt and parameters.
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+        if seed is not None:
+            torch.manual_seed(seed)
+        videos = self.videogen_pipeline(prompt,
+                                        video_length=seconds * fps,
+                                        height=size[0],
+                                        width=size[1],
+                                        num_inference_steps=self.config.model_config.num_sampling_steps,
+                                        guidance_scale=self.config.model_config.guidance_scale).video
+        return videos[0]

src/videogen_hub/infermodels/modelscope.py

@@ -0,0 +1,62 @@
+import os
+
+import torch
+from huggingface_hub import snapshot_download
+
+from videogen_hub import MODEL_PATH
+
+
+class ModelScope:
+    def __init__(self, device="gpu"):
+        """
+        1. Download the pretrained model and put it inside checkpoints/modelscope
+        2. Create Pipeline
+        Note: it seems that the model needed from model_dir cannot support cpu
+        Args:
+            device: 'gpu' or 'cpu' the device to use the model
+        """
+        from modelscope.pipelines import pipeline
+        from modelscope.models import Model
+
+        model_dir = snapshot_download(
+            repo_id="ali-vilab/modelscope-damo-text-to-video-synthesis",
+            local_dir=os.path.join(MODEL_PATH, "modelscope"),
+            
+        )
+        model = Model.from_pretrained(model_dir)
+        self.pipeline = pipeline("text-to-video-synthesis", model=model, device=device)
+
+    def infer_one_video(
+            self, prompt: str = None, seconds: int = 2, fps: int = 8, seed: int = 42
+    ):
+        """
+        Generates a single video based on the provided prompt and parameters.
+        The generated video always has resolution 256x256
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+        from modelscope.outputs import OutputKeys
+        from decord import VideoReader
+        from decord import cpu, gpu
+        import io
+
+        torch.manual_seed(seed)
+        self.pipeline.model.config.model.model_args.max_frames = fps * seconds
+
+        test_text = {
+            "text": prompt,
+        }
+        output_video_path = self.pipeline(
+            test_text,
+        )[OutputKeys.OUTPUT_VIDEO]
+        result = io.BytesIO(output_video_path)
+        result = VideoReader(result, ctx=cpu(0))
+        result = torch.from_numpy(result.get_batch(range(len(result))).asnumpy())
+        return result

src/videogen_hub/infermodels/opensora.py

@@ -0,0 +1,134 @@
+import os
+
+from huggingface_hub import snapshot_download, hf_hub_download
+
+from videogen_hub import MODEL_PATH
+
+
+class OpenSora:
+    def __init__(self, device="gpu"):
+        """
+        1. Download the pretrained model and put it inside MODEL_PATH/modelscope
+        2. Create Pipeline
+        Note: it seems that the model needed from model_dir cannot support cpu
+        Args:
+            device: 'gpu' or 'cpu' the device to use the model
+        """
+
+        from mmengine import Config as mmengine_config
+        from videogen_hub.pipelines.opensora.scripts.inference import main
+
+        self.pipeline = main
+        self.config = {
+            # Basic video frame settings
+            "num_frames": 32,  # Total number of frames in a clip
+            "frame_interval": 3,  # Interval between frames
+            "fps": 24,  # Frames per second
+            "image_size": [480, 854],  # Resolution of each frame (height, width)
+            # Model configuration for multi-resolution and specific model parameters
+            "multi_resolution": "STDiT2",  # Multi-resolution model type
+            "model": {
+                "type": "STDiT2-XL/2",  # Model type and size
+                "from_pretrained": os.path.join(MODEL_PATH, "STDiT2-XL_2"),  # Path to pretrained checkpoint
+                "file_name": "model.safetensors",  # Name of the model file
+                "input_sq_size": 512,  # Input square size for the model
+                "qk_norm": True,  # Whether to normalize query-key in attention
+                "enable_flashattn": False,  # Enable flash attention mechanism, require flash_attn package
+                "enable_layernorm_kernel": False,  # Enable layer normalization in kernel, requires apex package
+            },
+            # Variational Autoencoder (VAE) specific settings
+            "vae": {
+                "type": "VideoAutoencoderKL",  # Type of the autoencoder
+                "from_pretrained": "stabilityai/sd-vae-ft-ema",  # Pretrained model from Hugging Face
+                "cache_dir": os.path.join(MODEL_PATH, "sd-vae-ft-ema"),  # Local cache directory for model weights
+                "micro_batch_size": 4,  # Batch size for processing
+            },
+            # Text encoder settings for embedding textual information
+            "text_encoder": {
+                "type": "t5",  # Text encoder model type
+                "from_pretrained": "DeepFloyd/t5-v1_1-xxl",  # Pretrained model
+                "cache_dir": os.path.join(MODEL_PATH, "t5-v1_1-xxl"),  # Cache directory
+                "model_max_length": 200,  # Max length of text inputs
+            },
+            # Scheduler settings for diffusion models
+            "scheduler": {
+                "type": "iddpm",  # Type of scheduler for the diffusion process
+                "num_sampling_steps": 50,  # Number of sampling steps in diffusion
+                "cfg_scale": 7.0,  # Scale for classifier-free guidance
+                "cfg_channel": 3,  # Number of channels for guidance
+            },
+            # Additional settings for processing and output
+            "dtype": "bf16",  # Data type for computation (bfloat16)
+            # "prompt_path": "./assets/texts/t2v_samples.txt",  # Path to text prompts
+            "prompt_path": None,  # Path to text prompts
+            "prompt": [
+                "A beautiful sunset over the city"
+            ],  # List of prompts for generation
+            "batch_size": 1,  # Batch size for generation
+            "seed": 42,  # Seed for random number generators
+            "save_dir": "./samples/samples/",  # Directory to save generated samples
+            "config": "sample.py",  # Path to this configuration file
+            "prompt_as_path": False,  # Treat the prompt as a file path (True/False)
+            "reference_path": None,  # Path to reference image/video for conditioning
+            "loop": 1,  # Number of times to loop the processing
+            "sample_name": None,  # Specific name for the generated sample
+            "num_sample": 1,  # Number of samples to generate
+        }
+        self.config = mmengine_config(self.config)
+
+        hf_hub_download(
+            repo_id="hpcai-tech/OpenSora-STDiT-v2-stage2",
+            filename="model.safetensors",
+            local_dir=self.config.model.from_pretrained,
+        )
+
+        hf_hub_download(
+            repo_id="stabilityai/sd-vae-ft-ema",
+            filename="diffusion_pytorch_model.safetensors",
+            local_dir=self.config.vae.cache_dir,
+        )
+
+        hf_hub_download(
+            repo_id="DeepFloyd/t5-v1_1-xxl",
+            filename="pytorch_model-00001-of-00002.bin",
+            local_dir=self.config.text_encoder.cache_dir,
+        )
+
+    def infer_one_video(
+        self,
+        prompt: str = None,
+        size: list = [320, 512],
+        seconds: int = 2,
+        fps: int = 8,
+        seed: int = 42,
+    ):
+        """
+        Generates a single video based on the provided prompt and parameters.
+        The generated video always has resolution 256x256
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+
+        self.config.num_frames = fps * seconds
+        self.config.fps = fps
+        self.config.seed = seed
+        self.config.prompt = [prompt]
+        self.config.image_size = size
+
+        all_batch_samples = self.pipeline(self.config) 
+        
+        sample = all_batch_samples[0][0]
+        # sample is torch.Size([1, C, f, H, W])
+
+        output = sample.squeeze(0).permute(1, 2, 3, 0).cpu().float()
+        # torch.Size([1, C, f, H, W]) -> torch.Size([f, H, W, C])
+        # BFloat16 -> Float
+
+        return output

src/videogen_hub/infermodels/opensora_12.py

@@ -0,0 +1,139 @@
+import os
+
+from huggingface_hub import snapshot_download, hf_hub_download
+
+from videogen_hub import MODEL_PATH
+
+
+class OpenSora12:
+    def __init__(self, device="gpu"):
+        """
+        1. Download the pretrained model and put it inside MODEL_PATH/modelscope
+        2. Create Pipeline
+        Note: it seems that the model needed from model_dir cannot support cpu
+        Args:
+            device: 'gpu' or 'cpu' the device to use the model
+        """
+
+        from mmengine import Config as mmengine_config
+        from videogen_hub.pipelines.opensora.scripts.inference import main
+        model_path = snapshot_download("hpcai-tech/OpenSora-STDiT-v3",
+                                       local_dir=os.path.join(MODEL_PATH, 'OpenSora-STDiT-v3'))
+        self.pipeline = main
+        self.config = {
+            # Basic video frame settings
+            "num_frames": 51,  # Total number of frames in a clip
+            "frame_interval": 1,  # Interval between frames
+            "fps": 24,  # Frames per second
+            "image_size": [480, 854],  # Resolution of each frame (height, width)
+            # Model configuration for multi-resolution and specific model parameters
+            "multi_resolution": "STDiT2",  # Multi-resolution model type
+            "model": {
+                "type": "STDiT3-XL/2",  # Model type and size
+                "from_pretrained": os.path.join(MODEL_PATH, "STDiT3-XL_2"),  # Path to pretrained checkpoint
+                "file_name": "model.safetensors",  # Name of the model file
+                "input_sq_size": 512,  # Input square size for the model
+                "qk_norm": True,  # Whether to normalize query-key in attention
+                "enable_flashattn": False,  # Enable flash attention mechanism, require flash_attn package
+                "enable_layernorm_kernel": False,  # Enable layer normalization in kernel, requires apex package
+            },
+            # Variational Autoencoder (VAE) specific settings
+            "vae": {
+                "type": "OpenSoraVAE_V1_2",  # Type of the autoencoder
+                "from_pretrained": "hpcai-tech/OpenSora-VAE-v1.2",  # Pretrained model from Hugging Face
+                #"cache_dir": os.path.join(MODEL_PATH, "OpenSora-VAE-v1.2"),  # Local cache directory for model weights
+                "micro_frame_size": 17,
+                "micro_batch_size": 4,  # Batch size for processing
+            },
+            # Text encoder settings for embedding textual information
+            "text_encoder": {
+                "type": "t5",  # Text encoder model type
+                "from_pretrained": "DeepFloyd/t5-v1_1-xxl",  # Pretrained model
+                "cache_dir": os.path.join(MODEL_PATH, "t5-v1_1-xxl"),  # Cache directory
+                "model_max_length": 300,  # Max length of text inputs
+            },
+            # Scheduler settings for diffusion models
+            "scheduler": {
+                "type": "rflow",  # Type of scheduler for the diffusion process
+                "num_sampling_steps": 30,  # Number of sampling steps in diffusion
+                "cfg_scale": 7.0,  # Scale for classifier-free guidance
+                # "cfg_channel": 3,  # Number of channels for guidance
+            },
+            # Additional settings for processing and output
+            "dtype": "bf16",  # Data type for computation (bfloat16)
+            # "prompt_path": "./assets/texts/t2v_samples.txt",  # Path to text prompts
+            "prompt_path": None,  # Path to text prompts
+            "prompt": [
+                "A beautiful sunset over the city"
+            ],  # List of prompts for generation
+            "batch_size": 1,  # Batch size for generation
+            "seed": 42,  # Seed for random number generators
+            "save_dir": "./samples/samples/",  # Directory to save generated samples
+            "config": "sample.py",  # Path to this configuration file
+            "prompt_as_path": False,  # Treat the prompt as a file path (True/False)
+            "reference_path": None,  # Path to reference image/video for conditioning
+            "loop": 1,  # Number of times to loop the processing
+            "sample_name": None,  # Specific name for the generated sample
+            "num_sample": 1,  # Number of samples to generate
+            "aes": 6.5,
+            "flow": None,
+        }
+        self.config = mmengine_config(self.config)
+
+        hf_hub_download(
+            repo_id="hpcai-tech/OpenSora-STDiT-v3",
+            filename="model.safetensors",
+            local_dir=self.config.model.from_pretrained,
+        )
+
+        hf_hub_download(
+            repo_id="hpcai-tech/OpenSora-VAE-v1.2",
+            filename="model.safetensors",
+            local_dir=os.path.join(MODEL_PATH, "OpenSora-VAE-v1.2"),
+        )
+
+        hf_hub_download(
+            repo_id="DeepFloyd/t5-v1_1-xxl",
+            filename="pytorch_model-00001-of-00002.bin",
+            local_dir=self.config.text_encoder.cache_dir,
+        )
+
+    def infer_one_video(
+            self,
+            prompt: str = None,
+            size: list = [320, 512],
+            seconds: int = 2,
+            fps: int = 8,
+            seed: int = 42,
+    ):
+        """
+        Generates a single video based on the provided prompt and parameters.
+        The generated video always has resolution 256x256
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            size (list, optional): The resolution of the video. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+
+        self.config.num_frames = fps * seconds
+        self.config.fps = fps
+        self.config.seed = seed
+        self.config.prompt = [prompt]
+        self.config.image_size = size
+
+        all_batch_samples = self.pipeline(self.config)
+
+        sample = all_batch_samples[0][0]
+        # sample is torch.Size([1, C, f, H, W])
+
+        output = sample.squeeze(0).permute(1, 2, 3, 0).cpu().float()
+        # torch.Size([1, C, f, H, W]) -> torch.Size([f, H, W, C])
+        # BFloat16 -> Float
+
+        return output

src/videogen_hub/infermodels/opensora_plan.py

@@ -0,0 +1,73 @@
+import os
+
+from huggingface_hub import snapshot_download, hf_hub_download
+import torch
+
+from videogen_hub import MODEL_PATH
+
+
+class OpenSoraPlan():
+    def __init__(self, device="cuda"):
+        """
+    1. Download the pretrained model and put it inside MODEL_PATH
+    2. Create Pipeline
+    Note: it seems that the model needed from model_dir cannot support cpu
+    Args:
+        device: 'cuda' or 'cpu' the device to use the model
+    """
+        from videogen_hub.pipelines.opensora_plan.opensora.sample_t2v import OpenSoraPlanPipeline
+
+        model_path = snapshot_download('LanguageBind/Open-Sora-Plan-v1.1.0', local_dir = os.path.join(MODEL_PATH, 'Open-Sora-Plan-v1.1.0'))
+
+        arg_list = ['--model_path', model_path,
+                    '--version', '65x512x512',
+                    '--num_frames', '65',
+                    '--height', '512',
+                    '--width', '512',
+                    '--cache_dir', MODEL_PATH,
+                    '--text_encoder_name', 'DeepFloyd/t5-v1_1-xxl',
+                    '--text_prompt', 'prompt_list_0.txt',
+                    '--ae', 'CausalVAEModel_4x8x8',
+                    '--ae_path', "/remote-home1/yeyang/CausalVAEModel_4x8x8",
+                    '--save_img_path', "./sample_video_65x512x512",
+                    '--fps', '24',
+                    '--guidance_scale', '7.5',
+                    '--num_sampling_steps', '150',
+                    '--enable_tiling']
+        self.pipeline = OpenSoraPlanPipeline(arg_list, device)
+
+    def infer_one_video(
+            self,
+            prompt: str = None,
+            size: list = [320, 512],
+            seconds: int = 2,
+            fps: int = 8,
+            seed: int = 42,
+    ):
+        """
+    Generates a single video based on the provided prompt and parameters.
+    Note that there are only 3 available shapes: (1 or 65 or 221)xHxW
+    The output is of shape [frames, channels, height, width].
+    Args:
+        prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+        seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+        fps (int, optional): The frames per second of the video. Defaults to 8.
+        seed (int, optional): The seed for random number generation. Defaults to 42.
+
+    Returns:
+        torch.Tensor: The generated video as a tensor.
+    """
+
+        torch.manual_seed(seed)
+
+        self.pipeline.args.text_prompt = prompt
+        self.pipeline.args.num_frames = fps * seconds
+        self.pipeline.args.fps = fps
+        self.pipeline.args.height = size[0]
+        self.pipeline.args.width = size[1]
+
+        samples = self.pipeline.inference(save_output=False)
+        # samples is torch.Size([B, T, H, W, C])
+
+        output = samples.squeeze(0).permute(0, 3, 1, 2).cpu().float()
+        return output

src/videogen_hub/infermodels/seine.py

@@ -0,0 +1,52 @@
+import os
+
+import torch
+from PIL import Image
+from huggingface_hub import snapshot_download, hf_hub_download
+
+from videogen_hub import MODEL_PATH
+
+
+class SEINE():
+    def __init__(self):
+        """
+        1. Download the pretrained model and put it inside MODEL_PATH/SEINE
+        2. Create Pipeline.
+        """
+        from videogen_hub.pipelines.seine.SEINEPipeline import SEINEPipeline
+
+        seine_path = hf_hub_download(repo_id="Vchitect/SEINE", filename="seine.pt", local_dir=os.path.join(MODEL_PATH, "SEINE"))
+        pretrained_model_path = snapshot_download(repo_id="CompVis/stable-diffusion-v1-4",
+                                                  local_dir=os.path.join(MODEL_PATH, "SEINE", "stable-diffusion-v1-4"),
+                                                  ignore_patterns=["*pytorch_model.bin", "*fp16*", "*non_ema*"])
+
+        self.pipeline = SEINEPipeline(seine_path, pretrained_model_path,
+                                      'src/videogen_hub/pipelines/seine/sample_i2v.yaml')
+
+    def infer_one_video(self,
+                        input_image: Image.Image,
+                        prompt: str = None,
+                        size: list = [320, 512],
+                        seconds: int = 2,
+                        fps: int = 8,
+                        seed: int = 42):
+        """
+        Generates a single video based on a textual prompt and first frame image, using either a provided image or an image path as the starting point. The output is a tensor representing the video.
+    
+        Args:
+            input_image (PIL.Image.Image): The input image to use as the basis for video generation.
+            prompt (str, optional): The text prompt that guides the video generation. If not specified, the video generation will rely solely on the input image. Defaults to None.
+            size (list, optional): Specifies the resolution of the output video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The number of frames per second in the generated video. This determines how smooth the video appears. Defaults to 8.
+            seed (int, optional): A seed value for random number generation, ensuring reproducibility of the video generation process. Defaults to 42.
+    
+        Returns:
+            torch.Tensor: A tensor representing the generated video, structured as (time, channel, height, width).
+        """
+        video = self.pipeline.infer_one_video(input_image=input_image,
+                                              text_prompt=prompt,
+                                              output_size=size,
+                                              num_frames=seconds * fps,
+                                              seed=seed)
+        return video

src/videogen_hub/infermodels/show_one.py

@@ -0,0 +1,79 @@
+import os
+
+
+from videogen_hub import MODEL_PATH
+
+
+class ShowOne():
+    def __init__(self):
+        """
+        Initialize the Pipeline, which download all necessary models.
+        """
+        from videogen_hub.pipelines.show_1.run_inference import ShowOnePipeline
+        from huggingface_hub import snapshot_download
+
+        base_path = snapshot_download(
+            repo_id="showlab/show-1-base",
+            local_dir=os.path.join(MODEL_PATH, "showlab", "show-1-base"),
+            local_dir_use_symlinks = False
+        )
+
+        interp_path = snapshot_download(
+            repo_id="showlab/show-1-interpolation",
+            local_dir=os.path.join(MODEL_PATH, "showlab", "show-1-interpolation"),
+            
+        )
+
+        deepfloyd_path = snapshot_download(
+            repo_id="DeepFloyd/IF-II-L-v1.0",
+            local_dir=os.path.join(MODEL_PATH, "DeepFloyd/IF-II-L-v1.0"),
+            
+        )
+
+        sr1_path = snapshot_download(
+            repo_id="showlab/show-1-sr1",
+            local_dir=os.path.join(MODEL_PATH, "showlab", "show-1-sr1"),
+            
+        )
+
+        sr2_path = snapshot_download(
+            repo_id="showlab/show-1-sr2",
+            local_dir=os.path.join(MODEL_PATH, "showlab", "show-1-sr2"),
+            
+        )
+
+        self.pipeline = ShowOnePipeline(base_path, interp_path, deepfloyd_path, sr1_path, sr2_path)
+
+    def infer_one_video(self,
+                        prompt: str = None,
+                        size: list = [320, 512],
+                        seconds: int = 2,
+                        fps: int = 8,
+                        seed: int = 42):
+        """
+        Generates a single video based on a textual prompt. The output is a tensor representing the video.
+        Since the initial_num_frames is set to be 8 as shown in paper in the pipeline,
+        we need the (number of frames - 1) divisible by 7 to manage interpolation.
+    
+        Args:
+            prompt (str, optional): The text prompt that guides the video generation. If not specified, the video generation will rely solely on the input image. Defaults to None.
+            size (list, optional): Specifies the resolution of the output video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The number of frames per second in the generated video. This determines how smooth the video appears. Defaults to 8.
+            seed (int, optional): A seed value for random number generation, ensuring reproducibility of the video generation process. Defaults to 42.
+    
+        Returns:
+            torch.Tensor: A tensor representing the generated video, structured as (time, channel, height, width).
+        """
+        num_frames = fps * seconds
+
+        assert (num_frames - 1) % 7 == 0
+        scaling_factor = (num_frames - 1) // 7
+        video = self.pipeline.inference(prompt=prompt,
+                                        negative_prompt="",
+                                        output_size=size,
+                                        initial_num_frames=8,
+                                        scaling_factor=scaling_factor,
+                                        seed=seed)
+
+        return video

src/videogen_hub/infermodels/streamingt2v.py

@@ -0,0 +1,49 @@
+import os
+
+from huggingface_hub import hf_hub_download
+
+from videogen_hub import MODEL_PATH
+
+
+class StreamingT2V:
+    def __init__(self, device="cuda"):
+        """
+        Initializes the StreamingT2V model.
+
+        Args:
+            device (str, optional): The device to run the model on. Defaults to "cuda".
+        """
+
+        from videogen_hub.pipelines.streamingt2v.streamingt2v_pipeline import pipeline
+        # https://huggingface.co/spaces/PAIR/StreamingT2V/resolve/main/t2v_enhanced/checkpoints/streaming_t2v.ckpt?download=true
+        model_url = "https://huggingface.co/spaces/PAIR/StreamingT2V/resolve/main/t2v_enhanced/checkpoints/streaming_t2v.ckpt?download=true"
+        # Download the file
+        ckpt_file_streaming_t2v = hf_hub_download(repo_id="PAIR/StreamingT2V",
+                                                  filename="streaming_t2v.ckpt",
+                                                  local_dir=os.path.join(MODEL_PATH, "streamingtv2"))
+
+        self.pipeline = pipeline
+
+    def infer_one_video(
+            self,
+            prompt: str = None,
+            size: list = [320, 512],
+            seconds: int = 2,
+            fps: int = 8,
+            seed: int = 42,
+    ):
+        """
+        Generates a single video based on the provided prompt and parameters.
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor.
+        """
+
+        return self.pipeline(prompt, size, seconds, fps, seed)

src/videogen_hub/infermodels/t2v_turbo.py

@@ -0,0 +1,147 @@
+import os
+
+from huggingface_hub import hf_hub_download, snapshot_download
+import torch
+
+from videogen_hub import MODEL_PATH
+
+
+class T2VTurbo():
+    def __init__(self, base_model="vc2", merged=True, device="cuda"):
+        """
+    1. Download the pretrained model and put it inside MODEL_PATH
+    2. Create Pipeline
+    Args:
+        device: 'cuda' or 'cpu' the device to use the model
+    """
+        from videogen_hub.pipelines.t2v_turbo.inference_vc2 import T2VTurboVC2Pipeline1
+        from videogen_hub.pipelines.t2v_turbo.inference_ms import T2VTurboMSPipeline1
+
+        self.config = {
+            "model": {
+                "target": "lvdm.models.ddpm3d.LatentDiffusion",
+                "params": {
+                    "linear_start": 0.00085,
+                    "linear_end": 0.012,
+                    "num_timesteps_cond": 1,
+                    "timesteps": 1000,
+                    "first_stage_key": "video",
+                    "cond_stage_key": "caption",
+                    "cond_stage_trainable": False,
+                    "conditioning_key": "crossattn",
+                    "image_size": [320, 512],
+                    "channels": 4,
+                    "scale_by_std": False,
+                    "scale_factor": 0.18215,
+                    "use_ema": False,
+                    "uncond_type": "empty_seq",
+                    "use_scale": True,
+                    "scale_b": 0.7,
+                    "unet_config": {
+                        "target": "lvdm.modules.networks.openaimodel3d.UNetModel",
+                        "params": {
+                            "in_channels": 4,
+                            "out_channels": 4,
+                            "model_channels": 320,
+                            "attention_resolutions": [4, 2, 1],
+                            "num_res_blocks": 2,
+                            "channel_mult": [1, 2, 4, 4],
+                            "num_head_channels": 64,
+                            "transformer_depth": 1,
+                            "context_dim": 1024,
+                            "use_linear": True,
+                            "use_checkpoint": True,
+                            "temporal_conv": True,
+                            "temporal_attention": True,
+                            "temporal_selfatt_only": True,
+                            "use_relative_position": False,
+                            "use_causal_attention": False,
+                            "temporal_length": 16,
+                            "addition_attention": True,
+                            "fps_cond": True
+                        }
+                    },
+                    "first_stage_config": {
+                        "target": "lvdm.models.autoencoder.AutoencoderKL",
+                        "params": {
+                            "embed_dim": 4,
+                            "monitor": "val / rec_loss",
+                            "ddconfig": {
+                                "double_z": True,
+                                "z_channels": 4,
+                                "resolution": 512,
+                                "in_channels": 3,
+                                "out_ch": 3,
+                                "ch": 128,
+                                "ch_mult": [1, 2, 4, 4],
+                                "num_res_blocks": 2,
+                                "attn_resolutions": [],
+                                "dropout": 0.0
+                            },
+                            "lossconfig": {
+                                "target": "torch.nn.Identity"
+                            }
+                        }
+                    },
+                    "cond_stage_config": {
+                        "target": "lvdm.modules.encoders.condition.FrozenOpenCLIPEmbedder",
+                        "params": {
+                            "freeze": True,
+                            "layer": "penultimate"
+                        }
+                    }
+                }
+            }
+        }
+        if base_model == "vc2" and merged:
+            merged_model_path = hf_hub_download(repo_id="jiachenli-ucsb/T2V-Turbo-VC2-Merged",
+                                                filename="t2v_turbo_vc2.pt",
+                                                local_dir=os.path.join(MODEL_PATH, "T2V-Turbo-VC2"))
+            self.pipeline = T2VTurboVC2Pipeline1(self.config, merged, device, None, merged_model_path)
+
+        elif base_model == "vc2":
+            base_model_path = hf_hub_download(repo_id="VideoCrafter/VideoCrafter2",
+                                              filename="model.ckpt",
+                                              local_dir=os.path.join(MODEL_PATH, "videocrafter2"))
+
+            unet_lora_path = hf_hub_download(repo_id="jiachenli-ucsb/T2V-Turbo-VC2",
+                                             filename="unet_lora.pt",
+                                             local_dir=os.path.join(MODEL_PATH, "T2V-Turbo-VC2"))
+            # It uses the config provided above.
+            self.pipeline = T2VTurboVC2Pipeline1(self.config, merged, device, unet_lora_path, base_model_path)
+        else:
+            base_model_path = snapshot_download(repo_id="ali-vilab/text-to-video-ms-1.7b",
+                                                local_dir=os.path.join(MODEL_PATH, "modelscope_1.7b"))
+
+            unet_lora_path = hf_hub_download(repo_id="jiachenli-ucsb/T2V-Turbo-MS",
+                                             filename="unet_lora.pt",
+                                             local_dir=os.path.join(MODEL_PATH, "T2V-Turbo-MS"))
+
+            # It uses the config provided by base_model.
+            self.pipeline = T2VTurboMSPipeline1(device, unet_lora_path, base_model_path)
+
+    def infer_one_video(
+            self,
+            prompt: str = None,
+            size: list = [320, 512],
+            seconds: int = 2,
+            fps: int = 8,
+            seed: int = 42,
+    ):
+        """
+    Generates a single video based on the provided prompt and parameters.
+    The output is of shape [frames, channels, height, width].
+    Args:
+        prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+        seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+        fps (int, optional): The frames per second of the video. Defaults to 8.
+        seed (int, optional): The seed for random number generation. Defaults to 42.
+
+    Returns:
+        torch.Tensor: The generated video as a tensor.
+    """
+        output = self.pipeline.inference(prompt=prompt, height=size[0], width=size[1],
+                                         seed=seed, num_frames=seconds * fps, fps=fps, randomize_seed=False)
+        # [channels, frames, height, width] -> [frames, channels, height, width]
+        output = output.squeeze().permute(1, 0, 2, 3)
+        return output.cpu()

src/videogen_hub/infermodels/videocrafter.py

@@ -0,0 +1,63 @@
+import torch
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+import os
+
+from videogen_hub import MODEL_PATH
+
+
+class VideoCrafter2():
+    def __init__(self, device="cuda"):
+        """
+        1. Download the pretrained model and put it inside MODEL_PATH/videocrafter2
+        2. Create Pipeline
+        Args:
+            device: 'cuda' or 'cpu' the device to use the model
+        """
+        from videogen_hub.pipelines.videocrafter.inference import VideoCrafterPipeline
+
+        model_path = hf_hub_download(repo_id="VideoCrafter/VideoCrafter2",
+                                     filename="model.ckpt",
+                                     local_dir=os.path.join(MODEL_PATH, "videocrafter2"))
+        config_path = str(Path(__file__).parent.parent.absolute())
+        config_path = os.path.join(config_path, 'pipelines/videocrafter/inference_t2v_512_v2.0.yaml')
+
+        arg_list = ['--mode', 'base',
+                    '--ckpt_path', model_path,
+                    '--config', config_path,
+                    '--n_samples', '1',
+                    '--bs', '1',
+                    '--unconditional_guidance_scale', '12.0',
+                    '--ddim_steps', '50',
+                    '--ddim_eta', '1.0',
+                    '--fps', '8']
+
+        self.pipeline = VideoCrafterPipeline(arg_list, device, 0, 1)
+
+    def infer_one_video(self,
+                        prompt: str = None,
+                        size: list = [320, 512],
+                        seconds: int = 2,
+                        fps: int = 8,
+                        seed: int = 42):
+        """
+        Generates a single video based on the provided prompt and parameters.
+
+        Args:
+            prompt (str, optional): The text prompt to generate the video from. Defaults to None.
+            size (list, optional): The size of the video as [height, width]. Defaults to [320, 512].
+            seconds (int, optional): The duration of the video in seconds. Defaults to 2.
+            fps (int, optional): The frames per second of the video. Defaults to 8.
+            seed (int, optional): The seed for random number generation. Defaults to 42.
+
+        Returns:
+            torch.Tensor: The generated video as a tensor, the shape being [num_frames, 3, height, width]
+
+        """
+        torch.manual_seed(seed)
+        video = self.pipeline.run_inference(prompt,
+                                            video_length=seconds * fps,
+                                            height=size[0],
+                                            width=size[1])
+
+        return video.squeeze(0, 1).cpu().permute(1, 0, 2, 3)

src/videogen_hub/metrics/brisque_metric.py

@@ -0,0 +1,47 @@
+from brisque import BRISQUE
+from PIL import Image
+import numpy as np
+from typing import List
+
+ROUND_DIGIT=3
+NUM_ASPECT=5
+
+BRISQUE_POINT_LOW=10
+BRISQUE_POINT_MID=30
+BRISQUE_POINT_HIGH=50
+
+class MetricBRISQUE():
+    def __init__(self) -> None:
+        """
+        Initialize a class MetricBRISQUE for testing visual quality of a given video.
+        
+        """
+        None
+
+    def evaluate(self,frame_list:List[Image.Image]):
+        """
+        Calculate BRISQUE for visual quality for each frame of the given video and take the average value, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation
+            
+        Returns:
+            piqe_avg: float, the computed average BRISQUE among the frames
+            quantized_ans: int, the quantized value of the above avg score based on pre-defined thresholds.
+        """
+        brisque_list=[]
+        for frame in frame_list:
+            brisque_score=BRISQUE().score(frame)
+            brisque_list.append(brisque_score)
+        brisque_avg=np.mean(brisque_list)
+        quantized_ans=0
+        if brisque_avg < BRISQUE_POINT_LOW:
+            quantized_ans=4
+        elif brisque_avg < BRISQUE_POINT_MID:
+            quantized_ans=3
+        elif brisque_avg < BRISQUE_POINT_HIGH:
+            quantized_ans=2
+        else:
+            quantized_ans=1
+        return brisque_avg, quantized_ans

src/videogen_hub/metrics/clip-sim_metric.py

@@ -0,0 +1,63 @@
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from typing import List
+from transformers import CLIPProcessor, CLIPModel
+        
+ROUND_DIGIT=3
+NUM_ASPECT=5
+
+CLIP_POINT_HIGH=0.97
+CLIP_POINT_MID=0.9
+CLIP_POINT_LOW=0.8   
+
+
+class MetricCLIP_sim():
+    def __init__(self, device = "cuda") -> None: 
+        """
+        Initialize a class MetricCLIP_sim with the specified device for testing temporal consistency of a given video.
+        
+        Args:
+            device (str, optional): The device on which the model will run. Defaults to "cuda".
+        """
+        self.device = device
+        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.model.to(self.device)
+        self.tokenizer = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    def evaluate(self,frame_list:List[Image.Image]):
+        """
+        Calculate the cosine similarity between the CLIP features of adjacent frames of a given video to test temporal consistency, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            
+        Returns:
+            clip_frame_score: float, the computed CLIP feature cosine similarity between each adjacent pair of frames and then averaged among all the pairs.
+            quantized_ans: int, the quantized value of the above avg CLIP-Sim scores based on pre-defined thresholds.
+        """
+        
+        device=self.model.device
+        frame_sim_list=[]
+        for f_idx in range(len(frame_list)-1):
+            frame_1 = frame_list[f_idx]
+            frame_2 = frame_list[f_idx+1]
+            input_1 = self.tokenizer(images=frame_1, return_tensors="pt", padding=True).to(device)
+            input_2 = self.tokenizer(images=frame_2, return_tensors="pt", padding=True).to(device)
+            output_1 = self.model.get_image_features(**input_1).flatten()
+            output_2 = self.model.get_image_features(**input_2).flatten()
+            cos_sim = F.cosine_similarity(output_1, output_2, dim=0).item()
+            frame_sim_list.append(cos_sim)
+            
+        clip_frame_score = np.mean(frame_sim_list)
+        quantized_ans=0
+        if clip_frame_score >= CLIP_POINT_HIGH:
+            quantized_ans=4
+        elif clip_frame_score < CLIP_POINT_HIGH and clip_frame_score >= CLIP_POINT_MID:
+            quantized_ans=3
+        elif clip_frame_score < CLIP_POINT_MID and clip_frame_score >= CLIP_POINT_LOW:
+            quantized_ans=2
+        else:
+            quantized_ans=1
+        return clip_frame_score, quantized_ans

src/videogen_hub/metrics/clipscore_metric.py

@@ -0,0 +1,65 @@
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from typing import List
+from transformers import CLIPProcessor, CLIPModel
+
+NUM_ASPECT=5
+ROUND_DIGIT=3
+MAX_LENGTH = 76
+
+MAX_NUM_FRAMES=8
+
+CLIP_POINT_LOW=0.27
+CLIP_POINT_MID=0.31
+CLIP_POINT_HIGH=0.35
+
+
+class MetricCLIPScore():
+    def __init__(self, device="cuda") -> None: 
+        """
+        Initialize a MetricCLIPScore object with the specified device.
+        
+        Args:
+            device (str, optional): The device on which the model will run. Defaults to "cuda".
+        """
+        self.device = device
+        self.model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.model.to(self.device)
+        self.tokenizer = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+    def evaluate(self, frame_list:List[Image.Image], text:str,):
+        """
+        Calculate the cosine similarity of between CLIP features of text prompt and each frame of a given video to test text-to-video alignment, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            text:str, text prompt for generating the video.
+            
+        Returns:
+            clip_score_avg: float, the computed average CLIP-Score between each frame and the text prompt.
+            quantized_ans: int, the quantized value of the above avg SSIM scores based on pre-defined thresholds.
+        """
+        
+        device=self.model.device
+        input_t = self.tokenizer(text=text, max_length=MAX_LENGTH, truncation=True, return_tensors="pt", padding=True).to(device)
+        cos_sim_list=[]
+        for image in frame_list:
+            input_f = self.tokenizer(images=image, return_tensors="pt", padding=True).to(device)
+            output_t = self.model.get_text_features(**input_t).flatten()
+            output_f = self.model.get_image_features(**input_f).flatten()
+            cos_sim = F.cosine_similarity(output_t, output_f, dim=0).item()
+            cos_sim_list.append(cos_sim)
+        clip_score_avg=np.mean(cos_sim_list)
+        quantized_ans=0
+        if clip_score_avg < CLIP_POINT_LOW:
+            quantized_ans=1
+        elif clip_score_avg >= CLIP_POINT_LOW and clip_score_avg < CLIP_POINT_MID:
+            quantized_ans=2
+        elif clip_score_avg >= CLIP_POINT_MID and clip_score_avg < CLIP_POINT_HIGH:
+            quantized_ans=3
+        else:
+            quantized_ans=4
+        return clip_score_avg, quantized_ans
+

src/videogen_hub/metrics/dino-sim_metric.py

@@ -0,0 +1,71 @@
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from typing import List
+from torchvision.models import vit_b_16 
+import torchvision.transforms as transforms
+
+ROUND_DIGIT=3
+NUM_ASPECT=5
+
+DINO_POINT_HIGH=0.97
+DINO_POINT_MID=0.9
+DINO_POINT_LOW=0.8   
+
+
+class MetricDINO_sim():
+    def __init__(self, device="cuda") -> None: 
+        """
+        Initialize a class MetricDINO_sim with the specified device for testing temporal consistency of a given video.
+        
+        Args:
+            device (str, optional): The device on which the model will run. Defaults to "cuda".
+        """
+        self.device = device
+        self.model = vit_b_16(pretrained=True)
+        self.model.to(self.device).eval()  
+        self.preprocess = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def evaluate(self, frame_list:List[Image.Image]):
+        """
+        Calculate the cosine similarity between the DINO features of adjacent frames of a given video to test temporal consistency, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            
+        Returns:
+            dino_frame_score: float, the computed DINO feature cosine similarity between each adjacent pair of frames and then averaged among all the pairs.
+            quantized_ans: int, the quantized value of the above avg DINO-Sim scores based on pre-defined thresholds.
+        """
+        
+        device = self.device
+        frame_sim_list=[]
+        for f_idx in range(len(frame_list)-1):
+            frame_1=frame_list[f_idx]
+            frame_2=frame_list[f_idx+1]
+            frame_tensor_1 = self.preprocess(frame_1).unsqueeze(0).to(device)
+            frame_tensor_2 = self.preprocess(frame_2).unsqueeze(0).to(device)
+            with torch.no_grad():
+                feat_1 = self.model(frame_tensor_1).flatten()
+                feat_2 = self.model(frame_tensor_2).flatten()
+            cos_sim=F.cosine_similarity(feat_1, feat_2, dim=0).item()
+            frame_sim_list.append(cos_sim)   
+    
+        dino_frame_score = np.mean(frame_sim_list)
+        quantized_ans=0
+        if dino_frame_score >= DINO_POINT_HIGH:
+            quantized_ans=4
+        elif dino_frame_score < DINO_POINT_HIGH and dino_frame_score >= DINO_POINT_MID:
+            quantized_ans=3
+        elif dino_frame_score < DINO_POINT_MID and dino_frame_score >= DINO_POINT_LOW:
+            quantized_ans=2
+        else:
+            quantized_ans=1
+        return dino_frame_score, quantized_ans

src/videogen_hub/metrics/mse-dyn_metric.py

@@ -0,0 +1,59 @@
+import numpy as np
+import cv2
+from PIL import Image
+from typing import List
+from skimage.metrics import structural_similarity as ssim
+from skimage import io, color
+
+ROUND_DIGIT=3
+DYN_SAMPLE_STEP=4
+NUM_ASPECT=5
+
+MSE_POINT_HIGH=3000
+MSE_POINT_MID=1000
+MSE_POINT_LOW=100
+
+
+class MetricMSE_dyn():
+    def __init__(self) -> None: 
+        """
+        Initialize a class MetricMSE_dyn for testing dynamic degree of a given video.
+        
+        """
+        None
+
+    def evaluate(self, frame_list:List[Image.Image]):
+        """
+        Calculate the MSE (Mean Squared Error) between frames sampled at regular intervals of a given video to test dynamic_degree, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            
+        Returns:
+            mse_avg: float, the computed MSE between frames sampled at regular intervals and then averaged among all the pairs.
+            quantized_ans: int, the quantized value of the above avg MSE scores based on pre-defined thresholds.
+        """
+        
+        mse_list=[]
+        sampled_list = frame_list[::DYN_SAMPLE_STEP]
+        for f_idx in range(len(sampled_list)-1):        
+            imageA = cv2.cvtColor(np.array(sampled_list[f_idx]), cv2.COLOR_RGB2BGR)
+            imageB = cv2.cvtColor(np.array(sampled_list[f_idx+1]), cv2.COLOR_RGB2BGR)
+            
+            err = np.sum((imageA.astype("float") - imageB.astype("float")) ** 2)
+            err /= float(imageA.shape[0] * imageA.shape[1])
+            mse_value = err
+            mse_list.append(mse_value)
+        mse_avg=np.mean(mse_list)
+        quantized_ans=0
+        if mse_avg >= MSE_POINT_HIGH:
+            quantized_ans=4
+        elif mse_avg < MSE_POINT_HIGH and mse_avg >= MSE_POINT_MID:
+            quantized_ans=3
+        elif mse_avg < MSE_POINT_MID and mse_avg >= MSE_POINT_LOW:
+            quantized_ans=2
+        else:
+            quantized_ans=1
+            
+        return mse_avg, quantized_ans

src/videogen_hub/metrics/piqe_metric.py

@@ -0,0 +1,49 @@
+from pypiqe import piqe
+from PIL import Image
+import numpy as np
+from typing import List
+
+ROUND_DIGIT=3
+NUM_ASPECT=5
+
+PIQE_POINT_LOW=15
+PIQE_POINT_MID=30
+PIQE_POINT_HIGH=50
+
+class MetricPIQE():
+    def __init__(self) -> None:
+        """
+        Initialize a class MetricPIQE for testing visual quality of a given video.
+        
+        """
+        None
+
+    def evaluate(self,frame_list:List[Image.Image]):
+        """
+        Calculate PIQE for visual quality for each frame of the given video and take the average value, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            
+        Returns:
+            piqe_avg: float, the computed average PIQE among the frames.
+            quantized_ans: int, the quantized value of the above avg score based on pre-defined thresholds.
+        """
+        piqe_list=[]
+        for frame in frame_list:
+            frame=np.array(frame)
+            piqe_score, _,_,_ = piqe(frame)
+            piqe_list.append(piqe_score)
+        piqe_avg=np.mean(piqe_list)
+        quantized_ans=0
+        if piqe_avg < PIQE_POINT_LOW:
+            quantized_ans=4
+        elif piqe_avg < PIQE_POINT_MID:
+            quantized_ans=3
+        elif piqe_avg < PIQE_POINT_HIGH:
+            quantized_ans=2
+        else:
+            quantized_ans=1
+        return piqe_avg, quantized_ans
+    

src/videogen_hub/metrics/ssim-dyn_metric.py

@@ -0,0 +1,60 @@
+import numpy as np
+from PIL import Image
+from typing import List
+from skimage.metrics import structural_similarity as ssim
+from skimage import io, color
+
+ROUND_DIGIT=3
+DYN_SAMPLE_STEP=4
+NUM_ASPECT=5
+
+SSIM_POINT_HIGH=0.9
+SSIM_POINT_MID=0.7
+SSIM_POINT_LOW=0.5
+
+
+
+class MetricSSIM_dyn():
+    def __init__(self) -> None: 
+        """
+        Initialize a class MetricSSIM_dyn for testing dynamic degree of a given video.
+        
+        """
+        None
+
+    def evaluate(self, frame_list:List[Image.Image]):
+        """
+        Calculate the MSE between frames sampled at regular intervals of a given video to test dynamic_degree, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            
+        Returns:
+            ssim_avg: float, the computed SSIM between frames sampled at regular intervals and then averaged among all the pairs.
+            quantized_ans: int, the quantized value of the above avg SSIM scores based on pre-defined thresholds.
+        """
+        
+        ssim_list=[]
+        sampled_list = frame_list[::DYN_SAMPLE_STEP]
+        for f_idx in range(len(sampled_list)-1):
+            frame_1=sampled_list[f_idx]
+            frame_1_gray=color.rgb2gray(frame_1)
+            frame_2=sampled_list[f_idx+1]
+            frame_2_gray=color.rgb2gray(frame_2)
+
+            ssim_value, _ = ssim(frame_1_gray, frame_2_gray, full=True,\
+                                    data_range=frame_2_gray.max() - frame_2_gray.min())
+            ssim_list.append(ssim_value)
+        ssim_avg=np.mean(ssim_list)
+        
+        quantized_ans=0
+        if ssim_avg >= SSIM_POINT_HIGH:
+            quantized_ans=1
+        elif ssim_avg <= SSIM_POINT_HIGH and ssim_avg > SSIM_POINT_MID:
+            quantized_ans=2
+        elif ssim_avg <= SSIM_POINT_MID and ssim_avg > SSIM_POINT_LOW:
+            quantized_ans=3
+        else:
+            quantized_ans=4
+        return ssim_avg, quantized_ans

src/videogen_hub/metrics/ssim-sim_metric.py

@@ -0,0 +1,59 @@
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from typing import List
+from skimage.metrics import structural_similarity as ssim
+from skimage import io, color
+
+ROUND_DIGIT=3
+NUM_ASPECT=5
+
+TEM_SSIM_POINT_HIGH=0.9
+TEM_SSIM_POINT_MID=0.75
+TEM_SSIM_POINT_LOW=0.6 
+
+
+class MetricSSIM_sim():
+    def __init__(self) -> None: 
+        """
+        Initialize a class MetricSSIM_sim for testing temporal consistency of a given video.
+        
+        """
+        None
+
+    def evaluate(self, frame_list:List[Image.Image]):
+        """
+        Calculate the SSIM between adjacent frames of a given video to test temporal consistency, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            
+        Returns:
+            ssim_avg: float, the computed SSIM between each adjacent pair of frames and then averaged among all the pairs.
+            quantized_ans: int, the quantized value of the above avg SSIM scores based on pre-defined thresholds.
+        """
+        
+        ssim_list=[]
+        for f_idx in range(len(frame_list)-1):
+            frame_1=frame_list[f_idx]
+            frame_1_gray=color.rgb2gray(frame_1)
+            frame_2=frame_list[f_idx+1]
+            frame_2_gray=color.rgb2gray(frame_2)
+
+            ssim_value, _ = ssim(frame_1_gray, frame_2_gray, full=True,\
+                                    data_range=frame_2_gray.max() - frame_2_gray.min())
+            ssim_list.append(ssim_value)
+        ssim_avg=np.mean(ssim_list)
+        quantized_ans=0
+        if ssim_avg >= TEM_SSIM_POINT_HIGH:
+            quantized_ans=4
+        elif ssim_avg < TEM_SSIM_POINT_HIGH and ssim_avg >= TEM_SSIM_POINT_MID:
+            quantized_ans=3
+        elif ssim_avg < TEM_SSIM_POINT_MID and ssim_avg >= TEM_SSIM_POINT_LOW:
+            quantized_ans=2
+        else:
+            quantized_ans=1
+        return ssim_avg, quantized_ans
+

src/videogen_hub/metrics/xclipscore_metric.py

@@ -0,0 +1,72 @@
+import numpy as np
+from PIL import Image
+import torch.nn.functional as F
+from typing import List
+from transformers import AutoTokenizer, AutoModel, AutoProcessor
+
+NUM_ASPECT=5
+ROUND_DIGIT=3
+MAX_LENGTH = 76
+
+MAX_NUM_FRAMES=8
+
+X_CLIP_POINT_LOW=0.15
+X_CLIP_POINT_MID=0.225
+X_CLIP_POINT_HIGH=0.30
+
+
+def _read_video_frames(frames, max_frames):
+    total_frames = len(frames)
+    indices = np.linspace(0, total_frames - 1, num=max_frames).astype(int)
+
+    selected_frames = [np.array(frames[i]) for i in indices]
+    return np.stack(selected_frames)
+
+
+class MetricXCLIPScore():
+    def __init__(self, device="cuda") -> None: 
+        """
+        Initialize a MetricXCLIPScore object with the specified device.
+        
+        Args:
+            device (str, optional): The device on which the model will run. Defaults to "cuda".
+        """
+        
+        self.model = AutoModel.from_pretrained("microsoft/xclip-base-patch32")
+        self.processor = AutoProcessor.from_pretrained("microsoft/xclip-base-patch32")
+        self.tokenizer = AutoTokenizer.from_pretrained("microsoft/xclip-base-patch32")
+
+    def evaluate(self, frame_list:List[Image.Image], text:str,):
+        """
+        Calculate the cosine similarity of between X-CLIP features of text prompt and the given video to test text-to-video alignment, 
+        then quantize the orginal output based on some predefined thresholds.
+        
+        Args:
+            frame_list:List[Image.Image], frames of the video used in calculation.
+            text:str, text prompt for generating the video.
+            
+        Returns:
+            xclip_score_avg: float, the computed X-CLIP-Score between video and its text prompt.
+            quantized_ans: int, the quantized value of the above avg SSIM scores based on pre-defined thresholds.
+        """
+        
+        input_text = self.tokenizer([text], max_length=MAX_LENGTH, truncation=True, padding=True, return_tensors="pt")
+        text_feature = self.model.get_text_features(**input_text).flatten()
+
+        video=_read_video_frames(frame_list,MAX_NUM_FRAMES)
+        
+        input_video = self.processor(videos=list(video), return_tensors="pt")
+        video_feature = self.model.get_video_features(**input_video).flatten()
+        cos_sim=F.cosine_similarity(text_feature, video_feature, dim=0).item()
+        quantized_ans=0
+        if cos_sim < X_CLIP_POINT_LOW:
+            quantized_ans=1
+        elif cos_sim >= X_CLIP_POINT_LOW and cos_sim < X_CLIP_POINT_MID:
+            quantized_ans=2
+        elif cos_sim >= X_CLIP_POINT_MID and cos_sim < X_CLIP_POINT_HIGH:
+            quantized_ans=3
+        else:
+            quantized_ans=4
+        return cos_sim, quantized_ans
+
+

src/videogen_hub/utils/__init__.py

@@ -0,0 +1,17 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+
+def images_to_tensor(image_list):
+    """
+    Parse a list of PIL images and convert them to a PyTorch tensor in shape (T, C, H, W).
+    """
+    transform = transforms.ToTensor()
+    
+    # Convert each PIL image to tensor and store in a list
+    tensor_list = [transform(img) for img in image_list]
+    
+    # Stack the list of tensors along a new dimension to create the final tensor
+    tensor = torch.stack(tensor_list, dim=0)
+    
+    return tensor

src/videogen_hub/utils/file_helper.py

@@ -0,0 +1,24 @@
+import os
+from typing import Union, List, Optional
+from urllib.parse import urlparse
+import requests
+
+def get_file_path(filename: Union[str, os.PathLike], search_from: Union[str, os.PathLike] = "."):
+    """
+    Search for a file across a directory and return its absolute path.
+
+    Args:
+        filename (Union[str, os.PathLike]): The name of the file to search for.
+        search_from (Union[str, os.PathLike], optional): The directory from which to start the search. Defaults to ".".
+
+    Returns:
+        str: Absolute path to the found file.
+
+    Raises:
+        FileNotFoundError: If the file is not found.
+    """
+    for root, dirs, files in os.walk(search_from):
+        for name in files:
+            if name == filename:
+                return os.path.abspath(os.path.join(root, name))
+    raise FileNotFoundError(filename, "not found.")