del t2v

Former-commit-id: 5daa2952997a29e74e2b5d235001d6d9dabf9ff6

opensora/dataset/__init__.py

@@ -1,11 +1,12 @@
 from torchvision.transforms import Compose
-
+from transformers import AutoTokenizer
 
 from opensora.models.ae import vae, vqvae, videovae, videovqvae
 from .feature_datasets import LandscopeFeature, SkyFeature, LandscopeVideoFeature
 from torchvision import transforms
 from torchvision.transforms import Lambda
 
+from .t2v_datasets import T2V_dataset
 from .transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVideo, CenterCropResizeVideo
 from .ucf101 import UCF101
 from .sky_datasets import Sky
@@ -50,12 +51,21 @@ def getdataset(args):
         )
         return UCF101(args, transform=transform, temporal_sample=temporal_sample)
     elif args.dataset == 'sky':
-        transform_sky = transforms.Compose([
+        transform = transforms.Compose([
             ToTensorVideo(),
             CenterCropResizeVideo(args.max_image_size),
             RandomHorizontalFlipVideo(p=0.5),
             norm_fun
         ])
-        return Sky(args, transform=transform_sky, temporal_sample=temporal_sample)
+        return Sky(args, transform=transform, temporal_sample=temporal_sample)
+    elif args.dataset == 't2v':
+        transform = transforms.Compose([
+            ToTensorVideo(),
+            CenterCropResizeVideo(args.max_image_size),
+            RandomHorizontalFlipVideo(p=0.5),
+            norm_fun
+        ])
+        tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir='./cache_dir')
+        return T2V_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer)
     else:
         raise NotImplementedError(args.dataset)

opensora/dataset/t2v_datasets.py

@@ -0,0 +1,156 @@
+import os, io, csv, math, random
+import numpy as np
+import torchvision
+from einops import rearrange
+from decord import VideoReader
+
+import torch
+import torchvision.transforms as transforms
+from torch.utils.data.dataset import Dataset
+
+from opensora.utils.dataset_utils import DecordInit
+from opensora.utils.utils import text_preprocessing
+
+#
+# class T2V_dataset(Dataset):
+#     def __init__(self, args, transform, temporal_sample, tokenizer):
+#
+#         with open(args.data_path, 'r') as csvfile:
+#             self.samples = list(csv.DictReader(csvfile))
+#         self.video_folder = args.video_folder
+#         self.num_frames = args.num_frames
+#         self.transform = transform
+#         self.temporal_sample = temporal_sample
+#         self.tokenizer = tokenizer
+#         self.model_max_length = args.model_max_length
+#         self.v_decoder = DecordInit()
+#
+#     def __len__(self):
+#         return len(self.samples)
+#
+#     def __getitem__(self, idx):
+#         video_dict = self.samples[idx]
+#         try:
+#             video_path, text = video_dict['Filename'], video_dict['Video Description']
+#             video_path = os.path.join(self.video_folder, video_path)
+#
+#             video = self.tv_read(video_path)
+#             video = self.transform(video)  # T C H W -> T C H W
+#             video = video.transpose(0, 1)  # T C H W -> C T H
+#
+#
+#             text = text_preprocessing(text)
+#             text_tokens_and_mask = self.tokenizer(
+#                 text,
+#                 max_length=self.model_max_length,
+#                 padding='max_length',
+#                 truncation=True,
+#                 return_attention_mask=True,
+#                 add_special_tokens=True,
+#                 return_tensors='pt'
+#             )
+#
+#             input_ids = text_tokens_and_mask['input_ids']
+#             cond_mask = text_tokens_and_mask['attention_mask']
+#
+#             return video, input_ids, cond_mask
+#         except Exception as e:
+#             print(f'Error with {e}, {video_path}')
+#             return self.__getitem__(random.randint(0, self.__len__()-1))
+#
+#     def tv_read(self, path):
+#         vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+#         total_frames = len(vframes)
+#
+#         # Sampling video frames
+#         start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+#         # assert end_frame_ind - start_frame_ind >= self.num_frames
+#         frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+#         video = vframes[frame_indice]  # (T, C, H, W)
+#
+#         return video
+#
+#     def decord_read(self, path):
+#         decord_vr = self.v_decoder(path)
+#         total_frames = len(decord_vr)
+#         # Sampling video frames
+#         start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+#         # assert end_frame_ind - start_frame_ind >= self.num_frames
+#         frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+#
+#         video_data = decord_vr.get_batch(frame_indice).asnumpy()
+#         video_data = torch.from_numpy(video_data)
+#         video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+#         return video_data
+
+
+class T2V_dataset(Dataset):
+    def __init__(self, args, transform, temporal_sample, tokenizer):
+
+        # with open(args.data_path, 'r') as csvfile:
+        #     self.samples = list(csv.DictReader(csvfile))
+        self.video_folder = args.video_folder
+        self.num_frames = args.num_frames
+        self.transform = transform
+        self.temporal_sample = temporal_sample
+        self.tokenizer = tokenizer
+        self.model_max_length = args.model_max_length
+        self.v_decoder = DecordInit()
+
+    def __len__(self):
+        # return len(self.samples)
+        return 10000
+
+    def __getitem__(self, idx):
+        # video_dict = self.samples[idx]
+        try:
+            # video_path, text = video_dict['Filename'], video_dict['Video Description']
+            # video_path = os.path.join(self.video_folder, video_path)
+            #
+            # video = self.tv_read(video_path)
+            # video = self.transform(video)  # T C H W -> T C H W
+            # video = video.transpose(0, 1)  # T C H W -> C T H W
+            video = torch.rand(3, 16, 256, 256)
+            text = 'ni hao'
+            text = text_preprocessing(text)
+            text_tokens_and_mask = self.tokenizer(
+                text,
+                max_length=self.model_max_length,
+                padding='max_length',
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors='pt'
+            )
+            input_ids = text_tokens_and_mask['input_ids']
+            cond_mask = text_tokens_and_mask['attention_mask']
+
+            return video, input_ids, cond_mask
+        except Exception as e:
+            # print(f'Error with {e}, {video_path}')
+            return self.__getitem__(random.randint(0, self.__len__() - 1))
+
+    def tv_read(self, path):
+        vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+        total_frames = len(vframes)
+
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+        video = vframes[frame_indice]  # (T, C, H, W)
+
+        return video
+
+    def decord_read(self, path):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        # assert end_frame_ind - start_frame_ind >= self.num_frames
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.num_frames, dtype=int)
+
+        video_data = decord_vr.get_batch(frame_indice).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data

opensora/dataset/ucf101.py

@@ -44,7 +44,6 @@ class UCF101(Dataset):
 
     def __getitem__(self, idx):
         video_path, label = self.samples[idx]
-
         try:
             video = self.tv_read(video_path)
             video = self.transform(video)  # T C H W -> T C H W

opensora/models/diffusion/__init__.py

@@ -1,8 +1,9 @@
 from .dit.modeling_dit import DiT_models
-from .latte.modeling_latte import Latte_models
+from .latte.modeling_latte import Latte_models, LatteT2V_models
 
 Diffusion_models = {}
 Diffusion_models.update(DiT_models)
 Diffusion_models.update(Latte_models)
+Diffusion_models.update(LatteT2V_models)
 
     

opensora/models/diffusion/latte/common.py

@@ -98,14 +98,14 @@ class Attention(nn.Module):
         if self.attention_mode == 'xformers':  # require pytorch 2.0
             if attn_mask is not None:
                 attn_mask = self.make_attn_mask(attn_mask)
-                attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).to(q.dtype)
+                attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).bool()
             with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
                 x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                                    dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C)
 
         elif self.attention_mode == 'flash':  # require pytorch 2.0
             # https://github.com/PKU-YuanGroup/Open-Sora-Plan/issues/109
-            assert attn_mask is None or torch.all(attn_mask.bool()), 'flash-attention do not support attention_mask'
+            assert attn_mask is None, 'flash-attention do not support attention_mask'
             with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
                 x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C)
 
@@ -117,8 +117,8 @@ class Attention(nn.Module):
                 attn = attn + attn_bias
             attn = attn.softmax(dim=-1)
             if torch.any(torch.isnan(attn)):
-                print('torch.any(torch.isnan(attn))')
-                attn = attn.masked_fill(torch.isnan(attn), float(0.))
+                raise ValueError('Found NAN in attn')
+                # attn = attn.masked_fill(torch.isnan(attn), float(0.))
             attn = self.attn_drop(attn)
             x = (attn @ v).transpose(1, 2).reshape(B, N, C)
 
