""" ein notation: b - batch n - sequence nt - text sequence nw - raw wave length d - dimension dt - dimension text """ from __future__ import annotations from pathlib import Path from random import random from random import sample as random_sample from functools import partial from itertools import zip_longest from collections import namedtuple from typing import Literal, Callable import jaxtyping from beartype import beartype import torch import torch.nn.functional as F from torch import nn, tensor, Tensor, from_numpy from torch.nn import Module, ModuleList, Sequential, Linear from torch.nn.utils.rnn import pad_sequence import torchaudio from torchaudio.functional import DB_to_amplitude from torchdiffeq import odeint import einx from einops.layers.torch import Rearrange from einops import rearrange, repeat, reduce, pack, unpack from x_transformers import ( Attention, FeedForward, RMSNorm, AdaptiveRMSNorm, ) from x_transformers.x_transformers import RotaryEmbedding import sys sys.path.insert(0, "/zhanghaomin/codes3/vocos-main/") from vocos import Vocos from transformers import AutoTokenizer from transformers import T5EncoderModel from transformers import EncodecModel, AutoProcessor import os import math import traceback import numpy as np from moviepy.editor import AudioFileClip, VideoFileClip from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection #import open_clip from transformers import AutoImageProcessor, AutoModel from PIL import Image import time import warnings warnings.filterwarnings("ignore") def normalize_wav(waveform): waveform = waveform - torch.mean(waveform) waveform = waveform / (torch.max(torch.abs(waveform[0, :])) + 1e-8) return waveform * 0.5 def read_frames_with_moviepy(video_path, max_frame_nums=None): try: clip = VideoFileClip(video_path) duration = clip.duration frames = [] for frame in clip.iter_frames(): frames.append(frame) except: print("Error read_frames_with_moviepy", video_path) traceback.print_exc() return None, None if max_frame_nums is not None: frames_idx = np.linspace(0, len(frames) - 1, max_frame_nums, dtype=int) return np.array(frames)[frames_idx, ...], duration else: return np.array(frames), duration pad_sequence = partial(pad_sequence, batch_first = True) # constants class TorchTyping: def __init__(self, abstract_dtype): self.abstract_dtype = abstract_dtype def __getitem__(self, shapes: str): return self.abstract_dtype[Tensor, shapes] Float = TorchTyping(jaxtyping.Float) Int = TorchTyping(jaxtyping.Int) Bool = TorchTyping(jaxtyping.Bool) # named tuples LossBreakdown = namedtuple('LossBreakdown', ['flow', 'velocity_consistency']) E2TTSReturn = namedtuple('E2TTS', ['loss', 'cond', 'pred_flow', 'pred_data', 'loss_breakdown']) # helpers def exists(v): return v is not None def default(v, d): return v if exists(v) else d def divisible_by(num, den): return (num % den) == 0 def pack_one_with_inverse(x, pattern): packed, packed_shape = pack([x], pattern) def inverse(x, inverse_pattern = None): inverse_pattern = default(inverse_pattern, pattern) return unpack(x, packed_shape, inverse_pattern)[0] return packed, inverse class Identity(Module): def forward(self, x, **kwargs): return x # tensor helpers def project(x, y): x, inverse = pack_one_with_inverse(x, 'b *') y, _ = pack_one_with_inverse(y, 'b *') dtype = x.dtype x, y = x.double(), y.double() unit = F.normalize(y, dim = -1) parallel = (x * unit).sum(dim = -1, keepdim = True) * unit orthogonal = x - parallel return inverse(parallel).to(dtype), inverse(orthogonal).to(dtype) # simple utf-8 tokenizer, since paper went character based def list_str_to_tensor( text: list[str], padding_value = -1 ) -> Int['b nt']: list_tensors = [tensor([*bytes(t, 'UTF-8')]) for t in text] padded_tensor = pad_sequence(list_tensors, padding_value = -1) return padded_tensor # simple english phoneme-based tokenizer from g2p_en import G2p import jieba from pypinyin import lazy_pinyin, Style def get_g2p_en_encode(): g2p = G2p() # used by @lucasnewman successfully here # https://github.com/lucasnewman/e2-tts-pytorch/blob/ljspeech-test/e2_tts_pytorch/e2_tts.py phoneme_to_index = g2p.p2idx num_phonemes = len(phoneme_to_index) extended_chars = [' ', ',', '.', '-', '!', '?', '\'', '"', '...', '..', '. .', '. . .', '. . . .', '. . . . .', '. ...', '... .', '.. ..'] num_extended_chars = len(extended_chars) extended_chars_dict = {p: (num_phonemes + i) for i, p in enumerate(extended_chars)} phoneme_to_index = {**phoneme_to_index, **extended_chars_dict} def encode( text: list[str], padding_value = -1 ) -> Int['b nt']: phonemes = [g2p(t) for t in text] list_tensors = [tensor([phoneme_to_index[p] for p in one_phoneme]) for one_phoneme in phonemes] padded_tensor = pad_sequence(list_tensors, padding_value = -1) return padded_tensor return encode, (num_phonemes + num_extended_chars) def all_en(word): res = word.replace("'", "").encode('utf-8').isalpha() return res def all_ch(word): res = True for w in word: if not '\u4e00' <= w <= '\u9fff': res = False return res def get_g2p_zh_encode(): puncs = ['，', '。', '？', '、'] pinyins = ['a', 'a1', 'ai1', 'ai2', 'ai3', 'ai4', 'an1', 'an3', 'an4', 'ang1', 'ang2', 'ang4', 'ao1', 'ao2', 'ao3', 'ao4', 'ba', 'ba1', 'ba2', 'ba3', 'ba4', 'bai1', 'bai2', 'bai3', 'bai4', 'ban1', 'ban2', 'ban3', 'ban4', 'bang1', 'bang2', 'bang3', 'bang4', 'bao1', 'bao2', 'bao3', 'bao4', 'bei', 'bei1', 'bei2', 'bei3', 'bei4', 'ben1', 'ben2', 'ben3', 'ben4', 'beng1', 'beng2', 'beng4', 'bi1', 'bi2', 'bi3', 'bi4', 'bian1', 'bian2', 'bian3', 'bian4', 'biao1', 'biao2', 'biao3', 'bie1', 'bie2', 'bie3', 'bie4', 'bin1', 'bin4', 'bing1', 'bing2', 'bing3', 'bing4', 'bo', 'bo1', 'bo2', 'bo3', 'bo4', 'bu2', 'bu3', 'bu4', 'ca1', 'cai1', 'cai2', 'cai3', 'cai4', 'can1', 'can2', 'can3', 'can4', 'cang1', 'cang2', 'cao1', 'cao2', 'cao3', 'ce4', 'cen1', 'cen2', 'ceng1', 'ceng2', 'ceng4', 'cha1', 'cha2', 'cha3', 'cha4', 'chai1', 'chai2', 'chan1', 'chan2', 'chan3', 'chan4', 'chang1', 'chang2', 'chang3', 'chang4', 'chao1', 'chao2', 'chao3', 'che1', 'che2', 'che3', 'che4', 'chen1', 'chen2', 'chen3', 'chen4', 'cheng1', 'cheng2', 'cheng3', 'cheng4', 'chi1', 'chi2', 'chi3', 'chi4', 'chong1', 'chong2', 'chong3', 'chong4', 'chou1', 'chou2', 'chou3', 'chou4', 'chu1', 'chu2', 'chu3', 'chu4', 'chua1', 'chuai1', 'chuai2', 'chuai3', 'chuai4', 'chuan1', 'chuan2', 'chuan3', 'chuan4', 'chuang1', 'chuang2', 'chuang3', 'chuang4', 'chui1', 'chui2', 'chun1', 'chun2', 'chun3', 'chuo1', 'chuo4', 'ci1', 'ci2', 'ci3', 'ci4', 'cong1', 'cong2', 'cou4', 'cu1', 'cu4', 'cuan1', 'cuan2', 'cuan4', 'cui1', 'cui3', 'cui4', 'cun1', 'cun2', 'cun4', 'cuo1', 'cuo2', 'cuo4', 'da', 'da1', 'da2', 'da3', 'da4', 'dai1', 'dai3', 'dai4', 'dan1', 'dan2', 'dan3', 'dan4', 'dang1', 'dang2', 'dang3', 'dang4', 'dao1', 'dao2', 'dao3', 'dao4', 'de', 'de1', 'de2', 'dei3', 'den4', 'deng1', 'deng2', 'deng3', 'deng4', 'di1', 'di2', 'di3', 'di4', 'dia3', 'dian1', 'dian2', 'dian3', 'dian4', 'diao1', 'diao3', 'diao4', 'die1', 'die2', 'ding1', 'ding2', 'ding3', 'ding4', 'diu1', 'dong1', 'dong3', 'dong4', 'dou1', 'dou2', 'dou3', 'dou4', 'du1', 'du2', 'du3', 'du4', 'duan1', 'duan2', 'duan3', 'duan4', 'dui1', 'dui4', 'dun1', 'dun3', 'dun4', 'duo1', 'duo2', 'duo3', 'duo4', 'e1', 'e2', 'e3', 'e4', 'ei2', 'en1', 'en4', 'er', 'er2', 'er3', 'er4', 'fa1', 'fa2', 'fa3', 'fa4', 'fan1', 'fan2', 'fan3', 'fan4', 'fang1', 'fang2', 'fang3', 'fang4', 'fei1', 'fei2', 'fei3', 'fei4', 'fen1', 'fen2', 'fen3', 'fen4', 'feng1', 'feng2', 'feng3', 'feng4', 'fo2', 'fou2', 'fou3', 'fu1', 'fu2', 'fu3', 'fu4', 'ga1', 'ga2', 'ga4', 'gai1', 'gai3', 'gai4', 'gan1', 'gan2', 'gan3', 'gan4', 'gang1', 'gang2', 'gang3', 'gang4', 'gao1', 'gao2', 'gao3', 'gao4', 'ge1', 'ge2', 'ge3', 'ge4', 'gei2', 'gei3', 'gen1', 'gen2', 'gen3', 'gen4', 'geng1', 'geng3', 'geng4', 'gong1', 'gong3', 'gong4', 'gou1', 'gou2', 'gou3', 'gou4', 'gu', 'gu1', 'gu2', 'gu3', 'gu4', 'gua1', 'gua2', 'gua3', 'gua4', 'guai1', 'guai2', 'guai3', 'guai4', 'guan1', 'guan2', 'guan3', 'guan4', 'guang1', 'guang2', 'guang3', 'guang4', 'gui1', 'gui2', 'gui3', 'gui4', 'gun3', 'gun4', 'guo1', 'guo2', 'guo3', 'guo4', 'ha1', 'ha2', 'ha3', 'hai1', 'hai2', 'hai3', 'hai4', 'han1', 'han2', 'han3', 'han4', 'hang1', 'hang2', 'hang4', 'hao1', 'hao2', 'hao3', 'hao4', 'he1', 'he2', 'he4', 'hei1', 'hen2', 'hen3', 'hen4', 'heng1', 'heng2', 'heng4', 'hong1', 'hong2', 'hong3', 'hong4', 'hou1', 'hou2', 'hou3', 'hou4', 'hu1', 'hu2', 'hu3', 'hu4', 'hua1', 'hua2', 'hua4', 'huai2', 'huai4', 'huan1', 'huan2', 'huan3', 'huan4', 'huang1', 'huang2', 'huang3', 'huang4', 'hui1', 'hui2', 'hui3', 'hui4', 'hun1', 'hun2', 'hun4', 'huo', 'huo1', 'huo2', 'huo3', 'huo4', 'ji1', 'ji2', 'ji3', 'ji4', 'jia', 'jia1', 'jia2', 'jia3', 'jia4', 'jian1', 'jian2', 'jian3', 'jian4', 'jiang1', 'jiang2', 'jiang3', 'jiang4', 'jiao1', 'jiao2', 'jiao3', 'jiao4', 'jie1', 'jie2', 'jie3', 'jie4', 'jin1', 'jin2', 'jin3', 'jin4', 'jing1', 'jing2', 'jing3', 'jing4', 'jiong3', 'jiu1', 'jiu2', 'jiu3', 'jiu4', 'ju1', 'ju2', 'ju3', 'ju4', 'juan1', 'juan2', 'juan3', 'juan4', 'jue1', 'jue2', 'jue4', 'jun1', 'jun4', 'ka1', 'ka2', 'ka3', 'kai1', 'kai2', 'kai3', 'kai4', 'kan1', 'kan2', 'kan3', 'kan4', 'kang1', 'kang2', 'kang4', 'kao2', 'kao3', 'kao4', 'ke1', 'ke2', 'ke3', 'ke4', 'ken3', 'keng1', 'kong1', 'kong3', 'kong4', 'kou1', 'kou2', 'kou3', 'kou4', 'ku1', 'ku2', 'ku3', 'ku4', 'kua1', 'kua3', 'kua4', 'kuai3', 'kuai4', 'kuan1', 'kuan2', 'kuan3', 'kuang1', 'kuang2', 'kuang4', 'kui1', 'kui2', 'kui3', 'kui4', 'kun1', 'kun3', 'kun4', 'kuo4', 'la', 'la1', 'la2', 'la3', 'la4', 'lai2', 'lai4', 'lan2', 'lan3', 'lan4', 'lang1', 'lang2', 'lang3', 'lang4', 'lao1', 'lao2', 'lao3', 'lao4', 'le', 'le1', 'le4', 'lei', 'lei1', 'lei2', 'lei3', 'lei4', 'leng1', 'leng2', 'leng3', 'leng4', 'li', 'li1', 'li2', 'li3', 'li4', 'lia3', 'lian2', 'lian3', 'lian4', 'liang2', 'liang3', 'liang4', 'liao1', 'liao2', 'liao3', 'liao4', 'lie1', 'lie2', 'lie3', 'lie4', 'lin1', 'lin2', 'lin3', 'lin4', 'ling2', 'ling3', 'ling4', 'liu1', 'liu2', 'liu3', 'liu4', 'long1', 'long2', 'long3', 'long4', 'lou1', 'lou2', 'lou3', 'lou4', 'lu1', 'lu2', 'lu3', 'lu4', 'luan2', 'luan3', 'luan4', 'lun1', 'lun2', 'lun4', 'luo1', 'luo2', 'luo3', 'luo4', 'lv2', 'lv3', 'lv4', 'lve3', 'lve4', 'ma', 'ma1', 'ma2', 'ma3', 'ma4', 'mai2', 'mai3', 'mai4', 'man2', 'man3', 'man4', 'mang2', 'mang3', 'mao1', 'mao2', 'mao3', 'mao4', 'me', 'mei2', 'mei3', 'mei4', 'men', 'men1', 'men2', 'men4', 'meng1', 'meng2', 'meng3', 'meng4', 'mi1', 'mi2', 'mi3', 'mi4', 'mian2', 'mian3', 'mian4', 'miao1', 'miao2', 'miao3', 'miao4', 'mie1', 'mie4', 'min2', 'min3', 'ming2', 'ming3', 'ming4', 'miu4', 'mo1', 'mo2', 'mo3', 'mo4', 'mou1', 'mou2', 'mou3', 'mu2', 'mu3', 'mu4', 'n2', 'na1', 'na2', 'na3', 'na4', 'nai2', 'nai3', 'nai4', 'nan1', 'nan2', 'nan3', 'nan4', 'nang1', 'nang2', 'nao1', 'nao2', 'nao3', 'nao4', 'ne', 'ne2', 'ne4', 'nei3', 'nei4', 'nen4', 'neng2', 'ni1', 'ni2', 'ni3', 'ni4', 'nian1', 'nian2', 'nian3', 'nian4', 'niang2', 'niang4', 'niao2', 'niao3', 'niao4', 'nie1', 'nie4', 'nin2', 'ning2', 'ning3', 'ning4', 'niu1', 'niu2', 'niu3', 'niu4', 'nong2', 'nong4', 'nou4', 'nu2', 'nu3', 'nu4', 'nuan3', 'nuo2', 'nuo4', 'nv2', 'nv3', 'nve4', 'o1', 'o2', 'ou1', 'ou3', 'ou4', 'pa1', 'pa2', 'pa4', 'pai1', 'pai2', 'pai3', 'pai4', 'pan1', 'pan2', 'pan4', 'pang1', 'pang2', 'pang4', 'pao1', 'pao2', 'pao3', 'pao4', 'pei1', 'pei2', 'pei4', 'pen1', 'pen2', 'pen4', 'peng1', 'peng2', 'peng3', 'peng4', 'pi1', 'pi2', 'pi3', 'pi4', 'pian1', 'pian2', 'pian4', 'piao1', 'piao2', 'piao3', 'piao4', 'pie1', 'pie2', 'pie3', 'pin1', 'pin2', 'pin3', 'pin4', 'ping1', 'ping2', 'po1', 'po2', 'po3', 'po4', 'pou1', 'pu1', 'pu2', 'pu3', 'pu4', 'qi1', 'qi2', 'qi3', 'qi4', 'qia1', 'qia3', 'qia4', 'qian1', 'qian2', 'qian3', 'qian4', 'qiang1', 'qiang2', 'qiang3', 'qiang4', 'qiao1', 'qiao2', 'qiao3', 'qiao4', 'qie1', 'qie2', 'qie3', 'qie4', 'qin1', 'qin2', 'qin3', 'qin4', 'qing1', 'qing2', 'qing3', 'qing4', 'qiong1', 'qiong2', 'qiu1', 'qiu2', 'qiu3', 'qu1', 'qu2', 'qu3', 'qu4', 'quan1', 'quan2', 'quan3', 'quan4', 'que1', 'que2', 'que4', 'qun2', 'ran2', 'ran3', 'rang1', 'rang2', 'rang3', 'rang4', 'rao2', 'rao3', 'rao4', 're2', 're3', 're4', 'ren2', 'ren3', 'ren4', 'reng1', 'reng2', 'ri4', 'rong1', 'rong2', 'rong3', 'rou2', 'rou4', 'ru2', 'ru3', 'ru4', 'ruan2', 'ruan3', 'rui3', 'rui4', 'run4', 'ruo4', 'sa1', 'sa2', 'sa3', 'sa4', 'sai1', 'sai4', 'san1', 'san2', 'san3', 'san4', 'sang1', 'sang3', 'sang4', 'sao1', 'sao2', 'sao3', 'sao4', 'se4', 'sen1', 'seng1', 'sha1', 'sha2', 'sha3', 'sha4', 'shai1', 'shai2', 'shai3', 'shai4', 'shan1', 'shan3', 'shan4', 'shang', 'shang1', 'shang3', 'shang4', 'shao1', 'shao2', 'shao3', 'shao4', 'she1', 'she2', 'she3', 'she4', 'shei2', 'shen1', 'shen2', 'shen3', 'shen4', 'sheng1', 'sheng2', 'sheng3', 'sheng4', 'shi', 'shi1', 'shi2', 'shi3', 'shi4', 'shou1', 'shou2', 'shou3', 'shou4', 'shu1', 'shu2', 'shu3', 'shu4', 'shua1', 'shua2', 'shua3', 'shua4', 'shuai1', 'shuai3', 'shuai4', 'shuan1', 'shuan4', 'shuang1', 'shuang3', 'shui2', 'shui3', 'shui4', 'shun3', 'shun4', 'shuo1', 'shuo4', 'si1', 'si2', 'si3', 'si4', 'song1', 'song3', 'song4', 'sou1', 'sou3', 'sou4', 'su1', 'su2', 'su4', 'suan1', 'suan4', 'sui1', 'sui2', 'sui3', 'sui4', 'sun1', 'sun3', 'suo', 'suo1', 'suo2', 'suo3', 'ta1', 'ta3', 'ta4', 'tai1', 'tai2', 'tai4', 'tan1', 'tan2', 'tan3', 'tan4', 'tang1', 'tang2', 'tang3', 'tang4', 'tao1', 'tao2', 'tao3', 'tao4', 'te4', 'teng2', 'ti1', 'ti2', 'ti3', 'ti4', 'tian1', 'tian2', 'tian3', 'tiao1', 'tiao2', 'tiao3', 'tiao4', 'tie1', 'tie2', 'tie3', 'tie4', 'ting1', 'ting2', 'ting3', 'tong1', 'tong2', 'tong3', 'tong4', 'tou', 'tou1', 'tou2', 'tou4', 'tu1', 'tu2', 'tu3', 'tu4', 'tuan1', 'tuan2', 'tui1', 'tui2', 'tui3', 'tui4', 'tun1', 'tun2', 'tun4', 'tuo1', 'tuo2', 'tuo3', 'tuo4', 'wa', 'wa1', 'wa2', 'wa3', 'wa4', 'wai1', 'wai3', 'wai4', 'wan1', 'wan2', 'wan3', 'wan4', 'wang1', 'wang2', 'wang3', 'wang4', 'wei1', 'wei2', 'wei3', 'wei4', 'wen1', 'wen2', 'wen3', 'wen4', 'weng1', 'weng4', 'wo1', 'wo3', 'wo4', 'wu1', 'wu2', 'wu3', 'wu4', 'xi1', 'xi2', 'xi3', 'xi4', 'xia1', 'xia2', 'xia4', 'xian1', 'xian2', 'xian3', 'xian4', 'xiang1', 'xiang2', 'xiang3', 'xiang4', 'xiao1', 'xiao2', 'xiao3', 'xiao4', 'xie1', 'xie2', 'xie3', 'xie4', 'xin1', 'xin2', 'xin4', 'xing1', 'xing2', 'xing3', 'xing4', 'xiong1', 'xiong2', 'xiu1', 'xiu3', 'xiu4', 'xu', 'xu1', 'xu2', 'xu3', 'xu4', 'xuan1', 'xuan2', 'xuan3', 'xuan4', 'xue1', 'xue2', 'xue3', 'xue4', 'xun1', 'xun2', 'xun4', 'ya', 'ya1', 'ya2', 'ya3', 'ya4', 'yan1', 'yan2', 'yan3', 'yan4', 'yang1', 'yang2', 'yang3', 'yang4', 'yao1', 'yao2', 'yao3', 'yao4', 'ye1', 'ye2', 'ye3', 'ye4', 'yi1', 'yi2', 'yi3', 'yi4', 'yin1', 'yin2', 'yin3', 'yin4', 'ying1', 'ying2', 'ying3', 'ying4', 'yo1', 'yong1', 'yong3', 'yong4', 'you1', 'you2', 'you3', 'you4', 'yu1', 'yu2', 'yu3', 'yu4', 'yuan1', 'yuan2', 'yuan3', 'yuan4', 'yue1', 'yue4', 'yun1', 'yun2', 'yun3', 'yun4', 'za1', 'za2', 'za3', 'zai1', 'zai3', 'zai4', 'zan1', 'zan2', 'zan3', 'zan4', 'zang1', 'zang4', 'zao1', 'zao2', 'zao3', 'zao4', 'ze2', 'ze4', 'zei2', 'zen3', 'zeng1', 'zeng4', 'zha1', 'zha2', 'zha3', 'zha4', 'zhai1', 'zhai2', 'zhai3', 'zhai4', 'zhan1', 'zhan2', 'zhan3', 'zhan4', 'zhang1', 'zhang2', 'zhang3', 'zhang4', 'zhao1', 'zhao2', 'zhao3', 'zhao4', 'zhe', 'zhe1', 'zhe2', 'zhe3', 'zhe4', 'zhen1', 'zhen2', 'zhen3', 'zhen4', 'zheng1', 'zheng2', 'zheng3', 'zheng4', 'zhi1', 'zhi2', 'zhi3', 'zhi4', 'zhong1', 'zhong2', 'zhong3', 'zhong4', 'zhou1', 'zhou2', 'zhou3', 'zhou4', 'zhu1', 'zhu2', 'zhu3', 'zhu4', 'zhua1', 'zhua2', 'zhua3', 'zhuai1', 'zhuai3', 'zhuai4', 'zhuan1', 'zhuan2', 'zhuan3', 'zhuan4', 'zhuang1', 'zhuang4', 'zhui1', 'zhui4', 'zhun1', 'zhun2', 'zhun3', 'zhuo1', 'zhuo2', 'zi', 'zi1', 'zi2', 'zi3', 'zi4', 'zong1', 'zong2', 'zong3', 'zong4', 'zou1', 'zou2', 'zou3', 'zou4', 'zu1', 'zu2', 'zu3', 'zuan1', 'zuan3', 'zuan4', 'zui2', 'zui3', 'zui4', 'zun1', 'zuo1', 'zuo2', 'zuo3', 'zuo4'] ens = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', "'", ' '] ens_U = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z'] phoneme_to_index = {} num_phonemes = 0 for index, punc in enumerate(puncs): phoneme_to_index[punc] = index + num_phonemes num_phonemes += len(puncs) for index, pinyin in enumerate(pinyins): phoneme_to_index[pinyin] = index + num_phonemes num_phonemes += len(pinyins) for index, en in enumerate(ens): phoneme_to_index[en] = index + num_phonemes for index, en in enumerate(ens_U): phoneme_to_index[en] = index + num_phonemes num_phonemes += len(ens) #print(num_phonemes, phoneme_to_index) # 1342 def encode( text: list[str], padding_value = -1 ) -> Int['b nt']: phonemes = [] for t in text: one_phoneme = [] brk = False for word in jieba.cut(t): if all_ch(word): seg = lazy_pinyin(word, style=Style.TONE3, tone_sandhi=True) one_phoneme.extend(seg) elif all_en(word): for seg in word: one_phoneme.append(seg) elif word in ["，", "。", "？", "、", "'", " "]: one_phoneme.append(word) else: for ch in word: if all_ch(ch): seg = lazy_pinyin(ch, style=Style.TONE3, tone_sandhi=True) one_phoneme.extend(seg) elif all_en(ch): for seg in ch: one_phoneme.append(seg) else: brk = True break if brk: break if not brk: phonemes.append(one_phoneme) else: print("Error Tokenized", t, list(jieba.cut(t))) list_tensors = [tensor([phoneme_to_index[p] for p in one_phoneme]) for one_phoneme in phonemes] padded_tensor = pad_sequence(list_tensors, padding_value = -1) return padded_tensor return encode, num_phonemes # tensor helpers def log(t, eps = 1e-5): return t.clamp(min = eps).log() def lens_to_mask( t: Int['b'], length: int | None = None ) -> Bool['b n']: if not exists(length): length = t.amax() seq = torch.arange(length, device = t.device) return einx.less('n, b -> b n', seq, t) def mask_from_start_end_indices( seq_len: Int['b'], start: Int['b'], end: Int['b'] ): max_seq_len = seq_len.max().item() seq = torch.arange(max_seq_len, device = start.device).long() return einx.greater_equal('n, b -> b n', seq, start) & einx.less('n, b -> b n', seq, end) def mask_from_frac_lengths( seq_len: Int['b'], frac_lengths: Float['b'], max_length: int | None = None, val = False ): lengths = (frac_lengths * seq_len).long() max_start = seq_len - lengths if not val: rand = torch.rand_like(frac_lengths) else: rand = torch.tensor([0.5]*frac_lengths.shape[0], device=frac_lengths.device).float() start = (max_start * rand).long().clamp(min = 0) end = start + lengths out = mask_from_start_end_indices(seq_len, start, end) if exists(max_length): out = pad_to_length(out, max_length) return out def maybe_masked_mean( t: Float['b n d'], mask: Bool['b n'] | None = None ) -> Float['b d']: if not exists(mask): return t.mean(dim = 1) t = einx.where('b n, b n d, -> b n d', mask, t, 0.) num = reduce(t, 'b n d -> b d', 'sum') den = reduce(mask.float(), 'b n -> b', 'sum') return einx.divide('b d, b -> b d', num, den.clamp(min = 1.)) def pad_to_length( t: Tensor, length: int, value = None ): seq_len = t.shape[-1] if length > seq_len: t = F.pad(t, (0, length - seq_len), value = value) return t[..., :length] def interpolate_1d( x: Tensor, length: int, mode = 'bilinear' ): x = rearrange(x, 'n d -> 1 d n 1') x = F.interpolate(x, (length, 1), mode = mode) return rearrange(x, '1 d n 1 -> n d') # to mel spec class MelSpec(Module): def __init__( self, filter_length = 1024, hop_length = 256, win_length = 1024, n_mel_channels = 100, sampling_rate = 24_000, normalize = False, power = 1, norm = None, center = True, ): super().__init__() self.n_mel_channels = n_mel_channels self.sampling_rate = sampling_rate self.mel_stft = torchaudio.transforms.MelSpectrogram( sample_rate = sampling_rate, n_fft = filter_length, win_length = win_length, hop_length = hop_length, n_mels = n_mel_channels, power = power, center = center, normalized = normalize, norm = norm, ) self.register_buffer('dummy', tensor(0), persistent = False) def forward(self, inp): if len(inp.shape) == 3: inp = rearrange(inp, 'b 1 nw -> b nw') assert len(inp.shape) == 2 if self.dummy.device != inp.device: self.to(inp.device) mel = self.mel_stft(inp) mel = log(mel) return mel class EncodecWrapper(Module): def __init__(self, path): super().__init__() self.model = EncodecModel.from_pretrained(path) self.processor = AutoProcessor.from_pretrained(path) for param in self.model.parameters(): param.requires_grad = False self.model.eval() def forward(self, waveform): with torch.no_grad(): inputs = self.processor(raw_audio=waveform[0], sampling_rate=self.processor.sampling_rate, return_tensors="pt") emb = self.model.encoder(inputs.input_values) return emb def decode(self, emb): with torch.no_grad(): output = self.model.decoder(emb) return output[0] from audioldm.audio.stft import TacotronSTFT from audioldm.variational_autoencoder import AutoencoderKL from audioldm.utils import default_audioldm_config, get_metadata def build_pretrained_models(name): checkpoint = torch.load(get_metadata()[name]["path"], map_location="cpu") scale_factor = checkpoint["state_dict"]["scale_factor"].item() vae_state_dict = {k[18:]: v for k, v in checkpoint["state_dict"].items() if "first_stage_model." in k} config = default_audioldm_config(name) vae_config = config["model"]["params"]["first_stage_config"]["params"] vae_config["scale_factor"] = scale_factor vae = AutoencoderKL(**vae_config) vae.load_state_dict(vae_state_dict) fn_STFT = TacotronSTFT( config["preprocessing"]["stft"]["filter_length"], config["preprocessing"]["stft"]["hop_length"], config["preprocessing"]["stft"]["win_length"], config["preprocessing"]["mel"]["n_mel_channels"], config["preprocessing"]["audio"]["sampling_rate"], config["preprocessing"]["mel"]["mel_fmin"], config["preprocessing"]["mel"]["mel_fmax"], ) vae.eval() fn_STFT.eval() return vae, fn_STFT class VaeWrapper(Module): def __init__(self): super().__init__() vae, stft = build_pretrained_models("audioldm-s-full") vae.eval() stft.eval() stft = stft.cpu() self.vae = vae for param in self.vae.parameters(): param.requires_grad = False def forward(self, waveform): return None def decode(self, emb): with torch.no_grad(): b, d, l = emb.shape latents = emb.transpose(1,2).reshape(b, l, 8, 16).transpose(1,2) mel = self.vae.decode_first_stage(latents) wave = self.vae.decode_to_waveform(mel) return wave # convolutional positional generating module # taken from https://github.com/lucidrains/voicebox-pytorch/blob/main/voicebox_pytorch/voicebox_pytorch.py#L203 class DepthwiseConv(Module): def __init__( self, dim, *, kernel_size, groups = None ): super().__init__() assert not divisible_by(kernel_size, 2) groups = default(groups, dim) # full depthwise conv by default self.dw_conv1d = nn.Sequential( nn.Conv1d(dim, dim, kernel_size, groups = groups, padding = kernel_size // 2), nn.SiLU() ) def forward( self, x, mask = None ): if exists(mask): x = einx.where('b n, b n d, -> b n d', mask, x, 0.) x = rearrange(x, 'b n c -> b c n') x = self.dw_conv1d(x) out = rearrange(x, 'b c n -> b n c') if exists(mask): out = einx.where('b n, b n d, -> b n d', mask, out, 0.) return out # adaln zero from DiT paper class AdaLNZero(Module): def __init__( self, dim, dim_condition = None, init_bias_value = -2. ): super().__init__() dim_condition = default(dim_condition, dim) self.to_gamma = nn.Linear(dim_condition, dim) nn.init.zeros_(self.to_gamma.weight) nn.init.constant_(self.to_gamma.bias, init_bias_value) def forward(self, x, *, condition): if condition.ndim == 2: condition = rearrange(condition, 'b d -> b 1 d') gamma = self.to_gamma(condition).sigmoid() return x * gamma # random projection fourier embedding class RandomFourierEmbed(Module): def __init__(self, dim): super().__init__() assert divisible_by(dim, 2) self.register_buffer('weights', torch.randn(dim // 2)) def forward(self, x): freqs = einx.multiply('i, j -> i j', x, self.weights) * 2 * torch.pi fourier_embed, _ = pack((x, freqs.sin(), freqs.cos()), 'b *') return fourier_embed # character embedding class CharacterEmbed(Module): def __init__( self, dim, num_embeds = 256, ): super().__init__() self.dim = dim self.embed = nn.Embedding(num_embeds + 1, dim) # will just use 0 as the 'filler token' def forward( self, text: Int['b nt'], max_seq_len: int, **kwargs ) -> Float['b n d']: text = text + 1 # shift all other token ids up by 1 and use 0 as filler token text = text[:, :max_seq_len] # just curtail if character tokens are more than the mel spec tokens, one of the edge cases the paper did not address text = pad_to_length(text, max_seq_len, value = 0) return self.embed(text) class InterpolatedCharacterEmbed(Module): def __init__( self, dim, num_embeds = 256, ): super().__init__() self.dim = dim self.embed = nn.Embedding(num_embeds, dim) self.abs_pos_mlp = Sequential( Rearrange('... -> ... 1'), Linear(1, dim), nn.SiLU(), Linear(dim, dim) ) def forward( self, text: Int['b nt'], max_seq_len: int, mask: Bool['b n'] | None = None ) -> Float['b n d']: device = text.device mask = default(mask, (None,)) interp_embeds = [] interp_abs_positions = [] for one_text, one_mask in zip_longest(text, mask): valid_text = one_text >= 0 one_text = one_text[valid_text] one_text_embed = self.embed(one_text) # save the absolute positions text_seq_len = one_text.shape[0] # determine audio sequence length from mask audio_seq_len = max_seq_len if exists(one_mask): audio_seq_len = one_mask.sum().long().item() # interpolate text embedding to audio embedding length interp_text_embed = interpolate_1d(one_text_embed, audio_seq_len) interp_abs_pos = torch.linspace(0, text_seq_len, audio_seq_len, device = device) interp_embeds.append(interp_text_embed) interp_abs_positions.append(interp_abs_pos) interp_embeds = pad_sequence(interp_embeds) interp_abs_positions = pad_sequence(interp_abs_positions) interp_embeds = F.pad(interp_embeds, (0, 0, 0, max_seq_len - interp_embeds.shape[-2])) interp_abs_positions = pad_to_length(interp_abs_positions, max_seq_len) # pass interp absolute positions through mlp for implicit positions interp_embeds = interp_embeds + self.abs_pos_mlp(interp_abs_positions) if exists(mask): interp_embeds = einx.where('b n, b n d, -> b n d', mask, interp_embeds, 0.) return interp_embeds # text audio cross conditioning in multistream setup class TextAudioCrossCondition(Module): def __init__( self, dim, dim_text, cond_audio_to_text = True, ): super().__init__() self.text_to_audio = nn.Linear(dim_text + dim, dim, bias = False) nn.init.zeros_(self.text_to_audio.weight) self.cond_audio_to_text = cond_audio_to_text if cond_audio_to_text: self.audio_to_text = nn.Linear(dim + dim_text, dim_text, bias = False) nn.init.zeros_(self.audio_to_text.weight) def forward( self, audio: Float['b n d'], text: Float['b n dt'] ): audio_text, _ = pack((audio, text), 'b n *') text_cond = self.text_to_audio(audio_text) audio_cond = self.audio_to_text(audio_text) if self.cond_audio_to_text else 0. return audio + text_cond, text + audio_cond # attention and transformer backbone # for use in both e2tts as well as duration module class Transformer(Module): @beartype def __init__( self, *, dim, dim_text = None, # will default to half of audio dimension depth = 8, heads = 8, dim_head = 64, ff_mult = 4, text_depth = None, text_heads = None, text_dim_head = None, text_ff_mult = None, cond_on_time = True, abs_pos_emb = True, max_seq_len = 8192, kernel_size = 31, dropout = 0.1, num_registers = 32, attn_kwargs: dict = dict( gate_value_heads = True, softclamp_logits = True, ), ff_kwargs: dict = dict(), if_text_modules = True, if_cross_attn = True, if_audio_conv = True, if_text_conv = False ): super().__init__() assert divisible_by(depth, 2), 'depth needs to be even' # absolute positional embedding self.max_seq_len = max_seq_len self.abs_pos_emb = nn.Embedding(max_seq_len, dim) if abs_pos_emb else None self.dim = dim dim_text = default(dim_text, dim // 2) self.dim_text = dim_text text_heads = default(text_heads, heads) text_dim_head = default(text_dim_head, dim_head) text_ff_mult = default(text_ff_mult, ff_mult) text_depth = default(text_depth, depth) assert 1 <= text_depth <= depth, 'must have at least 1 layer of text conditioning, but less than total number of speech layers' self.depth = depth self.layers = ModuleList([]) # registers self.num_registers = num_registers self.registers = nn.Parameter(torch.zeros(num_registers, dim)) nn.init.normal_(self.registers, std = 0.02) if if_text_modules: self.text_registers = nn.Parameter(torch.zeros(num_registers, dim_text)) nn.init.normal_(self.text_registers, std = 0.02) # rotary embedding self.rotary_emb = RotaryEmbedding(dim_head) self.text_rotary_emb = RotaryEmbedding(dim_head) # time conditioning # will use adaptive rmsnorm self.cond_on_time = cond_on_time rmsnorm_klass = RMSNorm if not cond_on_time else AdaptiveRMSNorm postbranch_klass = Identity if not cond_on_time else partial(AdaLNZero, dim = dim) self.time_cond_mlp = Identity() if cond_on_time: self.time_cond_mlp = Sequential( RandomFourierEmbed(dim), Linear(dim + 1, dim), nn.SiLU() ) for ind in range(depth): is_later_half = ind >= (depth // 2) has_text = ind < text_depth # speech related if if_audio_conv: speech_conv = DepthwiseConv(dim, kernel_size = kernel_size) attn_norm = rmsnorm_klass(dim) attn = Attention(dim = dim, heads = heads, dim_head = dim_head, dropout = dropout, **attn_kwargs) attn_adaln_zero = postbranch_klass() if if_cross_attn: attn_norm2 = rmsnorm_klass(dim) attn2 = Attention(dim = dim, heads = heads, dim_head = dim_head, dropout = dropout, **attn_kwargs) attn_adaln_zero2 = postbranch_klass() ff_norm = rmsnorm_klass(dim) ff = FeedForward(dim = dim, glu = True, mult = ff_mult, dropout = dropout, **ff_kwargs) ff_adaln_zero = postbranch_klass() skip_proj = Linear(dim * 2, dim, bias = False) if is_later_half else None if if_cross_attn: if if_audio_conv: speech_modules = ModuleList([ skip_proj, speech_conv, attn_norm, attn, attn_adaln_zero, attn_norm2, attn2, attn_adaln_zero2, ff_norm, ff, ff_adaln_zero, ]) else: speech_modules = ModuleList([ skip_proj, attn_norm, attn, attn_adaln_zero, attn_norm2, attn2, attn_adaln_zero2, ff_norm, ff, ff_adaln_zero, ]) else: if if_audio_conv: speech_modules = ModuleList([ skip_proj, speech_conv, attn_norm, attn, attn_adaln_zero, ff_norm, ff, ff_adaln_zero, ]) else: speech_modules = ModuleList([ skip_proj, attn_norm, attn, attn_adaln_zero, ff_norm, ff, ff_adaln_zero, ]) text_modules = None if has_text and if_text_modules: # text related if if_text_conv: text_conv = DepthwiseConv(dim_text, kernel_size = kernel_size) text_attn_norm = RMSNorm(dim_text) text_attn = Attention(dim = dim_text, heads = text_heads, dim_head = text_dim_head, dropout = dropout, **attn_kwargs) text_ff_norm = RMSNorm(dim_text) text_ff = FeedForward(dim = dim_text, glu = True, mult = text_ff_mult, dropout = dropout, **ff_kwargs) # cross condition is_last = ind == (text_depth - 1) cross_condition = TextAudioCrossCondition(dim = dim, dim_text = dim_text, cond_audio_to_text = not is_last) if if_text_conv: text_modules = ModuleList([ text_conv, text_attn_norm, text_attn, text_ff_norm, text_ff, cross_condition ]) else: text_modules = ModuleList([ text_attn_norm, text_attn, text_ff_norm, text_ff, cross_condition ]) self.layers.append(ModuleList([ speech_modules, text_modules ])) self.final_norm = RMSNorm(dim) self.if_cross_attn = if_cross_attn self.if_audio_conv = if_audio_conv self.if_text_conv = if_text_conv ####self.a_align = Linear(dim, 512) ####self.v_align = Linear(dim_text, 512) ####self.contrastive_loss = SupConLoss() #self.contrastive_loss = FactorCLSSL(None, [1024, 1280], None) self.contrastive_loss = FactorCLSUP(None, [1024, 1280], 6) def forward( self, x: Float['b n d'], times: Float['b'] | Float[''] | None = None, mask: Bool['b n'] | None = None, text_embed: Float['b n dt'] | None = None, context: Float['b nc dc'] | None = None, context_mask: Float['b nc'] | None = None ): batch, seq_len, device = *x.shape[:2], x.device assert not (exists(times) ^ self.cond_on_time), '`times` must be passed in if `cond_on_time` is set to `True` and vice versa' # handle absolute positions if needed if exists(self.abs_pos_emb): assert seq_len <= self.max_seq_len, f'{seq_len} exceeds the set `max_seq_len` ({self.max_seq_len}) on Transformer' seq = torch.arange(seq_len, device = device) x = x + self.abs_pos_emb(seq) # handle adaptive rmsnorm kwargs norm_kwargs = dict() if exists(times): if times.ndim == 0: times = repeat(times, ' -> b', b = batch) times = self.time_cond_mlp(times) norm_kwargs.update(condition = times) # register tokens registers = repeat(self.registers, 'r d -> b r d', b = batch) x, registers_packed_shape = pack((registers, x), 'b * d') if exists(mask): mask = F.pad(mask, (self.num_registers, 0), value = True) # rotary embedding rotary_pos_emb = self.rotary_emb.forward_from_seq_len(x.shape[-2]) # text related if exists(text_embed): text_rotary_pos_emb = self.text_rotary_emb.forward_from_seq_len(x.shape[-2]) text_registers = repeat(self.text_registers, 'r d -> b r d', b = batch) text_embed, _ = pack((text_registers, text_embed), 'b * d') # skip connection related stuff skips = [] # go through the layers loss_contra = 0 for ind, (speech_modules, text_modules) in enumerate(self.layers): layer = ind + 1 if self.if_cross_attn: if self.if_audio_conv: ( maybe_skip_proj, speech_conv, attn_norm, attn, maybe_attn_adaln_zero, attn_norm2, attn2, maybe_attn_adaln_zero2, ff_norm, ff, maybe_ff_adaln_zero ) = speech_modules else: ( maybe_skip_proj, attn_norm, attn, maybe_attn_adaln_zero, attn_norm2, attn2, maybe_attn_adaln_zero2, ff_norm, ff, maybe_ff_adaln_zero ) = speech_modules else: if self.if_audio_conv: ( maybe_skip_proj, speech_conv, attn_norm, attn, maybe_attn_adaln_zero, ff_norm, ff, maybe_ff_adaln_zero ) = speech_modules else: ( maybe_skip_proj, attn_norm, attn, maybe_attn_adaln_zero, ff_norm, ff, maybe_ff_adaln_zero ) = speech_modules # smaller text transformer if exists(text_embed) and exists(text_modules): if self.if_text_conv: ( text_conv, text_attn_norm, text_attn, text_ff_norm, text_ff, cross_condition ) = text_modules else: ( text_attn_norm, text_attn, text_ff_norm, text_ff, cross_condition ) = text_modules if self.if_text_conv: text_embed = text_conv(text_embed, mask = mask) + text_embed text_embed = text_attn(text_attn_norm(text_embed), rotary_pos_emb = text_rotary_pos_emb, mask = mask) + text_embed text_embed = text_ff(text_ff_norm(text_embed)) + text_embed if x.shape[0] >= 8: if layer in [1]: #print("contrastive learning", x.shape, text_embed.shape) # [16, 782, 1024] [16, 782, 1280] ####features1 = self.a_align(x) ####features2 = self.v_align(text_embed) features1 = x features2 = text_embed #b1, b2 = random_sample(range(2,8), 2) #t1, t2 = random_sample(range(x.shape[1]), 2) #loss_contra = torch.cosine_similarity(x[b1,t1:t1+1,:], y[b1,t2:t2+1,:]) + torch.cosine_similarity(x[b1,t1:t1+1,:], y[b2,t1:t1+1,:]) - torch.cosine_similarity(x[b1,t1:t1+1,:], y[b1,t1:t1+1,:]) - torch.cosine_similarity(x[b2,t2:t2+1,:], y[b2,t2:t2+1,:]) features1 = features1[2:8,32:,:] features2 = features2[2:8,32:,:] if self.training: ts = random_sample(range(features1.shape[1]), 1) #ts = random_sample(range(features1.shape[1]), 6) #ts = random_sample(range(features1.shape[1]), 14) else: ts = [350] #ts = [100, 200, 300, 400, 500, 600] #ts = [50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700] features1 = features1[:,ts,:] features2 = features2[:,ts,:] ####loss_contra += self.contrastive_loss(features1, features2, labels=None, mask=None) nf = 1 features1 = features1.reshape(6*nf, -1) features2 = features2.reshape(6*nf, -1) label = torch.tensor([0.0]*nf+[1.0]*nf+[2.0]*nf+[3.0]*nf+[4.0]*nf+[5.0]*nf, device=device).reshape(6*nf,1) loss_contra += self.contrastive_loss(features1, features2, label) else: loss_contra = torch.tensor(0.0, device=device) x, text_embed = cross_condition(x, text_embed) # skip connection logic is_first_half = layer <= (self.depth // 2) is_later_half = not is_first_half if is_first_half: skips.append(x) if is_later_half: skip = skips.pop() x = torch.cat((x, skip), dim = -1) x = maybe_skip_proj(x) # position generating convolution if self.if_audio_conv: x = speech_conv(x, mask = mask) + x # attention and feedforward blocks attn_out = attn(attn_norm(x, **norm_kwargs), rotary_pos_emb = rotary_pos_emb, mask = mask) x = x + maybe_attn_adaln_zero(attn_out, **norm_kwargs) if self.if_cross_attn: attn_out = attn2(attn_norm2(x, **norm_kwargs), rotary_pos_emb = rotary_pos_emb, mask = mask, context = context, context_mask = context_mask) x = x + maybe_attn_adaln_zero2(attn_out, **norm_kwargs) ff_out = ff(ff_norm(x, **norm_kwargs)) x = x + maybe_ff_adaln_zero(ff_out, **norm_kwargs) assert len(skips) == 0 _, x = unpack(x, registers_packed_shape, 'b * d') return self.final_norm(x), loss_contra from .multibench_model import FactorCLSUP, FactorCLSSL class SupConLoss(nn.Module): """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf. It also supports the unsupervised contrastive loss in SimCLR""" def __init__(self, temperature=0.07, contrast_mode='all', base_temperature=0.07): super(SupConLoss, self).__init__() self.temperature = temperature self.contrast_mode = contrast_mode self.base_temperature = base_temperature def forward(self, features, features2, labels=None, mask=None): """Compute loss for model. If both `labels` and `mask` are None, it degenerates to SimCLR unsupervised loss: https://arxiv.org/pdf/2002.05709.pdf Args: features: hidden vector of shape [bsz, n_views, ...]. labels: ground truth of shape [bsz]. mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j has the same class as sample i. Can be asymmetric. Returns: A loss scalar. """ features = F.normalize(features, dim=2) features2 = F.normalize(features2, dim=2) device = (torch.device('cuda') if features.is_cuda else torch.device('cpu')) if len(features.shape) < 3: raise ValueError('`features` needs to be [bsz, n_views, ...],' 'at least 3 dimensions are required') if len(features.shape) > 3: features = features.view(features.shape[0], features.shape[1], -1) features2 = features2.view(features2.shape[0], features2.shape[1], -1) batch_size = features.shape[0] if labels is not None and mask is not None: raise ValueError('Cannot define both `labels` and `mask`') elif labels is None and mask is None: mask = torch.eye(batch_size, dtype=torch.float32).to(device) elif labels is not None: labels = labels.contiguous().view(-1, 1) if labels.shape[0] != batch_size: raise ValueError('Num of labels does not match num of features') mask = torch.eq(labels, labels.T).float().to(device) else: mask = mask.float().to(device) contrast_count = features.shape[1] contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0) contrast_feature2 = torch.cat(torch.unbind(features2, dim=1), dim=0) if self.contrast_mode == 'one': anchor_feature = features[:, 0] anchor_feature2 = features2[:, 0] anchor_count = 1 elif self.contrast_mode == 'all': anchor_feature = contrast_feature anchor_feature2 = contrast_feature2 anchor_count = contrast_count else: raise ValueError('Unknown mode: {}'.format(self.contrast_mode)) # compute logits anchor_dot_contrast = torch.div( #torch.matmul(anchor_feature, contrast_feature.T), torch.matmul(anchor_feature, contrast_feature2.T), self.temperature) # for numerical stability logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True) logits = anchor_dot_contrast - logits_max.detach() # tile mask mask = mask.repeat(anchor_count, contrast_count) # mask-out self-contrast cases #logits_mask = torch.scatter( # torch.ones_like(mask), # 1, # torch.arange(batch_size * anchor_count).view(-1, 1).to(device), # 0 #) logits_mask = torch.ones_like(mask) mask = mask * logits_mask #print("logits", logits.shape, logits.min(), logits.max()) # compute log_prob exp_logits = torch.exp(logits) * logits_mask #print("exp_logits", exp_logits.shape, exp_logits.min(), exp_logits.max()) log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True)) #print("log_prob", log_prob.shape, log_prob.min(), log_prob.max()) # compute mean of log-likelihood over positive # modified to handle edge cases when there is no positive pair # for an anchor point. # Edge case e.g.:- # features of shape: [4,1,...] # labels: [0,1,1,2] # loss before mean: [nan, ..., ..., nan] mask_pos_pairs = mask.sum(1) mask_pos_pairs = torch.where(mask_pos_pairs < 1e-6, 1, mask_pos_pairs) mean_log_prob_pos = (mask * log_prob).sum(1) / mask_pos_pairs # loss loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos loss = loss.view(anchor_count, batch_size).mean() return loss # main classes class DurationPredictor(Module): @beartype def __init__( self, transformer: dict | Transformer, num_channels = None, mel_spec_kwargs: dict = dict(), char_embed_kwargs: dict = dict(), text_num_embeds = None, tokenizer: ( Literal['char_utf8', 'phoneme_en'] | Callable[[list[str]], Int['b nt']] ) = 'char_utf8' ): super().__init__() if isinstance(transformer, dict): transformer = Transformer( **transformer, cond_on_time = False ) # mel spec self.mel_spec = MelSpec(**mel_spec_kwargs) self.num_channels = default(num_channels, self.mel_spec.n_mel_channels) self.transformer = transformer dim = transformer.dim dim_text = transformer.dim_text self.dim = dim self.proj_in = Linear(self.num_channels, self.dim) # tokenizer and text embed if callable(tokenizer): assert exists(text_num_embeds), '`text_num_embeds` must be given if supplying your own tokenizer encode function' self.tokenizer = tokenizer elif tokenizer == 'char_utf8': text_num_embeds = 256 self.tokenizer = list_str_to_tensor elif tokenizer == 'phoneme_en': self.tokenizer, text_num_embeds = get_g2p_en_encode() elif tokenizer == 'phoneme_zh': self.tokenizer, text_num_embeds = get_g2p_zh_encode() else: raise ValueError(f'unknown tokenizer string {tokenizer}') self.embed_text = CharacterEmbed(dim_text, num_embeds = text_num_embeds, **char_embed_kwargs) # to prediction self.to_pred = Sequential( Linear(dim, 1, bias = False), nn.Softplus(), Rearrange('... 1 -> ...') ) def forward( self, x: Float['b n d'] | Float['b nw'], *, text: Int['b nt'] | list[str] | None = None, lens: Int['b'] | None = None, return_loss = True ): # raw wave if x.ndim == 2: x = self.mel_spec(x) x = rearrange(x, 'b d n -> b n d') assert x.shape[-1] == self.dim x = self.proj_in(x) batch, seq_len, device = *x.shape[:2], x.device # text text_embed = None if exists(text): if isinstance(text, list): text = list_str_to_tensor(text).to(device) assert text.shape[0] == batch text_embed = self.embed_text(text, seq_len) # handle lengths (duration) if not exists(lens): lens = torch.full((batch,), seq_len, device = device) mask = lens_to_mask(lens, length = seq_len) # if returning a loss, mask out randomly from an index and have it predict the duration if return_loss: rand_frac_index = x.new_zeros(batch).uniform_(0, 1) rand_index = (rand_frac_index * lens).long() seq = torch.arange(seq_len, device = device) mask &= einx.less('n, b -> b n', seq, rand_index) # attending x = self.transformer( x, mask = mask, text_embed = text_embed, ) x = maybe_masked_mean(x, mask) pred = self.to_pred(x) # return the prediction if not returning loss if not return_loss: return pred # loss return F.mse_loss(pred, lens.float()) class E2TTS(Module): @beartype def __init__( self, transformer: dict | Transformer = None, duration_predictor: dict | DurationPredictor | None = None, odeint_kwargs: dict = dict( #atol = 1e-5, #rtol = 1e-5, #method = 'midpoint' method = "euler" ), audiocond_drop_prob = 0.30, cond_drop_prob = 0.20, prompt_drop_prob = 0.10, num_channels = None, mel_spec_module: Module | None = None, char_embed_kwargs: dict = dict(), mel_spec_kwargs: dict = dict(), frac_lengths_mask: tuple[float, float] = (0.7, 1.), audiocond_snr: tuple[float, float] | None = None, concat_cond = False, interpolated_text = False, text_num_embeds: int | None = None, tokenizer: ( Literal['char_utf8', 'phoneme_en', 'phoneme_zh'] | Callable[[list[str]], Int['b nt']] ) = 'char_utf8', use_vocos = True, pretrained_vocos_path = 'charactr/vocos-mel-24khz', sampling_rate: int | None = None, frame_size: int = 320, velocity_consistency_weight = -1e-5, if_cond_proj_in = True, cond_proj_in_bias = True, if_embed_text = True, if_text_encoder2 = True, if_clip_encoder = False, video_encoder = "clip_vit" ): super().__init__() if isinstance(transformer, dict): transformer = Transformer( **transformer, cond_on_time = True ) if isinstance(duration_predictor, dict): duration_predictor = DurationPredictor(**duration_predictor) self.transformer = transformer dim = transformer.dim dim_text = transformer.dim_text self.dim = dim self.dim_text = dim_text self.frac_lengths_mask = frac_lengths_mask self.audiocond_snr = audiocond_snr self.duration_predictor = duration_predictor # sampling self.odeint_kwargs = odeint_kwargs # mel spec self.mel_spec = default(mel_spec_module, None) num_channels = default(num_channels, None) self.num_channels = num_channels self.sampling_rate = default(sampling_rate, None) self.frame_size = frame_size # whether to concat condition and project rather than project both and sum self.concat_cond = concat_cond if concat_cond: self.proj_in = nn.Linear(num_channels * 2, dim) else: self.proj_in = nn.Linear(num_channels, dim) self.cond_proj_in = nn.Linear(num_channels, dim, bias=cond_proj_in_bias) if if_cond_proj_in else None # to prediction self.to_pred = Linear(dim, num_channels) # tokenizer and text embed if callable(tokenizer): assert exists(text_num_embeds), '`text_num_embeds` must be given if supplying your own tokenizer encode function' self.tokenizer = tokenizer elif tokenizer == 'char_utf8': text_num_embeds = 256 self.tokenizer = list_str_to_tensor elif tokenizer == 'phoneme_en': self.tokenizer, text_num_embeds = get_g2p_en_encode() elif tokenizer == 'phoneme_zh': self.tokenizer, text_num_embeds = get_g2p_zh_encode() else: raise ValueError(f'unknown tokenizer string {tokenizer}') self.audiocond_drop_prob = audiocond_drop_prob self.cond_drop_prob = cond_drop_prob self.prompt_drop_prob = prompt_drop_prob # text embedding text_embed_klass = CharacterEmbed if not interpolated_text else InterpolatedCharacterEmbed self.embed_text = text_embed_klass(dim_text, num_embeds = text_num_embeds, **char_embed_kwargs) if if_embed_text else None # weight for velocity consistency self.register_buffer('zero', torch.tensor(0.), persistent = False) self.velocity_consistency_weight = velocity_consistency_weight # default vocos for mel -> audio if pretrained_vocos_path == 'charactr/vocos-mel-24khz': self.vocos = Vocos.from_pretrained(pretrained_vocos_path) if use_vocos else None elif pretrained_vocos_path == 'facebook/encodec_24khz': self.vocos = EncodecWrapper("facebook/encodec_24khz") if use_vocos else None elif pretrained_vocos_path == 'vae': self.vocos = VaeWrapper() if use_vocos else None if if_text_encoder2: self.tokenizer2 = AutoTokenizer.from_pretrained("/ailab-train/speech/zhanghaomin/models/flan-t5-large") self.text_encoder2 = T5EncoderModel.from_pretrained("/ailab-train/speech/zhanghaomin/models/flan-t5-large") for param in self.text_encoder2.parameters(): param.requires_grad = False self.text_encoder2.eval() self.proj_text = None if if_clip_encoder: if video_encoder == "clip_vit": self.image_processor = CLIPImageProcessor() #self.image_encoder = CLIPVisionModelWithProjection.from_pretrained("/ailab-train/speech/zhanghaomin/models/IP-Adapter/", subfolder="models/image_encoder") self.image_encoder = CLIPVisionModelWithProjection.from_pretrained("/ailab-train/speech/zhanghaomin/models/IP-Adapter/", subfolder="sdxl_models/image_encoder") elif video_encoder == "clip_vit2": self.image_processor = AutoProcessor.from_pretrained("/ailab-train/speech/zhanghaomin/models/clip-vit-large-patch14-336/") self.image_encoder = CLIPVisionModelWithProjection.from_pretrained("/ailab-train/speech/zhanghaomin/models/clip-vit-large-patch14-336/") elif video_encoder == "clip_convnext": self.image_encoder, _, self.image_processor = open_clip.create_model_and_transforms("hf-hub:laion/CLIP-convnext_xxlarge-laion2B-s34B-b82K-augreg-soup") elif video_encoder == "dinov2": self.image_processor = AutoImageProcessor.from_pretrained("/ailab-train/speech/zhanghaomin/models/dinov2-giant/") self.image_encoder = AutoModel.from_pretrained("/ailab-train/speech/zhanghaomin/models/dinov2-giant/") elif video_encoder == "mixed": #pass self.image_processor1 = CLIPImageProcessor() self.image_encoder1 = CLIPVisionModelWithProjection.from_pretrained("/ailab-train/speech/zhanghaomin/models/IP-Adapter/", subfolder="sdxl_models/image_encoder") self.image_processor2 = AutoProcessor.from_pretrained("/ailab-train/speech/zhanghaomin/models/clip-vit-large-patch14-336/") self.image_encoder2 = CLIPVisionModelWithProjection.from_pretrained("/ailab-train/speech/zhanghaomin/models/clip-vit-large-patch14-336/") self.image_encoder3, _, self.image_processor3 = open_clip.create_model_and_transforms("hf-hub:laion/CLIP-convnext_xxlarge-laion2B-s34B-b82K-augreg-soup") self.image_processor4 = AutoImageProcessor.from_pretrained("/ailab-train/speech/zhanghaomin/models/dinov2-giant/") self.image_encoder4 = AutoModel.from_pretrained("/ailab-train/speech/zhanghaomin/models/dinov2-giant/") else: self.image_processor = None self.image_encoder = None if video_encoder != "mixed": for param in self.image_encoder.parameters(): param.requires_grad = False self.image_encoder.eval() else: for param in self.image_encoder1.parameters(): param.requires_grad = False self.image_encoder1.eval() for param in self.image_encoder2.parameters(): param.requires_grad = False self.image_encoder2.eval() for param in self.image_encoder3.parameters(): param.requires_grad = False self.image_encoder3.eval() for param in self.image_encoder4.parameters(): param.requires_grad = False self.image_encoder4.eval() self.dim_text_raw = 4608 self.proj_text = Linear(self.dim_text_raw, dim_text) self.video_encoder = video_encoder for param in self.vocos.parameters(): param.requires_grad = False self.vocos.eval() @property def device(self): return next(self.parameters()).device def encode_text(self, prompt): device = self.device batch = self.tokenizer2(prompt, max_length=self.tokenizer2.model_max_length, padding=True, truncation=True, return_tensors="pt") input_ids, attention_mask = batch.input_ids.to(device), batch.attention_mask.to(device) with torch.no_grad(): encoder_hidden_states = self.text_encoder2(input_ids=input_ids, attention_mask=attention_mask)[0] boolean_encoder_mask = (attention_mask == 1).to(device) return encoder_hidden_states, boolean_encoder_mask def encode_video(self, video_paths, l): if self.proj_text is None: d = self.dim_text else: d = self.dim_text_raw device = self.device b = 20 with torch.no_grad(): video_embeddings = [] video_lens = [] for video_path in video_paths: if video_path is None: video_embeddings.append(None) video_lens.append(0) continue if isinstance(video_path, tuple): video_path, start_sample, max_sample = video_path else: start_sample = 0 max_sample = None if video_path.startswith("/ailab-train/speech/zhanghaomin/VGGSound/"): if self.video_encoder == "clip_vit": feature_path = video_path.replace("/video/", "/feature/").replace(".mp4", ".npz") elif self.video_encoder == "clip_vit2": feature_path = video_path.replace("/video/", "/feature_clip_vit2/").replace(".mp4", ".npz") elif self.video_encoder == "clip_convnext": feature_path = video_path.replace("/video/", "/feature_clip_convnext/").replace(".mp4", ".npz") elif self.video_encoder == "dinov2": feature_path = video_path.replace("/video/", "/feature_dinov2/").replace(".mp4", ".npz") elif self.video_encoder == "mixed": feature_path = video_path.replace("/video/", "/feature_mixed/").replace(".mp4", ".npz") else: raise Exception("Invalid video_encoder " + self.video_encoder) else: if self.video_encoder == "clip_vit": feature_path = video_path.replace(".mp4", ".generated.npz") elif self.video_encoder == "clip_vit2": feature_path = video_path.replace(".mp4", ".generated.clip_vit2.npz") elif self.video_encoder == "clip_convnext": feature_path = video_path.replace(".mp4", ".generated.clip_convnext.npz") elif self.video_encoder == "dinov2": feature_path = video_path.replace(".mp4", ".generated.dinov2.npz") elif self.video_encoder == "mixed": feature_path = video_path.replace(".mp4", ".generated.mixed.npz") else: raise Exception("Invalid video_encoder " + self.video_encoder) if not os.path.exists(feature_path): #print("video not exist", video_path) frames, duration = read_frames_with_moviepy(video_path, max_frame_nums=None) if frames is None: video_embeddings.append(None) video_lens.append(0) continue if self.video_encoder in ["clip_vit", "clip_vit2", "dinov2"]: images = self.image_processor(images=frames, return_tensors="pt").to(device) #print("images", images["pixel_values"].shape, images["pixel_values"].max(), images["pixel_values"].min(), torch.abs(images["pixel_values"]).mean()) elif self.video_encoder in ["clip_convnext"]: images = [] for i in range(frames.shape[0]): images.append(self.image_processor(Image.fromarray(frames[i])).unsqueeze(0)) images = torch.cat(images, dim=0).to(device) #print("images", images.shape, images.max(), images.min(), torch.abs(images).mean()) elif self.video_encoder in ["mixed"]: #images1 = self.image_processor1(images=frames, return_tensors="pt").to(device) images2 = self.image_processor2(images=frames, return_tensors="pt").to(device) images4 = self.image_processor4(images=frames, return_tensors="pt").to(device) images3 = [] for i in range(frames.shape[0]): images3.append(self.image_processor3(Image.fromarray(frames[i])).unsqueeze(0)) images3 = torch.cat(images3, dim=0).to(device) else: raise Exception("Invalid video_encoder " + self.video_encoder) image_embeddings = [] if self.video_encoder == "clip_vit": for i in range(math.ceil(images["pixel_values"].shape[0] / b)): image_embeddings.append(self.image_encoder(pixel_values=images["pixel_values"][i*b: (i+1)*b]).image_embeds.cpu()) elif self.video_encoder == "clip_vit2": for i in range(math.ceil(images["pixel_values"].shape[0] / b)): image_embeddings.append(self.image_encoder(pixel_values=images["pixel_values"][i*b: (i+1)*b]).image_embeds.cpu()) elif self.video_encoder == "clip_convnext": for i in range(math.ceil(images.shape[0] / b)): image_embeddings.append(self.image_encoder.encode_image(images[i*b: (i+1)*b]).cpu()) elif self.video_encoder == "dinov2": for i in range(math.ceil(images["pixel_values"].shape[0] / b)): image_embeddings.append(self.image_encoder(pixel_values=images["pixel_values"][i*b: (i+1)*b]).pooler_output.cpu()) elif self.video_encoder == "mixed": feature_path1 = feature_path.replace("/feature_mixed/", "/feature/") if not os.path.exists(feature_path1): image_embeddings1 = [] for i in range(math.ceil(images1["pixel_values"].shape[0] / b)): image_embeddings1.append(self.image_encoder1(pixel_values=images1["pixel_values"][i*b: (i+1)*b]).image_embeds.cpu()) image_embeddings1 = torch.cat(image_embeddings1, dim=0) #np.savez(feature_path1, image_embeddings1, duration) else: data1 = np.load(feature_path1) image_embeddings1 = torch.from_numpy(data1["arr_0"]) feature_path2 = feature_path.replace("/feature_mixed/", "/feature_clip_vit2/") if not os.path.exists(feature_path2): image_embeddings2 = [] for i in range(math.ceil(images2["pixel_values"].shape[0] / b)): image_embeddings2.append(self.image_encoder2(pixel_values=images2["pixel_values"][i*b: (i+1)*b]).image_embeds.cpu()) image_embeddings2 = torch.cat(image_embeddings2, dim=0) np.savez(feature_path2, image_embeddings2, duration) else: data2 = np.load(feature_path2) image_embeddings2 = torch.from_numpy(data2["arr_0"]) feature_path3 = feature_path.replace("/feature_mixed/", "/feature_clip_convnext/") if not os.path.exists(feature_path3): image_embeddings3 = [] for i in range(math.ceil(images3.shape[0] / b)): image_embeddings3.append(self.image_encoder3.encode_image(images3[i*b: (i+1)*b]).cpu()) image_embeddings3 = torch.cat(image_embeddings3, dim=0) np.savez(feature_path3, image_embeddings3, duration) else: data3 = np.load(feature_path3) image_embeddings3 = torch.from_numpy(data3["arr_0"]) feature_path4 = feature_path.replace("/feature_mixed/", "/feature_dinov2/") if not os.path.exists(feature_path4): image_embeddings4 = [] for i in range(math.ceil(images4["pixel_values"].shape[0] / b)): image_embeddings4.append(self.image_encoder4(pixel_values=images4["pixel_values"][i*b: (i+1)*b]).pooler_output.cpu()) image_embeddings4 = torch.cat(image_embeddings4, dim=0) np.savez(feature_path4, image_embeddings4, duration) else: data4 = np.load(feature_path4) image_embeddings4 = torch.from_numpy(data4["arr_0"]) mixed_l = min([image_embeddings1.shape[0], image_embeddings2.shape[0], image_embeddings3.shape[0], image_embeddings4.shape[0]]) for i in range(mixed_l): image_embeddings.append(torch.cat([image_embeddings1[i:i+1,:], image_embeddings2[i:i+1,:], image_embeddings3[i:i+1,:], image_embeddings4[i:i+1,:]], dim=1)) else: raise Exception("Invalid video_encoder " + self.video_encoder) image_embeddings = torch.cat(image_embeddings, dim=0) #print("image_embeddings", image_embeddings.shape, image_embeddings.max(), image_embeddings.min(), torch.abs(image_embeddings).mean()) np.savez(feature_path, image_embeddings, duration) else: #print("video exist", feature_path) data = np.load(feature_path) image_embeddings = torch.from_numpy(data["arr_0"]) #print("image_embeddings", image_embeddings.shape, image_embeddings.max(), image_embeddings.min(), torch.abs(image_embeddings).mean()) duration = data["arr_1"].item() if max_sample is None: max_sample = int(duration * self.sampling_rate) interpolated = [] for i in range(start_sample, max_sample, self.frame_size): j = min(round((i+self.frame_size//2) / self.sampling_rate / (duration / (image_embeddings.shape[0] - 1))), image_embeddings.shape[0] - 1) interpolated.append(image_embeddings[j:j+1]) if len(interpolated) >= l: break interpolated = torch.cat(interpolated, dim=0) #ll = list(range(start_sample, max_sample, self.frame_size)) #print("encode_video l", len(ll), l, round((ll[-1]+self.frame_size//2) / self.sampling_rate / (duration / (image_embeddings.shape[0] - 1))), image_embeddings.shape[0] - 1) #print("encode_video one", video_path, duration, image_embeddings.shape, interpolated.shape, l) video_embeddings.append(interpolated.unsqueeze(0)) video_lens.append(interpolated.shape[1]) max_length = max(video_lens) if max_length == 0: max_length = l else: max_length = l for i in range(len(video_embeddings)): if video_embeddings[i] is None: video_embeddings[i] = torch.zeros(1, max_length, d) continue if video_embeddings[i].shape[1] < max_length: video_embeddings[i] = torch.cat([video_embeddings[i], torch.zeros(1, max_length-video_embeddings[i].shape[1], d)], 1) video_embeddings = torch.cat(video_embeddings, 0) #print("encode_video", l, video_embeddings.shape, video_lens) return video_embeddings.to(device) def transformer_with_pred_head( self, x: Float['b n d'], cond: Float['b n d'] | None = None, times: Float['b'] | None = None, mask: Bool['b n'] | None = None, text: Int['b nt'] | Float['b nt dt'] | None = None, prompt = None, video_drop_prompt = None, audio_drop_prompt = None, drop_audio_cond: bool | None = None, drop_text_cond: bool | None = None, drop_text_prompt: bool | None = None, return_drop_conditions = False ): seq_len = x.shape[-2] bs = x.shape[0] drop_audio_cond = [default(drop_audio_cond, self.training and random() < self.audiocond_drop_prob) for _ in range(bs)] drop_text_cond = default(drop_text_cond, self.training and random() < self.cond_drop_prob) drop_text_prompt = [default(drop_text_prompt, self.training and random() < self.prompt_drop_prob) for _ in range(bs)] if cond is not None: for b in range(bs): if drop_audio_cond[b]: cond[b] = 0 if audio_drop_prompt is not None and audio_drop_prompt[b]: cond[b] = 0 if cond is not None: if self.concat_cond: # concat condition, given as using voicebox-like scheme x = torch.cat((cond, x), dim = -1) x = self.proj_in(x) if cond is not None: if not self.concat_cond: # an alternative is to simply sum the condition # seems to work fine cond = self.cond_proj_in(cond) x = x + cond # whether to use a text embedding text_embed = None if exists(text) and len(text.shape) == 3: text_embed = text.clone() if drop_text_cond: for b in range(bs): text_embed[b] = 0 elif exists(text) and not drop_text_cond: text_embed = self.embed_text(text, seq_len, mask = mask) context, context_mask = None, None if prompt is not None: #for b in range(bs): # if drop_text_prompt[b]: # prompt[b] = "" if video_drop_prompt is not None: for b in range(bs): if video_drop_prompt[b]: prompt[b] = "the sound of X X" context, context_mask = self.encode_text(prompt) for b in range(bs): if drop_text_prompt[b]: context[b] = 0 if video_drop_prompt is not None and video_drop_prompt[b]: context[b] = 0 #print("cross attention", context.shape, context_mask.shape, x.shape, mask.shape, text_embed.shape if text_embed is not None else None, torch.mean(torch.abs(text_embed), dim=(1,2))) #print("video_drop_prompt", prompt, video_drop_prompt, context.shape, torch.mean(torch.abs(context), dim=(1,2))) #print("audio_drop_prompt", audio_drop_prompt, cond.shape, torch.mean(torch.abs(cond), dim=(1,2))) if self.proj_text is not None: text_embed = self.proj_text(text_embed) # attend attended, loss_contra = self.transformer( x, times = times, mask = mask, text_embed = text_embed, context = context, context_mask = context_mask ) pred = self.to_pred(attended) if not return_drop_conditions: return pred, loss_contra return pred, loss_contra, drop_audio_cond, drop_text_cond, drop_text_prompt def cfg_transformer_with_pred_head( self, *args, cfg_strength: float = 1., remove_parallel_component: bool = True, keep_parallel_frac: float = 0., **kwargs, ): pred, _ = self.transformer_with_pred_head(*args, drop_audio_cond = False, drop_text_cond = False, drop_text_prompt = False, **kwargs) if cfg_strength < 1e-5: return pred null_pred, _ = self.transformer_with_pred_head(*args, drop_audio_cond = True, drop_text_cond = True, drop_text_prompt = True, **kwargs) cfg_update = pred - null_pred if remove_parallel_component: # https://arxiv.org/abs/2410.02416 parallel, orthogonal = project(cfg_update, pred) cfg_update = orthogonal + parallel * keep_parallel_frac return pred + cfg_update * cfg_strength def add_noise(self, signal, mask, val): if self.audiocond_snr is None: return signal if not val: snr = np.random.uniform(self.audiocond_snr[0], self.audiocond_snr[1]) else: snr = (self.audiocond_snr[0] + self.audiocond_snr[1]) / 2.0 #print("add_noise", self.audiocond_snr, snr, signal.shape, mask) # [True, ..., False] noise = torch.randn_like(signal) w = torch.abs(signal[mask]).mean() / (torch.abs(noise[mask]).mean() + 1e-6) / snr return signal + noise * w @torch.no_grad() def sample( self, cond: Float['b n d'] | Float['b nw'] | None = None, *, text: Int['b nt'] | list[str] | None = None, lens: Int['b'] | None = None, duration: int | Int['b'] | None = None, steps = 32, cfg_strength = 1., # they used a classifier free guidance strength of 1. remove_parallel_component = True, sway_sampling = True, max_duration = 4096, # in case the duration predictor goes haywire vocoder: Callable[[Float['b d n']], list[Float['_']]] | None = None, return_raw_output: bool | None = None, save_to_filename: str | None = None, prompt = None, video_drop_prompt = None, audio_drop_prompt = None, video_paths = None ) -> ( Float['b n d'], list[Float['_']] ): self.eval() # raw wave if cond.ndim == 2: cond = self.mel_spec(cond) cond = rearrange(cond, 'b d n -> b n d') assert cond.shape[-1] == self.num_channels batch, cond_seq_len, device = *cond.shape[:2], cond.device if not exists(lens): lens = torch.full((batch,), cond_seq_len, device = device, dtype = torch.long) if video_paths is not None: text = self.encode_video(video_paths, cond_seq_len) # text elif isinstance(text, list): text = self.tokenizer(text).to(device) assert text.shape[0] == batch if exists(text): text_lens = (text != -1).sum(dim = -1) lens = torch.maximum(text_lens, lens) # make sure lengths are at least those of the text characters # duration cond_mask = lens_to_mask(lens) if exists(duration): if isinstance(duration, int): duration = torch.full((batch,), duration, device = device, dtype = torch.long) elif exists(self.duration_predictor): duration = self.duration_predictor(cond, text = text, lens = lens, return_loss = False).long() duration = torch.maximum(lens, duration) # just add one token so something is generated duration = duration.clamp(max = max_duration) assert duration.shape[0] == batch max_duration = duration.amax() cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value = 0.) cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value = False) cond_mask = rearrange(cond_mask, '... -> ... 1') mask = lens_to_mask(duration) #print("####", duration, mask, mask.shape, lens, cond_mask, cond_mask.shape, text) # neural ode def fn(t, x): # at each step, conditioning is fixed if lens[0] == duration[0]: print("No cond", lens, duration) step_cond = None else: step_cond = torch.where(cond_mask, self.add_noise(cond, cond_mask, True), torch.zeros_like(cond)) # predict flow return self.cfg_transformer_with_pred_head( x, step_cond, times = t, text = text, mask = mask, prompt = prompt, video_drop_prompt = video_drop_prompt, audio_drop_prompt = audio_drop_prompt, cfg_strength = cfg_strength, remove_parallel_component = remove_parallel_component ) ####torch.manual_seed(0) y0 = torch.randn_like(cond) t = torch.linspace(0, 1, steps, device = self.device) if sway_sampling: t = t + -1.0 * (torch.cos(torch.pi / 2 * t) - 1 + t) #print("@@@@", t) trajectory = odeint(fn, y0, t, **self.odeint_kwargs) sampled = trajectory[-1] out = sampled if lens[0] != duration[0]: out = torch.where(cond_mask, cond, out) # able to return raw untransformed output, if not using mel rep if exists(return_raw_output) and return_raw_output: return out # take care of transforming mel to audio if `vocoder` is passed in, or if `use_vocos` is turned on if exists(vocoder): assert not exists(self.vocos), '`use_vocos` should not be turned on if you are passing in a custom `vocoder` on sampling' out = rearrange(out, 'b n d -> b d n') out = vocoder(out) elif exists(self.vocos): audio = [] for mel, one_mask in zip(out, mask): #one_out = DB_to_amplitude(mel[one_mask], ref = 1., power = 0.5) one_out = mel[one_mask] one_out = rearrange(one_out, 'n d -> 1 d n') one_audio = self.vocos.decode(one_out) one_audio = rearrange(one_audio, '1 nw -> nw') audio.append(one_audio) out = audio if exists(save_to_filename): assert exists(vocoder) or exists(self.vocos) assert exists(self.sampling_rate) path = Path(save_to_filename) parent_path = path.parents[0] parent_path.mkdir(exist_ok = True, parents = True) for ind, one_audio in enumerate(out): one_audio = rearrange(one_audio, 'nw -> 1 nw') if len(out) == 1: save_path = str(parent_path / f'{path.name}') else: save_path = str(parent_path / f'{ind + 1}.{path.name}') torchaudio.save(save_path, one_audio.detach().cpu(), sample_rate = self.sampling_rate) return out def forward( self, inp: Float['b n d'] | Float['b nw'], # mel or raw wave *, text: Int['b nt'] | list[str] | None = None, times: int | Int['b'] | None = None, lens: Int['b'] | None = None, velocity_consistency_model: E2TTS | None = None, velocity_consistency_delta = 1e-3, prompt = None, video_drop_prompt=None, audio_drop_prompt=None, val = False, video_paths=None ): need_velocity_loss = exists(velocity_consistency_model) and self.velocity_consistency_weight > 0. # handle raw wave if inp.ndim == 2: inp = self.mel_spec(inp) inp = rearrange(inp, 'b d n -> b n d') assert inp.shape[-1] == self.num_channels batch, seq_len, dtype, device = *inp.shape[:2], inp.dtype, self.device if video_paths is not None: text = self.encode_video(video_paths, seq_len) # handle text as string elif isinstance(text, list): text = self.tokenizer(text).to(device) #print("text tokenized", text[0]) assert text.shape[0] == batch # lens and mask if not exists(lens): lens = torch.full((batch,), seq_len, device = device) mask = lens_to_mask(lens, length = seq_len) # get a random span to mask out for training conditionally if not val: if self.audiocond_drop_prob > 1.0: frac_lengths = torch.zeros((batch,), device = self.device).float().uniform_(1.0,1.0) else: frac_lengths = torch.zeros((batch,), device = self.device).float().uniform_(*self.frac_lengths_mask) else: frac_lengths = torch.tensor([(0.7+1.0)/2.0]*batch, device = self.device).float() rand_span_mask = mask_from_frac_lengths(lens, frac_lengths, max_length = seq_len, val = val) if exists(mask): rand_span_mask &= mask # mel is x1 x1 = inp # main conditional flow training logic # just ~5 loc # x0 is gaussian noise if val: torch.manual_seed(0) x0 = torch.randn_like(x1) if val: torch.manual_seed(int(time.time()*1000)) # t is random times from above if times is None: times = torch.rand((batch,), dtype = dtype, device = self.device) else: times = torch.tensor((times,)*batch, dtype = dtype, device = self.device) t = rearrange(times, 'b -> b 1 1') # if need velocity consistency, make sure time does not exceed 1. if need_velocity_loss: t = t * (1. - velocity_consistency_delta) # sample xt (w in the paper) w = (1. - t) * x0 + t * x1 flow = x1 - x0 # only predict what is within the random mask span for infilling if self.audiocond_drop_prob > 1.0: cond = None else: cond = einx.where( 'b n, b n d, b n d -> b n d', rand_span_mask, torch.zeros_like(x1), self.add_noise(x1, ~rand_span_mask, val) ) # transformer and prediction head if not val: pred, loss_contra, did_drop_audio_cond, did_drop_text_cond, did_drop_text_prompt = self.transformer_with_pred_head( w, cond, times = times, text = text, mask = mask, prompt = prompt, video_drop_prompt = video_drop_prompt, audio_drop_prompt = audio_drop_prompt, return_drop_conditions = True ) else: pred, loss_contra, did_drop_audio_cond, did_drop_text_cond, did_drop_text_prompt = self.transformer_with_pred_head( w, cond, times = times, text = text, mask = mask, prompt = prompt, video_drop_prompt = video_drop_prompt, audio_drop_prompt = audio_drop_prompt, drop_audio_cond = False, drop_text_cond = False, drop_text_prompt = False, return_drop_conditions = True ) # maybe velocity consistency loss velocity_loss = self.zero if need_velocity_loss: t_with_delta = t + velocity_consistency_delta w_with_delta = (1. - t_with_delta) * x0 + t_with_delta * x1 with torch.no_grad(): ema_pred, _ = velocity_consistency_model.transformer_with_pred_head( w_with_delta, cond, times = times + velocity_consistency_delta, text = text, mask = mask, prompt = prompt, video_drop_prompt = video_drop_prompt, audio_drop_prompt = audio_drop_prompt, drop_audio_cond = did_drop_audio_cond, drop_text_cond = did_drop_text_cond, drop_text_prompt = did_drop_text_prompt ) velocity_loss = F.mse_loss(pred, ema_pred, reduction = 'none') velocity_loss = velocity_loss[rand_span_mask].mean() # flow matching loss loss = F.mse_loss(pred, flow, reduction = 'none') #print("loss", loss.shape, loss, "rand_span_mask", rand_span_mask.shape, rand_span_mask, "loss[rand_span_mask]", loss[rand_span_mask].shape, loss[rand_span_mask]) loss = loss[rand_span_mask].mean() # total loss and get breakdown #total_loss = loss + velocity_loss * self.velocity_consistency_weight #breakdown = LossBreakdown(loss, velocity_loss) #print("loss", loss, velocity_loss, self.velocity_consistency_weight) ####contra_weight = 0.04 contra_weight = 1.0 total_loss = loss + loss_contra * contra_weight breakdown = LossBreakdown(loss, loss_contra * contra_weight) #print("loss", loss, loss_contra * contra_weight) # return total loss and bunch of intermediates return E2TTSReturn(total_loss, cond if cond is not None else w, pred, x0 + pred, breakdown)