@@ -169,28 +169,29 @@ class CrossAttention(nn.Module):
 
     def forward(self, x, cond, attn_mask=None, cond_mask=None):
         B, N, C = x.shape
+        _, N_, _ = cond.shape
         q = self.q_linear(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3).contiguous()
-        kv = self.kv_linear(cond).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous()
+        kv = self.kv_linear(cond).reshape(B, N_, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous()
         k, v = kv.unbind(0)  # make torchscript happy (cannot use tensor as tuple) b h n c
 
         if self.attention_mode == 'xformers' or self.attention_mode == 'flash':  # require pytorch 2.0
-            attn_mask = self.make_attn_mask(attn_mask, cond_mask)
+            attn_mask = self.make_attn_mask(N, attn_mask, cond_mask)
             if attn_mask is not None:
-                attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).to(q.dtype)
+                attn_mask = attn_mask.repeat(1, self.num_heads, 1, 1).bool()
             with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=False, enable_mem_efficient=True):
                 x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                                    dropout_p=self.attn_drop.p, scale=self.scale).reshape(B, N, C)
 
         elif self.attention_mode == 'math':
             attn = (q @ k.transpose(-2, -1)) * self.scale
-            attn_bias = self.make_attn_bias(attn_mask, cond_mask)
+            attn_bias = self.make_attn_bias(N, attn_mask, cond_mask)
             if attn_bias is not None:
                 attn_bias = attn_bias.repeat(1, self.num_heads, 1, 1).to(q.dtype)
                 attn = attn + attn_bias
             attn = attn.softmax(dim=-1)
             if torch.any(torch.isnan(attn)):
-                print('torch.any(torch.isnan(attn))')
-                attn = attn.masked_fill(torch.isnan(attn), float(0.))
+                raise ValueError('Found NAN in attn')
+                # attn = attn.masked_fill(torch.isnan(attn), float(0.))
             attn = self.attn_drop(attn)
             x = (attn @ v).transpose(1, 2).reshape(B, N, C)
 
@@ -202,11 +203,11 @@ class CrossAttention(nn.Module):
 
         return x
 
-    def make_attn_mask(self, attn_mask, cond_mask):
+    def make_attn_mask(self, N, attn_mask, cond_mask):
         if attn_mask is None and cond_mask is None:
             return None
         if attn_mask is None and cond_mask is not None:
-            attn_mask = cond_mask
+            attn_mask = torch.ones(cond_mask.shape[0], N).to(device=cond_mask.device, dtype=cond_mask.dtype)
         if attn_mask is not None and cond_mask is None:
             raise NotImplementedError
         attn_mask = attn_mask.flatten(1).unsqueeze(-1)  # b n -> b n 1
@@ -215,10 +216,10 @@ class CrossAttention(nn.Module):
         attn_mask = attn_mask.unsqueeze(1)  # b n n -> b 1 n n
         return attn_mask
 
-    def make_attn_bias(self, attn_mask, cond_mask):
+    def make_attn_bias(self, N, attn_mask, cond_mask):
         # The numerical range of bfloat16, float16 can't conver -1e8
         # Refer to https://discuss.pytorch.org/t/runtimeerror-value-cannot-be-converted-to-type-at-half-without-overflow-1e-30/109768
-        attn_mask = self.make_attn_mask(attn_mask, cond_mask)
+        attn_mask = self.make_attn_mask(N, attn_mask, cond_mask)
         if attn_mask is None:
             return None
         attn_bias = torch.where(attn_mask == 0, -1e8 if attn_mask.dtype == torch.float32 else -1e4, attn_mask)
@@ -294,6 +295,36 @@ class LabelEmbedder(nn.Module):
         embeddings = self.embedding_table(labels)
         return embeddings
 
+class CaptionEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+
+    def __init__(self, in_channels, hidden_size, uncond_prob, act_layer=nn.GELU(approximate='tanh'), token_num=120):
+        super().__init__()
+        self.y_proj = Mlp(in_features=in_channels, hidden_features=hidden_size, out_features=hidden_size, act_layer=act_layer, drop=0)
+        self.register_buffer("y_embedding", nn.Parameter(torch.randn(token_num, in_channels) / in_channels ** 0.5))
+        self.uncond_prob = uncond_prob
+
+    def token_drop(self, caption, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, caption)
+        return caption
+
+    def forward(self, caption, train, force_drop_ids=None):
+        if train:
+            assert caption.shape[2:] == self.y_embedding.shape
+        use_dropout = self.uncond_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            caption = self.token_drop(caption, force_drop_ids)
+        caption = self.y_proj(caption)
+        return caption
 
 #################################################################################
 #                                 Core Latte Model                                #
@@ -322,6 +353,31 @@ class TransformerBlock(nn.Module):
         x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
         return x
 
+class TransformerCrossConditonBlock(nn.Module):
+    """
+    A Latte tansformer block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
+        self.cross_attn = CrossAttention(hidden_size, num_heads, **block_kwargs)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, cond, t, attn_mask=None, cond_mask=None):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), attn_mask)
+        x = x + self.cross_attn(x, cond, attn_mask, cond_mask)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
 class FinalLayer(nn.Module):
     """
     The final layer of Latte.

opensora/models/diffusion/latte/configuration_latte.py

@@ -41,7 +41,7 @@ class LatteConfiguration:
         self.class_dropout_prob = class_dropout_prob
         self.num_classes = num_classes
         self.learn_sigma = learn_sigma
-        self.extras = learn_sigma
+        self.extras = extras
         self.attention_mode = attention_mode
         self.compress_kv = compress_kv
         self.attention_pe_mode = attention_pe_mode
@@ -62,6 +62,14 @@ class LatteConfiguration:
             config_dict = json.load(json_file)
         return cls(**config_dict)
 
+class LatteT2VConfiguration(LatteConfiguration):
+    def __init__(self,
+                 caption_channels=4096,
+                 model_max_length=120,
+                 **kwargs):
+        super(LatteT2VConfiguration, self).__init__(**kwargs)
+        self.caption_channels = caption_channels
+        self.model_max_length = model_max_length
 
 def Latte_XL_122_Config(**kwargs):
     return LatteConfiguration(depth=56, hidden_size=1152, patch_size_t=1, patch_size=2, num_heads=16, **kwargs)
@@ -100,9 +108,52 @@ def Latte_S_188_Config(**kwargs):
     return LatteConfiguration(depth=24, hidden_size=384, patch_size_t=1, patch_size=8, num_heads=6, **kwargs)
 
 
+def LatteT2V_XL_122_Config(**kwargs):
+    return LatteT2VConfiguration(depth=56, hidden_size=1152, patch_size_t=1, patch_size=2, num_heads=16, **kwargs)
+
+def LatteT2V_XL_144_Config(**kwargs):
+    return LatteT2VConfiguration(depth=56, hidden_size=1152, patch_size_t=1, patch_size=4, num_heads=16, **kwargs)
+
+def LatteT2V_XL_188_Config(**kwargs):
+    return LatteT2VConfiguration(depth=56, hidden_size=1152, patch_size_t=1, patch_size=8, num_heads=16, **kwargs)
+
+def LatteT2V_L_122_Config(**kwargs):
+    return LatteT2VConfiguration(depth=48, hidden_size=1024, patch_size_t=1, patch_size=2, num_heads=16, **kwargs)
+
+def LatteT2V_L_144_Config(**kwargs):
+    return LatteT2VConfiguration(depth=48, hidden_size=1024, patch_size_t=1, patch_size=4, num_heads=16, **kwargs)
+
+def LatteT2V_L_188_Config(**kwargs):
+    return LatteT2VConfiguration(depth=48, hidden_size=1024, patch_size_t=1, patch_size=8, num_heads=16, **kwargs)
+
+def LatteT2V_B_122_Config(**kwargs):
+    return LatteT2VConfiguration(depth=24, hidden_size=768, patch_size_t=1, patch_size=2, num_heads=12, **kwargs)
+
+def LatteT2V_B_144_Config(**kwargs):
+    return LatteT2VConfiguration(depth=24, hidden_size=768, patch_size_t=1, patch_size=4, num_heads=12, **kwargs)
+
+def LatteT2V_B_188_Config(**kwargs):
+    return LatteT2VConfiguration(depth=24, hidden_size=768, patch_size_t=1, patch_size=8, num_heads=12, **kwargs)
+
+def LatteT2V_S_122_Config(**kwargs):
+    return LatteT2VConfiguration(depth=24, hidden_size=384, patch_size_t=1, patch_size=2, num_heads=6, **kwargs)
+
+def LatteT2V_S_144_Config(**kwargs):
+    return LatteT2VConfiguration(depth=24, hidden_size=384, patch_size_t=1, patch_size=4, num_heads=6, **kwargs)
+
+def LatteT2V_S_188_Config(**kwargs):
+    return LatteT2VConfiguration(depth=24, hidden_size=384, patch_size_t=1, patch_size=8, num_heads=6, **kwargs)
+
 Latte_configs = {
     "Latte-XL/122": Latte_XL_122_Config, "Latte-XL/144": Latte_XL_144_Config, "Latte-XL/188": Latte_XL_188_Config,
     "Latte-L/122": Latte_L_122_Config, "Latte-L/144": Latte_L_144_Config, "Latte-L/188": Latte_L_188_Config,
     "Latte-B/122": Latte_B_122_Config, "Latte-B/144": Latte_B_144_Config, "Latte-B/188": Latte_B_188_Config,
     "Latte-S/122": Latte_S_122_Config, "Latte-S/144": Latte_S_144_Config, "Latte-S/188": Latte_S_188_Config,
 }
+
+LatteT2V_configs = {
+    "LatteT2V-XL/122": LatteT2V_XL_122_Config, "LatteT2V-XL/144": LatteT2V_XL_144_Config, "LatteT2V-XL/188": LatteT2V_XL_188_Config,
+    "LatteT2V-L/122": LatteT2V_L_122_Config, "LatteT2V-L/144": LatteT2V_L_144_Config, "LatteT2V-L/188": LatteT2V_L_188_Config,
+    "LatteT2V-B/122": LatteT2V_B_122_Config, "LatteT2V-B/144": LatteT2V_B_144_Config, "LatteT2V-B/188": LatteT2V_B_188_Config,
+    "LatteT2V-S/122": LatteT2V_S_122_Config, "LatteT2V-S/144": LatteT2V_S_144_Config, "LatteT2V-S/188": LatteT2V_S_188_Config,
+}

opensora/models/diffusion/latte/modeling_latte.py

@@ -3,9 +3,10 @@ from einops import rearrange, repeat
 from timm.layers import PatchEmbed
 from torch import nn
 import torch
-from .pos import get_1d_sincos_temp_embed, get_2d_sincos_pos_embed
-from .common import TimestepEmbedder, LabelEmbedder, TransformerBlock, FinalLayer
-from .configuration_latte import LatteConfiguration
+from opensora.models.diffusion.latte.pos import get_1d_sincos_temp_embed, get_2d_sincos_pos_embed
+from opensora.models.diffusion.latte.common import TimestepEmbedder, LabelEmbedder, TransformerBlock, FinalLayer, \
+    TransformerCrossConditonBlock, CaptionEmbedder
+from opensora.models.diffusion.latte.configuration_latte import LatteConfiguration, LatteT2VConfiguration
 
 class Latte(nn.Module):
     """
@@ -26,7 +27,7 @@ class Latte(nn.Module):
         class_dropout_prob = config.class_dropout_prob
         num_classes = config.num_classes
         learn_sigma = config.learn_sigma
-        extras = config.learn_sigma
+        extras = config.extras
         attention_mode = config.attention_mode
         compress_kv = config.compress_kv
         attention_pe_mode = config.attention_pe_mode
@@ -51,7 +52,6 @@ class Latte(nn.Module):
         self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
         self.t_embedder = TimestepEmbedder(hidden_size)
 
-        assert self.extras == 1 or self.extras == 2
         if self.extras == 2:
             self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
 
@@ -246,16 +246,135 @@ class Latte(nn.Module):
         eps = torch.cat([half_eps, half_eps], dim=0) 
         return torch.cat([eps, rest], dim=2)
 
+class LatteT2V(Latte):
+    def __init__(self, config: LatteT2VConfiguration):
+        super(LatteT2V, self).__init__(config)
+
+        input_size = config.input_size
+        patch_size = config.patch_size
+        patch_size_t = config.patch_size_t
+        in_channels = config.in_channels
+        hidden_size = config.hidden_size
+        depth = config.depth
+        num_heads = config.num_heads
+        mlp_ratio = config.mlp_ratio
+        num_frames = config.num_frames
+        class_dropout_prob = config.class_dropout_prob
+        num_classes = config.num_classes
+        learn_sigma = config.learn_sigma
+        extras = config.extras
+        attention_mode = config.attention_mode
+        compress_kv = config.compress_kv
+        attention_pe_mode = config.attention_pe_mode
+        pt_input_size = config.pt_input_size
+        pt_num_frames = config.pt_num_frames
+        intp_vfreq = config.intp_vfreq
+        caption_channels = config.caption_channels
+        model_max_length = config.model_max_length
+
+        self.config = config
+
+        if pt_input_size is None:
+            pt_input_size = input_size
+        if pt_num_frames is None:
+            pt_num_frames = num_frames
+        blocks = []
+        for i in range(depth):
+            if i % 2 == 0:
+                m = TransformerCrossConditonBlock(
+                    hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode,
+                    attention_pe_mode=attention_pe_mode,
+                    hw=(input_size[0] // patch_size, input_size[1] // patch_size),
+                    pt_hw=(pt_input_size[0] // patch_size, pt_input_size[1] // patch_size),
+                    intp_vfreq=intp_vfreq, compress_kv=compress_kv
+                )
+            else:
+                m = TransformerBlock(
+                    hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode,
+                    attention_pe_mode=attention_pe_mode,
+                    hw=(num_frames, 1),
+                    pt_hw=(pt_num_frames, 1),
+                    intp_vfreq=intp_vfreq, compress_kv=compress_kv
+                )
+            blocks.append(m)
+        self.blocks = nn.ModuleList(blocks)
+
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.y_embedder = CaptionEmbedder(in_channels=caption_channels, hidden_size=hidden_size,
+                                          uncond_prob=class_dropout_prob, act_layer=approx_gelu,
+                                          token_num=model_max_length)
+
+    def forward(self,
+                x,
+                t,
+                cond,
+                attn_mask=None,
+                cond_mask=None):
+        """
+        Forward pass of Latte.
+        x: (N, F, C, H, W) tensor of video inputs
+        t: (N,) tensor of diffusion timesteps
+        cond: (N, 1, L, C_) tensor of text conditions
+        attn_mask: (N, F, H, W)
+        cond_mask: (N, L)
+        """
+        attn_mask_temproal, attn_mask_spatial = None, None
+        if attn_mask is not None:
+            attn_mask_spatial = rearrange(attn_mask, 'b t h w -> (b t) (h w)')
+            attn_mask_temproal = rearrange(attn_mask, 'b t h w -> (b h w) t')
+
+        batches, frames, channels, high, weight = x.shape
+
+        x = rearrange(x, 'b f c h w -> (b f) c h w')
+
+        x = self.x_embedder(x) + self.pos_embed
+        t = self.t_embedder(t)
+        cond = self.y_embedder(cond, self.training)  # (N, 1, L, D)
+
+        # timestep condition
+        timestep_spatial = repeat(t, 'n d -> (n c) d', c=self.temp_embed.shape[1])
+        timestep_temp = repeat(t, 'n d -> (n c) d', c=self.pos_embed.shape[1])
+        cond_spatial = repeat(cond, 'b 1 l d -> (b f) l d', f=self.temp_embed.shape[1]).contiguous()
+        cond_mask_spatial = repeat(cond_mask, 'b l -> (b f) l', f=self.temp_embed.shape[1]).contiguous()
+
+        for i in range(0, len(self.blocks), 2):
+            spatial_block, temp_block = self.blocks[i:i + 2]
+
+            # cond only apply to spatial attention
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(spatial_block), x, cond_spatial, timestep_spatial, attn_mask_spatial, cond_mask_spatial)
+            else:
+                x = spatial_block(x, cond_spatial, timestep_spatial, attn_mask_spatial, cond_mask_spatial)
+
+            x = rearrange(x, '(b f) t d -> (b t) f d', b=batches)
+            # Add Time Embedding
+            if i == 0:
+                x = x + self.temp_embed
+
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(temp_block), x, timestep_temp, attn_mask_temproal)
+            else:
+                x = temp_block(x, timestep_temp, attn_mask_temproal)
+            x = rearrange(x, '(b t) f d -> (b f) t d', b=batches)
+
+        x = self.final_layer(x, timestep_spatial)
+        x = self.unpatchify(x)
+        x = rearrange(x, '(b f) c h w -> b f c h w', b=batches)
+        return x
 
 #################################################################################
 #                                   Latte Configs                                  #
 #################################################################################
 
-from .configuration_latte import (
+from opensora.models.diffusion.latte.configuration_latte import (
     Latte_XL_122_Config, Latte_XL_144_Config, Latte_XL_188_Config,
     Latte_L_122_Config, Latte_L_144_Config, Latte_L_188_Config,
     Latte_B_122_Config, Latte_B_144_Config, Latte_B_188_Config,
-    Latte_S_122_Config, Latte_S_144_Config, Latte_S_188_Config, LatteConfiguration,
+    Latte_S_122_Config, Latte_S_144_Config, Latte_S_188_Config,
+    LatteT2V_XL_122_Config, LatteT2V_XL_144_Config, LatteT2V_XL_188_Config,
+    LatteT2V_L_122_Config, LatteT2V_L_144_Config, LatteT2V_L_188_Config,
+    LatteT2V_B_122_Config, LatteT2V_B_144_Config, LatteT2V_B_188_Config,
+    LatteT2V_S_122_Config, LatteT2V_S_144_Config, LatteT2V_S_188_Config,
 )
 
 def Latte_XL_122(**kwargs):
@@ -295,6 +414,42 @@ def Latte_S_188(**kwargs):
     return Latte(Latte_S_188_Config(**kwargs))
 
 
+def LatteT2V_XL_122(**kwargs):
+    return LatteT2V(LatteT2V_XL_122_Config(**kwargs))
+
+def LatteT2V_XL_144(**kwargs):
+    return LatteT2V(LatteT2V_XL_144_Config(**kwargs))
+
+def LatteT2V_XL_188(**kwargs):
+    return LatteT2V(LatteT2V_XL_188_Config(**kwargs))
+
+def LatteT2V_L_122(**kwargs):
+    return LatteT2V(LatteT2V_L_122_Config(**kwargs))
+
+def LatteT2V_L_144(**kwargs):
+    return LatteT2V(LatteT2V_L_144_Config(**kwargs))
+
+def LatteT2V_L_188(**kwargs):
+    return LatteT2V(LatteT2V_L_188_Config(**kwargs))
+
+def LatteT2V_B_122(**kwargs):
+    return LatteT2V(LatteT2V_B_122_Config(**kwargs))
+
+def LatteT2V_B_144(**kwargs):
+    return LatteT2V(LatteT2V_B_144_Config(**kwargs))
+
+def LatteT2V_B_188(**kwargs):
+    return LatteT2V(LatteT2V_B_188_Config(**kwargs))
+
+def LatteT2V_S_122(**kwargs):
+    return LatteT2V(LatteT2V_S_122_Config(**kwargs))
+
+def LatteT2V_S_144(**kwargs):
+    return LatteT2V(LatteT2V_S_144_Config(**kwargs))
+
+def LatteT2V_S_188(**kwargs):
+    return LatteT2V(LatteT2V_S_188_Config(**kwargs))
+
 Latte_models = {
     "Latte-XL/122": Latte_XL_122, "Latte-XL/144": Latte_XL_144, "Latte-XL/188": Latte_XL_188,
     "Latte-L/122": Latte_L_122, "Latte-L/144": Latte_L_144, "Latte-L/188": Latte_L_188,
@@ -303,22 +458,34 @@ Latte_models = {
 }
 
 
+LatteT2V_models = {
+    "LatteT2V-XL/122": LatteT2V_XL_122, "LatteT2V-XL/144": LatteT2V_XL_144, "LatteT2V-XL/188": LatteT2V_XL_188,
+    "LatteT2V-L/122": LatteT2V_L_122, "LatteT2V-L/144": LatteT2V_L_144, "LatteT2V-L/188": LatteT2V_L_188,
+    "LatteT2V-B/122": LatteT2V_B_122, "LatteT2V-B/144": LatteT2V_B_144, "LatteT2V-B/188": LatteT2V_B_188,
+    "LatteT2V-S/122": LatteT2V_S_122, "LatteT2V-S/144": LatteT2V_S_144, "LatteT2V-S/188": LatteT2V_S_188,
+}
+
 if __name__ == '__main__':
 
     import torch
-
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    """        
+    x: (N, F, C, H, W) tensor of video inputs
+    t: (N,) tensor of diffusion timesteps
+    cond: (N, L, C_) tensor of text conditions
+    attn_mask: (N, F, H, W)
+    cond_mask: (N, L)
+    """
+    x = torch.randn(2, 16, 4, 32, 32).to(device)
+    t = torch.tensor([1, 2]).to(device)
+    cond = torch.randn(2, 1, 120, 4096).to(device)
+    attn_mask = torch.ones(2, 16, 32//2, 32//2).to(device)
+    cond_mask = torch.ones(2, 120).to(device)
+    model = LatteT2V_XL_122().to(device)
 
-    img = torch.randn(3, 16, 4, 32, 32).to(device)
-    t = torch.tensor([1, 2, 3]).to(device)
-    y = torch.tensor([1, 2, 3]).to(device)
-    network = Latte_XL_122().to(device)
-    from thop import profile 
-    flops, params = profile(network, inputs=(img, t))
-    print('FLOPs = ' + str(flops/1000**3) + 'G')
-    print('Params = ' + str(params/1000**2) + 'M')
-    # y_embeder = LabelEmbedder(num_classes=101, hidden_size=768, dropout_prob=0.5).to(device)
-    # lora.mark_only_lora_as_trainable(network)
-    # out = y_embeder(y, True)
-    # out = network(img, t, y)
-    # print(out.shape)
+    print(f"Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
+    for n, p in model.named_parameters():
+        print(f"Training Parameters: {n}, {p.shape}")
+    with torch.no_grad():
+        out = model(x, t, cond, attn_mask, cond_mask)
+        print(out.shape)  # torch.Size([2, 16, 8, 32, 32])

opensora/models/text_encoder/__init__.py

@@ -0,0 +1,39 @@
+import torch
+from torch import nn
+from transformers import T5EncoderModel
+
+from opensora.utils.utils import get_precision
+
+
+class T5Wrapper(nn.Module):
+    def __init__(self, args):
+        super(T5Wrapper, self).__init__()
+        self.model_name = args.text_encoder_name
+        dtype = get_precision(args)
+        t5_model_kwargs = {'cache_dir': './cache_dir', 'low_cpu_mem_usage': True, 'torch_dtype': dtype}
+        self.text_enc = T5EncoderModel.from_pretrained(self.model_name, **t5_model_kwargs).eval()
+
+    def forward(self, input_ids, attention_mask):
+        text_encoder_embs = self.text_enc(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
+        return text_encoder_embs.detach()
+
+class CLIPWrapper(nn.Module):
+    def __init__(self, args, device, cache_dir='./cache_dir', **kwargs):
+        super(CLIPWrapper, self).__init__()
+        raise NotImplementedError
+
+    def forward(self, prompts):  # prompts = ['I am a test caption', 'Test twice']
+        raise NotImplementedError
+
+
+text_encoder = {
+    'DeepFloyd/t5-v1_1-xxl': T5Wrapper,
+    'clip': CLIPWrapper
+}
+
+
+def get_text_enc(args):
+    """deprecation"""
+    text_enc = text_encoder.get(args.text_encoder_name, None)
+    assert text_enc is not None
+    return text_enc(args)

opensora/models/text_encoder/t5.py

@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+import os
+import re
+import html
+import urllib.parse as ul
+
+import ftfy
+import torch
+from bs4 import BeautifulSoup
+from transformers import T5EncoderModel, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+class T5Embedder:
+
+    available_models = ['t5-v1_1-xxl']
+    bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
+
+    def __init__(self, device, dir_or_name='t5-v1_1-xxl', *, cache_dir='./cache_dir', hf_token=None, use_text_preprocessing=True,
+                 t5_model_kwargs=None, torch_dtype=None, model_max_length=120):
+        self.device = torch.device(device)
+        self.torch_dtype = torch_dtype or torch.bfloat16
+        if t5_model_kwargs is None:
+            t5_model_kwargs = {'low_cpu_mem_usage': True, 'torch_dtype': self.torch_dtype}
+            t5_model_kwargs['device_map'] = {'shared': self.device, 'encoder': self.device}
+
+        self.use_text_preprocessing = use_text_preprocessing
+        self.hf_token = hf_token
+        self.cache_dir = cache_dir
+        self.dir_or_name = dir_or_name
+        cache_dir = os.path.join(self.cache_dir, 't5-v1_1-xxl')
+        for filename in ['config.json', 'special_tokens_map.json', 'spiece.model', 'tokenizer_config.json',
+                         'pytorch_model-00001-of-00002.bin', 'pytorch_model-00002-of-00002.bin', 'pytorch_model.bin.index.json']:
+            hf_hub_download(repo_id='DeepFloyd/t5-v1_1-xxl', filename=filename, cache_dir=cache_dir,
+                            force_filename=filename, token=self.hf_token)
+
+        print(cache_dir)
+        self.tokenizer = AutoTokenizer.from_pretrained(cache_dir)
+        self.model = T5EncoderModel.from_pretrained(cache_dir, **t5_model_kwargs).eval()
+        self.model_max_length = model_max_length
+
+    def get_text_embeddings(self, texts):
+        texts = [self.text_preprocessing(text) for text in texts]
+
+        text_tokens_and_mask = self.tokenizer(
+            texts,
+            max_length=self.model_max_length,
+            padding='max_length',
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors='pt'
+        )
+
+        text_tokens_and_mask['input_ids'] = text_tokens_and_mask['input_ids']
+        text_tokens_and_mask['attention_mask'] = text_tokens_and_mask['attention_mask']
+
+        with torch.no_grad():
+            text_encoder_embs = self.model(
+                input_ids=text_tokens_and_mask['input_ids'].to(self.device),
+                attention_mask=text_tokens_and_mask['attention_mask'].to(self.device),
+            )['last_hidden_state'].detach()
+        return text_encoder_embs, text_tokens_and_mask['attention_mask'].to(self.device)
+
+    def text_preprocessing(self, text):
+        if self.use_text_preprocessing:
+            # The exact text cleaning as was in the training stage:
+            text = self.clean_caption(text)
+            text = self.clean_caption(text)
+            return text
+        else:
+            return text.lower().strip()
+
+    @staticmethod
+    def basic_clean(text):
+        text = ftfy.fix_text(text)
+        text = html.unescape(html.unescape(text))
+        return text.strip()
+
+    def clean_caption(self, caption):
+        caption = str(caption)
+        caption = ul.unquote_plus(caption)
+        caption = caption.strip().lower()
+        caption = re.sub('<person>', 'person', caption)
+        # urls:
+        caption = re.sub(
+            r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+            '', caption)  # regex for urls
+        caption = re.sub(
+            r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+            '', caption)  # regex for urls
+        # html:
+        caption = BeautifulSoup(caption, features='html.parser').text
+
+        # @<nickname>
+        caption = re.sub(r'@[\w\d]+\b', '', caption)
+
+        # 31C0—31EF CJK Strokes
+        # 31F0—31FF Katakana Phonetic Extensions
+        # 3200—32FF Enclosed CJK Letters and Months
+        # 3300—33FF CJK Compatibility
+        # 3400—4DBF CJK Unified Ideographs Extension A
+        # 4DC0—4DFF Yijing Hexagram Symbols
+        # 4E00—9FFF CJK Unified Ideographs
+        caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
+        caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
+        caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
+        caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
+        caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
+        caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
+        caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
+        #######################################################
+
+        # все виды тире / all types of dash --> "-"
+        caption = re.sub(
+            r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
+            '-', caption)
+
+        # кавычки к одному стандарту
+        caption = re.sub(r'[`´«»“”¨]', '"', caption)
+        caption = re.sub(r'[‘’]', "'", caption)
+
+        # &quot;
+        caption = re.sub(r'&quot;?', '', caption)
+        # &amp
+        caption = re.sub(r'&amp', '', caption)
+
+        # ip adresses:
+        caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
+
+        # article ids:
+        caption = re.sub(r'\d:\d\d\s+$', '', caption)
+
+        # \n
+        caption = re.sub(r'\\n', ' ', caption)
+
+        # "#123"
+        caption = re.sub(r'#\d{1,3}\b', '', caption)
+        # "#12345.."
+        caption = re.sub(r'#\d{5,}\b', '', caption)
+        # "123456.."
+        caption = re.sub(r'\b\d{6,}\b', '', caption)
+        # filenames:
+        caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
+
+        #
+        caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
+        caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
+
+        caption = re.sub(self.bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+        caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
+
+        # this-is-my-cute-cat / this_is_my_cute_cat
+        regex2 = re.compile(r'(?:\-|\_)')
+        if len(re.findall(regex2, caption)) > 3:
+            caption = re.sub(regex2, ' ', caption)
+
+        caption = self.basic_clean(caption)
+
+        caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
+        caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
+        caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
+
+        caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
+        caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
+        caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
+        caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
+        caption = re.sub(r'\bpage\s+\d+\b', '', caption)
+
+        caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
+
+        caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
+
+        caption = re.sub(r'\b\s+\:\s+', r': ', caption)
+        caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
+        caption = re.sub(r'\s+', ' ', caption)
+
+        caption.strip()
+
+        caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
+        caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
+        caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
+        caption = re.sub(r'^\.\S+$', '', caption)
+
+        return caption.strip()
+
+if __name__ == '__main__':
+    t5 = T5Embedder(device="cuda", cache_dir='./cache_dir', torch_dtype=torch.float)
+    device = t5.device
+    prompts = ['I am a test caption', 'Test twice']
+    with torch.no_grad():
+        caption_embs, emb_masks = t5.get_text_embeddings(prompts)
+        emb_dict = {
+            'caption_feature': caption_embs.float().cpu().data.numpy(),
+            'attention_mask': emb_masks.cpu().data.numpy(),
+        }
+    import ipdb;ipdb.set_trace()
+    print()

opensora/sample/sample.py

@@ -78,7 +78,7 @@ def main(args):
             y_null = torch.tensor([args.num_classes] * 1, device=device)
             y = torch.cat([y, y_null], dim=0)
             model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
-            sample_fn = model.forward_with_cfg
+            sample_fn = model.module.forward_with_cfg
         else:
             sample_fn = model.forward
             model_kwargs = dict(y=None)

opensora/train/train.py

@@ -15,7 +15,7 @@ from accelerate import Accelerator
 from torch import nn
 
 from opensora.utils.utils import get_experiment_dir, create_logger, requires_grad, update_ema, write_tensorboard, \
-    cleanup, create_tensorboard
+    cleanup, create_tensorboard, get_precision
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
@@ -50,6 +50,7 @@ def main(args):
     # Setup accelerator:
     accelerator = Accelerator(mixed_precision=args.mixed_precision, gradient_accumulation_steps=args.gradient_accumulation_steps)
     device = accelerator.device
+    dtype = get_precision(args)
 
     # Setup an experiment folder:
     if args.resume_from_checkpoint:
@@ -104,17 +105,16 @@ def main(args):
         pixelart_alpha = torch.load(args.pretrained, map_location='cpu')['state_dict']
         checkpoint = {}
         for k, v in pixelart_alpha.items():
-            if 'x_embedder' in k or 't_embedder' in k:
+            if 'x_embedder' in k or 't_embedder' in k or 'y_embedder' in k:
                 checkpoint[k] = v
             if k.startswith('blocks'):
-                if 'cross_attn' not in k and 'drop_path' not in k and 'scale_shift_table' not in k:
-                    if 'q_norm' not in k and 'k_norm' not in k:
-                        blk_id = k[7]
-                        new_k = f'blocks.{str(int(blk_id) * 2)}{k[8:]}'
-                        checkpoint[new_k] = v
-        total_load = set(model.state_dict().keys()) & set(checkpoint.keys())
-        msg = model.load_state_dict(checkpoint, strict=False)
-        logger.info(f'Successfully load {len(total_load)} keys from {args.pretrained}!')
+                k_spilt = k.split('.')
+                blk_id = str(int(k_spilt[1]) * 2)
+                k_spilt[1] = blk_id
+                new_k = '.'.join(k_spilt)
+                checkpoint[new_k] = v
+        missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+        logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)} keys from {args.pretrained}!')
 
     model = model.to(device)
     # Note that parameter initialization is done within the DiT constructor
@@ -204,13 +204,6 @@ def main(args):
     if accelerator.is_main_process:
         logger.info(f"Training for {num_train_epochs} epochs...")
 
-    if args.mixed_precision == "bf16":
-        dtype = torch.bfloat16
-    elif args.mixed_precision == "fp16":
-        dtype = torch.float16
-    else:
-        dtype = torch.float32
-
     for epoch in range(first_epoch, num_train_epochs):
         logger.info(f"Beginning epoch {epoch}...")
         for step, (x, y, attn_mask) in enumerate(loader):
@@ -218,31 +211,29 @@ def main(args):
                 continue
             x = x.to(device)  # B C T H W
             y = y.to(device)
-            attn_mask = attn_mask.to(device)  # B T H W
-            with torch.no_grad():
-                # Map input images to latent space + normalize latents
-                x = ae.encode(x)  # B C T H W -> B T C H W
-
-            if args.extras == 78: # text-to-video
-                raise 'T2V training are Not supported at this moment!'
-            elif args.extras == 2:
+            # attn_mask = attn_mask.to(device)  # B T H W
+            assert torch.all(attn_mask.bool()), 'do not enable dynamic input'
+            attn_mask = None
+            if args.extras == 2:
                 model_kwargs = dict(y=y, attn_mask=attn_mask)
-            else:
+            elif args.extras == 1:
                 model_kwargs = dict(y=None, attn_mask=attn_mask)
+            else:
+                raise NotImplementedError
             with torch.autocast(device_type='cuda', dtype=dtype):
+                with torch.no_grad():
+                    # Map input images to latent space + normalize latents
+                    x = ae.encode(x)  # B C T H W -> B T C H W
                 t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
                 loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
                 loss = loss_dict["loss"].mean()
                 opt.zero_grad()
 
             accelerator.backward(loss)
-
-            # if train_steps < args.start_clip_iter: # if train_steps >= start_clip_iter, will clip gradient
-            #     gradient_norm = clip_grad_norm_(model.module.parameters(), args.clip_max_norm, clip_grad=False)
-            # else:
-            #     gradient_norm = clip_grad_norm_(model.module.parameters(), args.clip_max_norm, clip_grad=True)
-
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
             opt.step()
+            lr_scheduler.step()
 
             if args.use_ema:
                 update_ema(ema, model.module)
@@ -311,7 +302,7 @@ if __name__ == "__main__":
 
 
     # --------------------------------------
-    parser.add_argument("--ae", type=str, default="ucf101_stride4x4x4")
+    parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse")
     parser.add_argument("--extras", type=int, default=2, choices=[1, 2, 78])
     parser.add_argument("--pretrained", type=str, default=None)
     parser.add_argument("--sample-rate", type=int, default=4)
@@ -327,7 +318,7 @@ if __name__ == "__main__":
     parser.add_argument("--attention-mode", type=str, choices=['xformers', 'math', 'flash'], default="math")
     parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
     parser.add_argument("--mixed-precision", type=str, default=None, choices=[None, "fp16", "bf16"])
-    parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)")
+    parser.add_argument("--clip-grad-norm", default=1.0, type=float, help="the maximum gradient norm (default None)")
     # --------------------------------------
 
 

opensora/train/train_t2v.py

@@ -0,0 +1,341 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+A minimal training script for DiT using PyTorch DDP.
+"""
+import math
+
+import torch
+# the first flag below was False when we tested this script but True makes A100 training a lot faster:
+from accelerate import Accelerator
+from torch import nn
+
+from opensora.models.text_encoder import get_text_enc
+from opensora.utils.utils import get_experiment_dir, create_logger, requires_grad, update_ema, write_tensorboard, \
+    cleanup, create_tensorboard, get_precision
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+import torch.distributed as dist
+from diffusers.optimization import get_scheduler
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from collections import OrderedDict
+from copy import deepcopy
+from glob import glob
+from time import time
+import argparse
+import logging
+import os
+from opensora.dataset import getdataset
+from opensora.models.ae import getae
+from opensora.utils.dataset_utils import Collate
+from opensora.models.ae import ae_stride_config, ae_channel_config
+from opensora.models.diffusion import Diffusion_models
+from opensora.models.diffusion.diffusion import create_diffusion
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):
+    """
+    Trains a new DiT model.
+    """
+    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
+
+    # Setup accelerator:
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, gradient_accumulation_steps=args.gradient_accumulation_steps)
+    device = accelerator.device
+    dtype = get_precision(args)
+
+    # Setup an experiment folder:
+    if args.resume_from_checkpoint:
+        experiment_dir = f"{args.results_dir}/{args.resume_from_checkpoint}"
+    else:
+        experiment_index = len(glob(f"{args.results_dir}/*"))
+        model_string_name = args.model.replace("/", "-")  # e.g., Latte-XL/122 --> Latte-XL-122 (for naming folders)
+        ae_string_name = args.ae.split("/")[-1]  # e.g., stabilityai/sd-vae-ft-mse --> sd-vae-ft-mse (for naming folders)
+        num_frame_string = 'F' + str(args.num_frames) + 'S' + str(args.sample_rate)
+        experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}-{ae_string_name}-{num_frame_string}-{args.dataset}"  # Create an experiment folder
+        experiment_dir = get_experiment_dir(experiment_dir, args)
+    checkpoint_dir = f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
+    if accelerator.is_main_process:
+        os.makedirs(args.results_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        logger = create_logger(experiment_dir)
+        tb_writer = create_tensorboard(experiment_dir)
+        logger.info(f"Experiment directory created at {experiment_dir}")
+        logger.info(f'{args}')
+    else:
+        logger = create_logger(None)
+        tb_writer = None
+
+    # Create model:
+    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
+    args.ae_stride_t, args.ae_stride_h, args.ae_stride_w = ae_stride_t, ae_stride_h, ae_stride_w
+    args.ae_stride = args.ae_stride_h
+    patch_size = args.model[-3:]
+    patch_size_t, patch_size_h, patch_size_w = int(patch_size[0]), int(patch_size[1]), int(patch_size[2])
+    args.patch_size = patch_size_h
+    args.patch_size_t, args.patch_size_h, args.patch_size_w = patch_size_t, patch_size_h, patch_size_w
+    assert ae_stride_h == ae_stride_w, f"Support only ae_stride_h == ae_stride_w now, but found ae_stride_h ({ae_stride_h}), ae_stride_w ({ae_stride_w})"
+    assert patch_size_h == patch_size_w, f"Support only patch_size_h == patch_size_w now, but found patch_size_h ({patch_size_h}), patch_size_w ({patch_size_w})"
+    assert args.num_frames % ae_stride_t == 0, f"Num_frames must be divisible by ae_stride_t, but found num_frames ({args.num_frames}), ae_stride_t ({ae_stride_t})."
+    assert args.max_image_size % ae_stride_h == 0, f"Image size must be divisible by ae_stride_h, but found max_image_size ({args.max_image_size}),  ae_stride_h ({ae_stride_h})."
+
+    latent_size = (args.max_image_size // ae_stride_h, args.max_image_size // ae_stride_w)
+
+    model = Diffusion_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes,
+        in_channels=ae_channel_config[args.ae],
+        extras=args.extras,
+        num_frames=args.num_frames // ae_stride_t,
+        attention_mode=args.attention_mode
+    )
+    model.gradient_checkpointing = args.gradient_checkpointing
+
+    # # use pretrained model?
+    if args.pretrained:
+        # load from pixart-alpha
+        pixelart_alpha = torch.load(args.pretrained, map_location='cpu')['state_dict']
+        checkpoint = {}
+        for k, v in pixelart_alpha.items():
+            if 'x_embedder' in k or 't_embedder' in k or 'y_embedder' in k:
+                checkpoint[k] = v
+            if k.startswith('blocks'):
+                k_spilt = k.split('.')
+                blk_id = str(int(k_spilt[1]) * 2)
+                k_spilt[1] = blk_id
+                new_k = '.'.join(k_spilt)
+                checkpoint[new_k] = v
+        missing_keys, unexpected_keys = model.load_state_dict(checkpoint, strict=False)
+        logger.info(f'Successfully load {len(model.state_dict()) - len(missing_keys)} keys from {args.pretrained}!')
+
+    model = model.to(device)
+    # Note that parameter initialization is done within the DiT constructor
+    if args.use_ema:
+        ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+        requires_grad(ema, False)
+    diffusion = create_diffusion(timestep_respacing="")  # default: 1000 steps, linear noise schedule
+
+    if args.use_compile:
+        model = torch.compile(model)
+
+    ae = getae(args).to(device)
+    text_enc = get_text_enc(args).to(device)
+    if accelerator.is_main_process:
+        logger.info(f"{model}")
+        logger.info(f"Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
+        for n, p in model.named_parameters():
+            if p.requires_grad:
+                logger.info(f"Training Parameters: {n}")
+
+    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0)
+
+    # Setup data:
+    dataset = getdataset(args)
+    loader = DataLoader(
+        dataset,
+        batch_size=int(args.local_batch_size),
+        shuffle=False,
+        # sampler=sampler,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        collate_fn=Collate(args)
+    )
+    logger.info(f"Dataset contains {len(dataset):,} videos ({args.data_path})")
+
+    # Scheduler
+    lr_scheduler = get_scheduler(
+        name="constant",
+        optimizer=opt,
+        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+
+    # Prepare models for training:
+    if args.use_ema:
+        update_ema(ema, model, decay=0)  # Ensure EMA is initialized with synced weights
+        ema.eval()  # EMA model should always be in eval mode
+    model.train()  # important! This enables embedding dropout for classifier-free guidance
+    model, opt, loader, lr_scheduler = accelerator.prepare(model, opt, loader, lr_scheduler)
+
+    # Variables for monitoring/logging purposes:
+    train_steps = 0
+    log_steps = 0
+    running_loss = 0
+    first_epoch = 0
+    start_time = time()
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(loader))
+    # Afterwards we recalculate our number of training epochs
+    num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        # Get the most recent checkpoint
+        dirs = os.listdir(os.path.join(experiment_dir, 'checkpoints'))
+        dirs = [d for d in dirs if d.endswith("pt")]
+        dirs = sorted(dirs, key=lambda x: int(x.split(".")[0]))
+        path = dirs[-1]
+        state_dict = torch.load(os.path.join(os.path.join(experiment_dir, 'checkpoints'), path), map_location='cpu')
+        print(state_dict['args'])
+        state_dict = {'module.'+k: v for k, v in state_dict['model'].items()}
+        logger.info(f"Resuming from checkpoint {path}")
+        model.load_state_dict(state_dict)
+        train_steps = int(path.split(".")[0])
+
+        first_epoch = train_steps // num_update_steps_per_epoch
+        resume_step = train_steps % num_update_steps_per_epoch
+
+    if args.pretrained:
+        try:
+            train_steps = int(args.pretrained.split("/")[-1].split('.')[0])
+        except:
+            train_steps = 0
+
+    if accelerator.is_main_process:
+        logger.info(f"Training for {num_train_epochs} epochs...")
+
+
+    for epoch in range(first_epoch, num_train_epochs):
+        logger.info(f"Beginning epoch {epoch}...")
+        for step, (x, attn_mask, input_ids, cond_mask) in enumerate(loader):
+            if args.resume_from_checkpoint and epoch == first_epoch and step < resume_step:
+                continue
+            x = x.to(device)  # B C T H W
+            # attn_mask = attn_mask.to(device)  # B T H W
+            assert torch.all(attn_mask.bool()), 'do not enable dynamic input'
+            attn_mask = None
+            input_ids = input_ids.to(device)  # B L
+            cond_mask = cond_mask.to(device)  # B L
+
+            with torch.autocast(device_type='cuda', dtype=dtype):
+                with torch.no_grad():
+                    # Map input images to latent space + normalize latents
+                    x = ae.encode(x)  # B C T H W -> B T C H W
+                    cond = text_enc(input_ids, cond_mask)  # B L D
+                    cond = cond[:, None]  # B L D -> B 1 L D
+                    cond_mask = cond_mask.to(dtype)
+
+                # b, t, c, h, w = 1, 64, 4, 32, 32
+                # x = torch.rand(b, t, c, h, w).to(device)
+                # cond = torch.rand(b, 1, 120, 4096).to(device)
+                # cond_mask = torch.ones(b, 120).to(device)
+                # model_kwargs = dict(cond=cond, attn_mask=None, cond_mask=cond_mask)
+
+                model_kwargs = dict(cond=cond, attn_mask=attn_mask, cond_mask=cond_mask)
+                t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device)
+                loss_dict = diffusion.training_losses(model, x, t, model_kwargs)
+                loss = loss_dict["loss"].mean()
+                opt.zero_grad()
+
+            accelerator.backward(loss)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(model.parameters(), args.clip_grad_norm)
+            opt.step()
+            lr_scheduler.step()
+
+            if args.use_ema:
+                update_ema(ema, model.module)
+
+            # Log loss values:
+            running_loss += loss.item()
+            log_steps += 1
+            train_steps += 1
+            if train_steps % args.log_every == 0:
+                # Measure training speed:
+                torch.cuda.synchronize()
+                end_time = time()
+                steps_per_sec = log_steps / (end_time - start_time)
+                # Reduce loss history over all processes:
+                avg_loss = torch.tensor(running_loss / log_steps, device=device)
+                dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+                avg_loss = avg_loss.item() / accelerator.num_processes
+                # logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
+                if accelerator.is_main_process:
+                    # logger.info(f"(step={train_steps:07d}/epoch={epoch:04d}) Train Loss: {avg_loss:.4f}, Gradient Norm: {gradient_norm:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
+                    logger.info(
+                        f"(step={train_steps:07d}/epoch={epoch:04d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
+
+                write_tensorboard(tb_writer, 'Train Loss', avg_loss, train_steps)
+                # Reset monitoring variables:
+                running_loss = 0
+                log_steps = 0
+                start_time = time()
+
+            # Save Model checkpoint:
+            if train_steps % args.ckpt_every == 0 and train_steps > 0:
+                if accelerator.is_main_process:
+                    checkpoint = {
+                        "model": model.module.state_dict(),
+                        "opt": opt.state_dict(),
+                        "args": args
+                    }
+                    if args.use_ema:
+                        checkpoint.update({"ema": ema.state_dict()})
+                    checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+                    torch.save(checkpoint, checkpoint_path)
+                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+
+    model.eval()  # important! This disables randomized embedding dropout
+    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+
+    if accelerator.is_main_process:
+        logger.info("Done!")
+    cleanup()
+
+
+if __name__ == "__main__":
+    # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters).
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, required=True)
+    parser.add_argument("--data-path", type=str, required=True)
+    parser.add_argument("--results-dir", type=str, default="results")
+    parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="DiT-XL/122")
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--max-train-steps", type=int, default=1000000)
+    parser.add_argument("--local-batch-size", type=int, default=256)
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--num-workers", type=int, default=8)
+    parser.add_argument("--log-every", type=int, default=100)
+    parser.add_argument("--ckpt-every", type=int, default=50_000)
+
+
+    # --------------------------------------
+    parser.add_argument("--ae", type=str, default="ucf101_stride4x4x4")
+    parser.add_argument("--extras", type=int, default=3, choices=[3])
+    parser.add_argument("--pretrained", type=str, default=None)
+    parser.add_argument("--sample-rate", type=int, default=4)
+    parser.add_argument("--num-frames", type=int, default=16)
+    parser.add_argument("--max-image-size", type=int, default=128)
+    parser.add_argument("--dynamic-frames", action="store_true")
+    parser.add_argument("--resume-from-checkpoint", type=str, default=None)
+    parser.add_argument("--gradient-checkpointing", action="store_true")
+    parser.add_argument("--use-compile", action="store_true")
+    parser.add_argument("--use-ema", action="store_true")
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--lr-warmup-steps", type=int, default=0)
+    parser.add_argument("--attention-mode", type=str, choices=['xformers', 'math', 'flash'], default="math")
+    parser.add_argument("--gradient-accumulation-steps", type=int, default=1)
+    parser.add_argument("--mixed-precision", type=str, default=None, choices=[None, "fp16", "bf16"])
+    parser.add_argument("--clip-grad-norm", default=1.0, type=float, help="the maximum gradient norm (default None)")
+    # --------------------------------------
+
+    # --------------------------------------
+    parser.add_argument("--video-folder", type=str, default='')
+    parser.add_argument("--text-encoder-name", type=str, default='DeepFloyd/t5-v1_1-xxl')
+    parser.add_argument("--model-max-length", type=int, default=120)
+    # --------------------------------------
+
+    args = parser.parse_args()
+    main(args)

opensora/utils/dataset_utils.py

@@ -52,8 +52,16 @@ class Collate:
         self.num_frames = args.num_frames
 
     def __call__(self, batch):
-        batch_tubes, labels = tuple(zip(*batch))
-        labels = torch.as_tensor(labels).to(torch.long)
+        unzip = tuple(zip(*batch))
+        if len(unzip) == 2:
+            batch_tubes, labels = unzip
+            labels = torch.as_tensor(labels).to(torch.long)
+        elif len(unzip) == 3:
+            batch_tubes, input_ids, cond_mask = unzip
+            input_ids = torch.stack(input_ids).squeeze(1)
+            cond_mask = torch.stack(cond_mask).squeeze(1)
+        else:
+            raise NotImplementedError
         ds_stride = self.ae_stride * self.patch_size
         t_ds_stride = self.ae_stride_t * self.patch_size_t
 
@@ -91,5 +99,8 @@ class Collate:
                                  0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size]
         attention_mask = torch.stack(attention_mask)
 
-        return pad_batch_tubes, labels, attention_mask
+        if len(unzip) == 2:
+            return pad_batch_tubes, labels, attention_mask
+        elif len(unzip) == 3:
+            return pad_batch_tubes, attention_mask, input_ids, cond_mask
 

opensora/utils/utils.py

@@ -162,6 +162,14 @@ def get_experiment_dir(root_dir, args):
     root_dir += f'-{args.max_image_size}'
     return root_dir
 
+def get_precision(args):
+    if args.mixed_precision == "bf16":
+        dtype = torch.bfloat16
+    elif args.mixed_precision == "fp16":
+        dtype = torch.float16
+    else:
+        dtype = torch.float32
+    return dtype
 
 #################################################################################
 #                             Training Logger                                   #
@@ -325,56 +333,37 @@ def collect_env():
 #                          Pixart-alpha  Utils                                  #
 #################################################################################
 
-bad_punct_regex = re.compile(
-    r"[" + "#®•©™&@·º½¾¿¡§~" + "\)" + "\(" + "\]" + "\[" + "\}" + "\{" + "\|" + "\\" + "\/" + "\*" + r"]{1,}"
-)
 
+bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
 
-def text_preprocessing(text, clean_caption=False):
-    if clean_caption and not is_bs4_available():
-        clean_caption = False
+def text_preprocessing(text):
+    # The exact text cleaning as was in the training stage:
+    text = clean_caption(text)
+    text = clean_caption(text)
+    return text
 
-    if clean_caption and not is_ftfy_available():
-        clean_caption = False
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
 
-    if not isinstance(text, (tuple, list)):
-        text = [text]
-
-    def process(text: str):
-        if clean_caption:
-            text = clean_caption(text)
-            text = clean_caption(text)
-        else:
-            text = text.lower().strip()
-        return text
-
-    return [process(t) for t in text]
-
-
-# Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
 def clean_caption(caption):
     caption = str(caption)
     caption = ul.unquote_plus(caption)
     caption = caption.strip().lower()
-    caption = re.sub("<person>", "person", caption)
+    caption = re.sub('<person>', 'person', caption)
     # urls:
     caption = re.sub(
-        r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",
-        # noqa
-        "",
-        caption,
-    )  # regex for urls
+        r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+        '', caption)  # regex for urls
     caption = re.sub(
-        r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",
-        # noqa
-        "",
-        caption,
-    )  # regex for urls
+        r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
+        '', caption)  # regex for urls
     # html:
-    caption = BeautifulSoup(caption, features="html.parser").text
+    caption = BeautifulSoup(caption, features='html.parser').text
 
     # @<nickname>
-    caption = re.sub(r"@[\w\d]+\b", "", caption)
+    caption = re.sub(r'@[\w\d]+\b', '', caption)
 
     # 31C0—31EF CJK Strokes
     # 31F0—31FF Katakana Phonetic Extensions
@@ -383,94 +372,88 @@ def clean_caption(caption):
     # 3400—4DBF CJK Unified Ideographs Extension A
     # 4DC0—4DFF Yijing Hexagram Symbols
     # 4E00—9FFF CJK Unified Ideographs
-    caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
-    caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
-    caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
-    caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
-    caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
-    caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
-    caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
+    caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
+    caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
+    caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
+    caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
+    caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
+    caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
+    caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
     #######################################################
 
     # все виды тире / all types of dash --> "-"
     caption = re.sub(
-        r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",
-        # noqa
-        "-",
-        caption,
-    )
+        r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
+        '-', caption)
 
     # кавычки к одному стандарту
-    caption = re.sub(r"[`´«»“”¨]", '"', caption)
-    caption = re.sub(r"[‘’]", "'", caption)
+    caption = re.sub(r'[`´«»“”¨]', '"', caption)
+    caption = re.sub(r'[‘’]', "'", caption)
 
     # &quot;
-    caption = re.sub(r"&quot;?", "", caption)
+    caption = re.sub(r'&quot;?', '', caption)
     # &amp
-    caption = re.sub(r"&amp", "", caption)
+    caption = re.sub(r'&amp', '', caption)
 
     # ip adresses:
-    caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
+    caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
 
     # article ids:
-    caption = re.sub(r"\d:\d\d\s+$", "", caption)
+    caption = re.sub(r'\d:\d\d\s+$', '', caption)
 
     # \n
-    caption = re.sub(r"\\n", " ", caption)
+    caption = re.sub(r'\\n', ' ', caption)
 
     # "#123"
-    caption = re.sub(r"#\d{1,3}\b", "", caption)
+    caption = re.sub(r'#\d{1,3}\b', '', caption)
     # "#12345.."
-    caption = re.sub(r"#\d{5,}\b", "", caption)
+    caption = re.sub(r'#\d{5,}\b', '', caption)
     # "123456.."
-    caption = re.sub(r"\b\d{6,}\b", "", caption)
+    caption = re.sub(r'\b\d{6,}\b', '', caption)
     # filenames:
-    caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
+    caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
 
     #
-    caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
-    caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
+    caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
+    caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
 
-    caption = re.sub(bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
-    caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
+    caption = re.sub(bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
+    caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
 
     # this-is-my-cute-cat / this_is_my_cute_cat
-    regex2 = re.compile(r"(?:\-|\_)")
+    regex2 = re.compile(r'(?:\-|\_)')
     if len(re.findall(regex2, caption)) > 3:
-        caption = re.sub(regex2, " ", caption)
+        caption = re.sub(regex2, ' ', caption)
 
-    caption = ftfy.fix_text(caption)
-    caption = html.unescape(html.unescape(caption))
+    caption = basic_clean(caption)
 
-    caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
-    caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
-    caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
+    caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
+    caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
+    caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
 
-    caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
-    caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
-    caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
-    caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
-    caption = re.sub(r"\bpage\s+\d+\b", "", caption)
+    caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
+    caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
+    caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
+    caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
+    caption = re.sub(r'\bpage\s+\d+\b', '', caption)
 
-    caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
+    caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
 
-    caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
+    caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
 
-    caption = re.sub(r"\b\s+\:\s+", r": ", caption)
-    caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
-    caption = re.sub(r"\s+", " ", caption)
+    caption = re.sub(r'\b\s+\:\s+', r': ', caption)
+    caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
+    caption = re.sub(r'\s+', ' ', caption)
 
     caption.strip()
 
-    caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
-    caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
-    caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
-    caption = re.sub(r"^\.\S+$", "", caption)
+    caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
+    caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
+    caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
+    caption = re.sub(r'^\.\S+$', '', caption)
 
     return caption.strip()
 
 
 
 
-
-

pyproject.toml

@@ -20,7 +20,8 @@ dependencies = [
     "pydub==0.25.1", "pytorch-lightning==1.4.2", "pytorchvideo==0.1.5", "PyYAML==6.0.1", "regex==2023.12.25",
     "requests==2.31.0", "scikit-learn==1.3.2", "scipy==1.10.1", "six==1.16.0", "tensorboard==2.14.0", "test-tube==0.7.5",
     "timm==0.9.16", "torchdiffeq==0.2.3", "torchmetrics==0.5.0", "tqdm==4.66.2", "urllib3==2.2.1", "uvicorn==0.27.1",
-    "diffusers==0.24.0", "scikit-video==1.1.11", "imageio-ffmpeg==0.4.9"
+    "diffusers==0.24.0", "scikit-video==1.1.11", "imageio-ffmpeg==0.4.9", "sentencepiece==0.1.99, beautifulsoup4=4.12.3",
+    'ftfy==6.1.3'
 ]
 
 [project.optional-dependencies]