work on pip package
This commit is contained in:
mrfakename
@@ -0,0 +1,94 @@
|
||||
import torch
|
||||
|
||||
|
||||
class MomentumBuffer:
|
||||
def __init__(self, momentum: float = -0.75):
|
||||
self.momentum = momentum
|
||||
self.running_average = 0
|
||||
|
||||
def update(self, update_value: torch.Tensor):
|
||||
new_average = self.momentum * self.running_average
|
||||
self.running_average = update_value + new_average
|
||||
|
||||
|
||||
def project(
|
||||
v0: torch.Tensor, # [B, C, H, W]
|
||||
v1: torch.Tensor, # [B, C, H, W]
|
||||
dims=[-1, -2],
|
||||
):
|
||||
dtype = v0.dtype
|
||||
if v0.device.type == "mps":
|
||||
v0, v1 = v0.float(), v1.float()
|
||||
else:
|
||||
v0, v1 = v0.double(), v1.double()
|
||||
v1 = torch.nn.functional.normalize(v1, dim=dims)
|
||||
v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
|
||||
v0_orthogonal = v0 - v0_parallel
|
||||
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
||||
|
||||
|
||||
def apg_forward(
|
||||
pred_cond: torch.Tensor, # [B, C, H, W]
|
||||
pred_uncond: torch.Tensor, # [B, C, H, W]
|
||||
guidance_scale: float,
|
||||
momentum_buffer: MomentumBuffer = None,
|
||||
eta: float = 0.0,
|
||||
norm_threshold: float = 2.5,
|
||||
dims=[-1, -2],
|
||||
):
|
||||
diff = pred_cond - pred_uncond
|
||||
if momentum_buffer is not None:
|
||||
momentum_buffer.update(diff)
|
||||
diff = momentum_buffer.running_average
|
||||
|
||||
if norm_threshold > 0:
|
||||
ones = torch.ones_like(diff)
|
||||
diff_norm = diff.norm(p=2, dim=dims, keepdim=True)
|
||||
scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
|
||||
diff = diff * scale_factor
|
||||
|
||||
diff_parallel, diff_orthogonal = project(diff, pred_cond, dims)
|
||||
normalized_update = diff_orthogonal + eta * diff_parallel
|
||||
pred_guided = pred_cond + (guidance_scale - 1) * normalized_update
|
||||
return pred_guided
|
||||
|
||||
|
||||
def cfg_forward(cond_output, uncond_output, cfg_strength):
|
||||
return uncond_output + cfg_strength * (cond_output - uncond_output)
|
||||
|
||||
|
||||
def cfg_double_condition_forward(
|
||||
cond_output,
|
||||
uncond_output,
|
||||
only_text_cond_output,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
):
|
||||
return (1 - guidance_scale_text) * uncond_output + (guidance_scale_text - guidance_scale_lyric) * only_text_cond_output + guidance_scale_lyric * cond_output
|
||||
|
||||
|
||||
def optimized_scale(positive_flat, negative_flat):
|
||||
|
||||
# Calculate dot production
|
||||
dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True)
|
||||
|
||||
# Squared norm of uncondition
|
||||
squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8
|
||||
|
||||
# st_star = v_cond^T * v_uncond / ||v_uncond||^2
|
||||
st_star = dot_product / squared_norm
|
||||
|
||||
return st_star
|
||||
|
||||
|
||||
def cfg_zero_star(noise_pred_with_cond, noise_pred_uncond, guidance_scale, i, zero_steps=1, use_zero_init=True):
|
||||
bsz = noise_pred_with_cond.shape[0]
|
||||
positive_flat = noise_pred_with_cond.view(bsz, -1)
|
||||
negative_flat = noise_pred_uncond.view(bsz, -1)
|
||||
alpha = optimized_scale(positive_flat, negative_flat)
|
||||
alpha = alpha.view(bsz, 1, 1, 1)
|
||||
if (i <= zero_steps) and use_zero_init:
|
||||
noise_pred = noise_pred_with_cond * 0.
|
||||
else:
|
||||
noise_pred = noise_pred_uncond * alpha + guidance_scale * (noise_pred_with_cond - noise_pred_uncond * alpha)
|
||||
return noise_pred
|
||||
@@ -0,0 +1,21 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
import random
|
||||
|
||||
|
||||
DEFAULT_ROOT_DIR = "examples/input_params"
|
||||
|
||||
|
||||
class DataSampler:
|
||||
def __init__(self, root_dir=DEFAULT_ROOT_DIR):
|
||||
self.root_dir = root_dir
|
||||
self.input_params_files = list(Path(self.root_dir).glob("*.json"))
|
||||
|
||||
def load_json(self, file_path):
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
def sample(self):
|
||||
json_path = random.choice(self.input_params_files)
|
||||
json_data = self.load_json(json_path)
|
||||
return json_data
|
||||
@@ -0,0 +1,43 @@
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint_path", type=str, default="")
|
||||
parser.add_argument("--server_name", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=7865)
|
||||
parser.add_argument("--device_id", type=int, default=0)
|
||||
parser.add_argument("--share", type=bool, default=False)
|
||||
parser.add_argument("--bf16", type=bool, default=True)
|
||||
parser.add_argument("--torch_compile", type=bool, default=False)
|
||||
args = parser.parse_args()
|
||||
|
||||
import os
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
|
||||
|
||||
|
||||
from ui.components import create_main_demo_ui
|
||||
from pipeline_ace_step import ACEStepPipeline
|
||||
from data_sampler import DataSampler
|
||||
|
||||
|
||||
def main(args):
|
||||
|
||||
model_demo = ACEStepPipeline(
|
||||
checkpoint_dir=args.checkpoint_path,
|
||||
dtype="bfloat16" if args.bf16 else "float32",
|
||||
torch_compile=args.torch_compile
|
||||
)
|
||||
data_sampler = DataSampler()
|
||||
|
||||
demo = create_main_demo_ui(
|
||||
text2music_process_func=model_demo.__call__,
|
||||
sample_data_func=data_sampler.sample,
|
||||
)
|
||||
demo.launch(
|
||||
server_name=args.server_name,
|
||||
server_port=args.port,
|
||||
share=args.share
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(args)
|
||||
@@ -0,0 +1,866 @@
|
||||
"""
|
||||
This file bundles language identification functions.
|
||||
|
||||
Modifications (fork): Copyright (c) 2021, Adrien Barbaresi.
|
||||
|
||||
Original code: Copyright (c) 2011 Marco Lui <saffsd@gmail.com>.
|
||||
Based on research by Marco Lui and Tim Baldwin.
|
||||
|
||||
See LICENSE file for more info.
|
||||
https://github.com/adbar/py3langid
|
||||
|
||||
Projects:
|
||||
https://github.com/juntaosun/LangSegment
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import numpy as np
|
||||
from collections import Counter
|
||||
from collections import defaultdict
|
||||
|
||||
# import langid
|
||||
# import py3langid as langid
|
||||
# pip install py3langid==0.2.2
|
||||
|
||||
# 启用语言预测概率归一化，概率预测的分数。因此，实现重新规范化 产生 0-1 范围内的输出。
|
||||
# langid disables probability normalization by default. For command-line usages of , it can be enabled by passing the flag.
|
||||
# For probability normalization in library use, the user must instantiate their own . An example of such usage is as follows:
|
||||
from py3langid.langid import LanguageIdentifier, MODEL_FILE
|
||||
|
||||
# Digital processing
|
||||
try:from .utils.num import num2str
|
||||
except ImportError:
|
||||
try:from utils.num import num2str
|
||||
except ImportError as e:
|
||||
raise e
|
||||
|
||||
# -----------------------------------
|
||||
# 更新日志：新版本分词更加精准。
|
||||
# Changelog: The new version of the word segmentation is more accurate.
|
||||
# チェンジログ:新しいバージョンの単語セグメンテーションはより正確です。
|
||||
# Changelog: 분할이라는 단어의 새로운 버전이 더 정확합니다.
|
||||
# -----------------------------------
|
||||
|
||||
|
||||
# Word segmentation function:
|
||||
# automatically identify and split the words (Chinese/English/Japanese/Korean) in the article or sentence according to different languages,
|
||||
# making it more suitable for TTS processing.
|
||||
# This code is designed for front-end text multi-lingual mixed annotation distinction, multi-language mixed training and inference of various TTS projects.
|
||||
# This processing result is mainly for (Chinese = zh, Japanese = ja, English = en, Korean = ko), and can actually support up to 97 different language mixing processing.
|
||||
|
||||
#===========================================================================================================
|
||||
#分かち書き機能:文章や文章の中の例えば（中国語/英語/日本語/韓国語）を、異なる言語で自動的に認識して分割し、TTS処理により適したものにします。
|
||||
#このコードは、さまざまなTTSプロジェクトのフロントエンドテキストの多言語混合注釈区別、多言語混合トレーニング、および推論のために特別に作成されています。
|
||||
#===========================================================================================================
|
||||
#(1)自動分詞:「韓国語では何を読むのですかあなたの体育の先生は誰ですか?今回の発表会では、iPhone 15シリーズの4機種が登場しました」
|
||||
#（2）手动分词:“あなたの名前は<ja>佐々木ですか?<ja>ですか?”
|
||||
#この処理結果は主に（中国語=ja、日本語=ja、英語=en、韓国語=ko）を対象としており、実際には最大97の異なる言語の混合処理をサポートできます。
|
||||
#===========================================================================================================
|
||||
|
||||
#===========================================================================================================
|
||||
# 단어 분할 기능: 기사 또는 문장에서 단어(중국어/영어/일본어/한국어)를 다른 언어에 따라 자동으로 식별하고 분할하여 TTS 처리에 더 적합합니다.
|
||||
# 이 코드는 프런트 엔드 텍스트 다국어 혼합 주석 분화, 다국어 혼합 교육 및 다양한 TTS 프로젝트의 추론을 위해 설계되었습니다.
|
||||
#===========================================================================================================
|
||||
# (1) 자동 단어 분할: "한국어로 무엇을 읽습니까? 스포츠 씨? 이 컨퍼런스는 4개의 iPhone 15 시리즈 모델을 제공합니다."
|
||||
# (2) 수동 참여: "이름이 <ja>Saki입니까? <ja>?"
|
||||
# 이 처리 결과는 주로 (중국어 = zh, 일본어 = ja, 영어 = en, 한국어 = ko)를 위한 것이며 실제로 혼합 처리를 위해 최대 97개의 언어를 지원합니다.
|
||||
#===========================================================================================================
|
||||
|
||||
# ===========================================================================================================
|
||||
# 分词功能：将文章或句子里的例如（中/英/日/韩），按不同语言自动识别并拆分，让它更适合TTS处理。
|
||||
# 本代码专为各种 TTS 项目的前端文本多语种混合标注区分，多语言混合训练和推理而编写。
|
||||
# ===========================================================================================================
|
||||
# （1）自动分词：“韩语中的오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型”
|
||||
# （2）手动分词：“你的名字叫<ja>佐々木？<ja>吗？”
|
||||
# 本处理结果主要针对（中文=zh , 日文=ja , 英文=en , 韩语=ko）, 实际上可支持多达 97 种不同的语言混合处理。
|
||||
# ===========================================================================================================
|
||||
|
||||
|
||||
# 手动分词标签规范：<语言标签>文本内容</语言标签>
|
||||
# 수동 단어 분할 태그 사양: <언어 태그> 텍스트 내용</언어 태그>
|
||||
# Manual word segmentation tag specification: <language tags> text content </language tags>
|
||||
# 手動分詞タグ仕様:<言語タグ>テキスト内容</言語タグ>
|
||||
# ===========================================================================================================
|
||||
# For manual word segmentation, labels need to appear in pairs, such as:
|
||||
# 如需手动分词，标签需要成对出现，例如：“<ja>佐々木<ja>” 或者 “<ja>佐々木</ja>”
|
||||
# 错误示范：“你的名字叫<ja>佐々木。” 此句子中出现的单个<ja>标签将被忽略，不会处理。
|
||||
# Error demonstration: "Your name is <ja>佐々木。" Single <ja> tags that appear in this sentence will be ignored and will not be processed.
|
||||
# ===========================================================================================================
|
||||
|
||||
|
||||
# ===========================================================================================================
|
||||
# 语音合成标记语言 SSML , 这里只支持它的标签（非 XML）Speech Synthesis Markup Language SSML, only its tags are supported here (not XML)
|
||||
# 想支持更多的 SSML 标签？欢迎 PR！ Want to support more SSML tags? PRs are welcome!
|
||||
# 说明：除了中文以外，它也可改造成支持多语种 SSML ，不仅仅是中文。
|
||||
# Note: In addition to Chinese, it can also be modified to support multi-language SSML, not just Chinese.
|
||||
# ===========================================================================================================
|
||||
# 中文实现：Chinese implementation:
|
||||
# 【SSML】<number>=中文大写数字读法（单字）
|
||||
# 【SSML】<telephone>=数字转成中文电话号码大写汉字（单字）
|
||||
# 【SSML】<currency>=按金额发音。
|
||||
# 【SSML】<date>=按日期发音。支持 2024年08月24, 2024/8/24, 2024-08, 08-24, 24 等输入。
|
||||
# ===========================================================================================================
|
||||
class LangSSML:
|
||||
|
||||
def __init__(self):
|
||||
# 纯数字
|
||||
self._zh_numerals_number = {
|
||||
'0': '零',
|
||||
'1': '一',
|
||||
'2': '二',
|
||||
'3': '三',
|
||||
'4': '四',
|
||||
'5': '五',
|
||||
'6': '六',
|
||||
'7': '七',
|
||||
'8': '八',
|
||||
'9': '九'
|
||||
}
|
||||
|
||||
# 将2024/8/24, 2024-08, 08-24, 24 标准化“年月日”
|
||||
# Standardize 2024/8/24, 2024-08, 08-24, 24 to "year-month-day"
|
||||
def _format_chinese_data(self, date_str:str):
|
||||
# 处理日期格式
|
||||
input_date = date_str
|
||||
if date_str is None or date_str.strip() == "":return ""
|
||||
date_str = re.sub(r"[\/\._|年|月]","-",date_str)
|
||||
date_str = re.sub(r"日",r"",date_str)
|
||||
date_arrs = date_str.split(' ')
|
||||
if len(date_arrs) == 1 and ":" in date_arrs[0]:
|
||||
time_str = date_arrs[0]
|
||||
date_arrs = []
|
||||
else:
|
||||
time_str = date_arrs[1] if len(date_arrs) >=2 else ""
|
||||
def nonZero(num,cn,func=None):
|
||||
if func is not None:num=func(num)
|
||||
return f"{num}{cn}" if num is not None and num != "" and num != "0" else ""
|
||||
f_number = self.to_chinese_number
|
||||
f_currency = self.to_chinese_currency
|
||||
# year, month, day
|
||||
year_month_day = ""
|
||||
if len(date_arrs) > 0:
|
||||
year, month, day = "","",""
|
||||
parts = date_arrs[0].split('-')
|
||||
if len(parts) == 3: # 格式为 YYYY-MM-DD
|
||||
year, month, day = parts
|
||||
elif len(parts) == 2: # 格式为 MM-DD 或 YYYY-MM
|
||||
if len(parts[0]) == 4: # 年-月
|
||||
year, month = parts
|
||||
else:month, day = parts # 月-日
|
||||
elif len(parts[0]) > 0: # 仅有月-日或年
|
||||
if len(parts[0]) == 4:
|
||||
year = parts[0]
|
||||
else:day = parts[0]
|
||||
year,month,day = nonZero(year,"年",f_number),nonZero(month,"月",f_currency),nonZero(day,"日",f_currency)
|
||||
year_month_day = re.sub(r"([年|月|日])+",r"\1",f"{year}{month}{day}")
|
||||
# hours, minutes, seconds
|
||||
time_str = re.sub(r"[\/\.\-：_]",":",time_str)
|
||||
time_arrs = time_str.split(":")
|
||||
hours, minutes, seconds = "","",""
|
||||
if len(time_arrs) == 3: # H/M/S
|
||||
hours, minutes, seconds = time_arrs
|
||||
elif len(time_arrs) == 2:# H/M
|
||||
hours, minutes = time_arrs
|
||||
elif len(time_arrs[0]) > 0:hours = f'{time_arrs[0]}点' # H
|
||||
if len(time_arrs) > 1:
|
||||
hours, minutes, seconds = nonZero(hours,"点",f_currency),nonZero(minutes,"分",f_currency),nonZero(seconds,"秒",f_currency)
|
||||
hours_minutes_seconds = re.sub(r"([点|分|秒])+",r"\1",f"{hours}{minutes}{seconds}")
|
||||
output_date = f"{year_month_day}{hours_minutes_seconds}"
|
||||
return output_date
|
||||
|
||||
# 【SSML】number=中文大写数字读法（单字）
|
||||
# Chinese Numbers(single word)
|
||||
def to_chinese_number(self, num:str):
|
||||
pattern = r'(\d+)'
|
||||
zh_numerals = self._zh_numerals_number
|
||||
arrs = re.split(pattern, num)
|
||||
output = ""
|
||||
for item in arrs:
|
||||
if re.match(pattern,item):
|
||||
output += ''.join(zh_numerals[digit] if digit in zh_numerals else "" for digit in str(item))
|
||||
else:output += item
|
||||
output = output.replace(".","点")
|
||||
return output
|
||||
|
||||
# 【SSML】telephone=数字转成中文电话号码大写汉字（单字）
|
||||
# Convert numbers to Chinese phone numbers in uppercase Chinese characters(single word)
|
||||
def to_chinese_telephone(self, num:str):
|
||||
output = self.to_chinese_number(num.replace("+86","")) # zh +86
|
||||
output = output.replace("一","幺")
|
||||
return output
|
||||
|
||||
# 【SSML】currency=按金额发音。
|
||||
# Digital processing from GPT_SoVITS num.py （thanks）
|
||||
def to_chinese_currency(self, num:str):
|
||||
pattern = r'(\d+)'
|
||||
arrs = re.split(pattern, num)
|
||||
output = ""
|
||||
for item in arrs:
|
||||
if re.match(pattern,item):
|
||||
output += num2str(item)
|
||||
else:output += item
|
||||
output = output.replace(".","点")
|
||||
return output
|
||||
|
||||
# 【SSML】date=按日期发音。支持 2024年08月24, 2024/8/24, 2024-08, 08-24, 24 等输入。
|
||||
def to_chinese_date(self, num:str):
|
||||
chinese_date = self._format_chinese_data(num)
|
||||
return chinese_date
|
||||
|
||||
|
||||
class LangSegment:
|
||||
|
||||
def __init__(self):
|
||||
|
||||
self.langid = LanguageIdentifier.from_pickled_model(MODEL_FILE, norm_probs=True)
|
||||
|
||||
self._text_cache = None
|
||||
self._text_lasts = None
|
||||
self._text_langs = None
|
||||
self._lang_count = None
|
||||
self._lang_eos = None
|
||||
|
||||
# 可自定义语言匹配标签：カスタマイズ可能な言語対応タグ:사용자 지정 가능한 언어 일치 태그:
|
||||
# Customizable language matching tags: These are supported，이 표현들은 모두 지지합니다
|
||||
# <zh>你好<zh> , <ja>佐々木</ja> , <en>OK<en> , <ko>오빠</ko> 这些写法均支持
|
||||
self.SYMBOLS_PATTERN = r'(<([a-zA-Z|-]*)>(.*?)<\/*[a-zA-Z|-]*>)'
|
||||
|
||||
# 语言过滤组功能, 可以指定保留语言。不在过滤组中的语言将被清除。您可随心搭配TTS语音合成所支持的语言。
|
||||
# 언어 필터 그룹 기능을 사용하면 예약된 언어를 지정할 수 있습니다. 필터 그룹에 없는 언어는 지워집니다. TTS 텍스트에서 지원하는 언어를 원하는 대로 일치시킬 수 있습니다.
|
||||
# 言語フィルターグループ機能では、予約言語を指定できます。フィルターグループに含まれていない言語はクリアされます。TTS音声合成がサポートする言語を自由に組み合わせることができます。
|
||||
# The language filter group function allows you to specify reserved languages.
|
||||
# Languages not in the filter group will be cleared. You can match the languages supported by TTS Text To Speech as you like.
|
||||
# 排名越前，优先级越高，The higher the ranking, the higher the priority，ランキングが上位になるほど、優先度が高くなります。
|
||||
|
||||
# 系统默认过滤器。System default filter。(ISO 639-1 codes given)
|
||||
# ----------------------------------------------------------------------------------------------------------------------------------
|
||||
# "zh"中文=Chinese ,"en"英语=English ,"ja"日语=Japanese ,"ko"韩语=Korean ,"fr"法语=French ,"vi"越南语=Vietnamese , "ru"俄语=Russian
|
||||
# "th"泰语=Thai
|
||||
# ----------------------------------------------------------------------------------------------------------------------------------
|
||||
self.DEFAULT_FILTERS = ["zh", "ja", "ko", "en"]
|
||||
|
||||
# 用户可自定义过滤器。User-defined filters
|
||||
self.Langfilters = self.DEFAULT_FILTERS[:] # 创建副本
|
||||
|
||||
# 合并文本
|
||||
self.isLangMerge = True
|
||||
|
||||
# 试验性支持：您可自定义添加："fr"法语 , "vi"越南语。Experimental: You can customize to add: "fr" French, "vi" Vietnamese.
|
||||
# 请使用API启用：self.setfilters(["zh", "en", "ja", "ko", "fr", "vi" , "ru" , "th"]) # 您可自定义添加，如："fr"法语 , "vi"越南语。
|
||||
|
||||
# 预览版功能，自动启用或禁用，无需设置
|
||||
# Preview feature, automatically enabled or disabled, no settings required
|
||||
self.EnablePreview = False
|
||||
|
||||
# 除此以外，它支持简写过滤器，只需按不同语种任意组合即可。
|
||||
# In addition to that, it supports abbreviation filters, allowing for any combination of different languages.
|
||||
# 示例：您可以任意指定多种组合，进行过滤
|
||||
# Example: You can specify any combination to filter
|
||||
|
||||
# 中/日语言优先级阀值（评分范围为 0 ~ 1）:评分低于设定阀值 <0.89 时，启用 filters 中的优先级。\n
|
||||
# 중/일본어 우선 순위 임계값(점수 범위 0-1): 점수가 설정된 임계값 <0.89보다 낮을 때 필터에서 우선 순위를 활성화합니다.
|
||||
# 中国語/日本語の優先度しきい値（スコア範囲0〜1）:スコアが設定されたしきい値<0.89未満の場合、フィルターの優先度が有効になります。\n
|
||||
# Chinese and Japanese language priority threshold (score range is 0 ~ 1): The default threshold is 0.89. \n
|
||||
# Only the common characters between Chinese and Japanese are processed with confidence and priority. \n
|
||||
self.LangPriorityThreshold = 0.89
|
||||
|
||||
# Langfilters = ["zh"] # 按中文识别
|
||||
# Langfilters = ["en"] # 按英文识别
|
||||
# Langfilters = ["ja"] # 按日文识别
|
||||
# Langfilters = ["ko"] # 按韩文识别
|
||||
# Langfilters = ["zh_ja"] # 中日混合识别
|
||||
# Langfilters = ["zh_en"] # 中英混合识别
|
||||
# Langfilters = ["ja_en"] # 日英混合识别
|
||||
# Langfilters = ["zh_ko"] # 中韩混合识别
|
||||
# Langfilters = ["ja_ko"] # 日韩混合识别
|
||||
# Langfilters = ["en_ko"] # 英韩混合识别
|
||||
# Langfilters = ["zh_ja_en"] # 中日英混合识别
|
||||
# Langfilters = ["zh_ja_en_ko"] # 中日英韩混合识别
|
||||
|
||||
# 更多过滤组合，请您随意。。。For more filter combinations, please feel free to......
|
||||
# より多くのフィルターの組み合わせ、お気軽に。。。더 많은 필터 조합을 원하시면 자유롭게 해주세요. .....
|
||||
|
||||
# 可选保留：支持中文数字拼音格式，更方便前端实现拼音音素修改和推理，默认关闭 False 。
|
||||
# 开启后 True ，括号内的数字拼音格式均保留，并识别输出为："zh"中文。
|
||||
self.keepPinyin = False
|
||||
|
||||
# DEFINITION
|
||||
self.PARSE_TAG = re.compile(r'(⑥\$*\d+[\d]{6,}⑥)')
|
||||
|
||||
self.LangSSML = LangSSML()
|
||||
|
||||
def _clears(self):
|
||||
self._text_cache = None
|
||||
self._text_lasts = None
|
||||
self._text_langs = None
|
||||
self._text_waits = None
|
||||
self._lang_count = None
|
||||
self._lang_eos = None
|
||||
|
||||
def _is_english_word(self, word):
|
||||
return bool(re.match(r'^[a-zA-Z]+$', word))
|
||||
|
||||
def _is_chinese(self, word):
|
||||
for char in word:
|
||||
if '\u4e00' <= char <= '\u9fff':
|
||||
return True
|
||||
return False
|
||||
|
||||
def _is_japanese_kana(self, word):
|
||||
pattern = re.compile(r'[\u3040-\u309F\u30A0-\u30FF]+')
|
||||
matches = pattern.findall(word)
|
||||
return len(matches) > 0
|
||||
|
||||
def _insert_english_uppercase(self, word):
|
||||
modified_text = re.sub(r'(?<!\b)([A-Z])', r' \1', word)
|
||||
modified_text = modified_text.strip('-')
|
||||
return modified_text + " "
|
||||
|
||||
def _split_camel_case(self, word):
|
||||
return re.sub(r'(?<!^)(?=[A-Z])', ' ', word)
|
||||
|
||||
def _statistics(self, language, text):
|
||||
# Language word statistics:
|
||||
# Chinese characters usually occupy double bytes
|
||||
if self._lang_count is None or not isinstance(self._lang_count, defaultdict):
|
||||
self._lang_count = defaultdict(int)
|
||||
lang_count = self._lang_count
|
||||
if not "|" in language:
|
||||
lang_count[language] += int(len(text)*2) if language == "zh" else len(text)
|
||||
self._lang_count = lang_count
|
||||
|
||||
def _clear_text_number(self, text):
|
||||
if text == "\n":return text,False # Keep Line Breaks
|
||||
clear_text = re.sub(r'([^\w\s]+)','',re.sub(r'\n+','',text)).strip()
|
||||
is_number = len(re.sub(re.compile(r'(\d+)'),'',clear_text)) == 0
|
||||
return clear_text,is_number
|
||||
|
||||
def _saveData(self, words,language:str,text:str,score:float,symbol=None):
|
||||
# Pre-detection
|
||||
clear_text , is_number = self._clear_text_number(text)
|
||||
# Merge the same language and save the results
|
||||
preData = words[-1] if len(words) > 0 else None
|
||||
if symbol is not None:pass
|
||||
elif preData is not None and preData["symbol"] is None:
|
||||
if len(clear_text) == 0:language = preData["lang"]
|
||||
elif is_number == True:language = preData["lang"]
|
||||
_ , pre_is_number = self._clear_text_number(preData["text"])
|
||||
if (preData["lang"] == language):
|
||||
self._statistics(preData["lang"],text)
|
||||
text = preData["text"] + text
|
||||
preData["text"] = text
|
||||
return preData
|
||||
elif pre_is_number == True:
|
||||
text = f'{preData["text"]}{text}'
|
||||
words.pop()
|
||||
elif is_number == True:
|
||||
priority_language = self._get_filters_string()[:2]
|
||||
if priority_language in "ja-zh-en-ko-fr-vi":language = priority_language
|
||||
data = {"lang":language,"text": text,"score":score,"symbol":symbol}
|
||||
filters = self.Langfilters
|
||||
if filters is None or len(filters) == 0 or "?" in language or \
|
||||
language in filters or language in filters[0] or \
|
||||
filters[0] == "*" or filters[0] in "alls-mixs-autos":
|
||||
words.append(data)
|
||||
self._statistics(data["lang"],data["text"])
|
||||
return data
|
||||
|
||||
def _addwords(self, words,language,text,score,symbol=None):
|
||||
if text == "\n":pass # Keep Line Breaks
|
||||
elif text is None or len(text.strip()) == 0:return True
|
||||
if language is None:language = ""
|
||||
language = language.lower()
|
||||
if language == 'en':text = self._insert_english_uppercase(text)
|
||||
# text = re.sub(r'[(（）)]', ',' , text) # Keep it.
|
||||
text_waits = self._text_waits
|
||||
ispre_waits = len(text_waits)>0
|
||||
preResult = text_waits.pop() if ispre_waits else None
|
||||
if preResult is None:preResult = words[-1] if len(words) > 0 else None
|
||||
if preResult and ("|" in preResult["lang"]):
|
||||
pre_lang = preResult["lang"]
|
||||
if language in pre_lang:preResult["lang"] = language = language.split("|")[0]
|
||||
else:preResult["lang"]=pre_lang.split("|")[0]
|
||||
if ispre_waits:preResult = self._saveData(words,preResult["lang"],preResult["text"],preResult["score"],preResult["symbol"])
|
||||
pre_lang = preResult["lang"] if preResult else None
|
||||
if ("|" in language) and (pre_lang and not pre_lang in language and not "…" in language):language = language.split("|")[0]
|
||||
if "|" in language:self._text_waits.append({"lang":language,"text": text,"score":score,"symbol":symbol})
|
||||
else:self._saveData(words,language,text,score,symbol)
|
||||
return False
|
||||
|
||||
def _get_prev_data(self, words):
|
||||
data = words[-1] if words and len(words) > 0 else None
|
||||
if data:return (data["lang"] , data["text"])
|
||||
return (None,"")
|
||||
|
||||
def _match_ending(self, input , index):
|
||||
if input is None or len(input) == 0:return False,None
|
||||
input = re.sub(r'\s+', '', input)
|
||||
if len(input) == 0 or abs(index) > len(input):return False,None
|
||||
ending_pattern = re.compile(r'([「」“”‘’"\':：。.！!?．？])')
|
||||
return ending_pattern.match(input[index]),input[index]
|
||||
|
||||
def _cleans_text(self, cleans_text):
|
||||
cleans_text = re.sub(r'(.*?)([^\w]+)', r'\1 ', cleans_text)
|
||||
cleans_text = re.sub(r'(.)\1+', r'\1', cleans_text)
|
||||
return cleans_text.strip()
|
||||
|
||||
def _mean_processing(self, text:str):
|
||||
if text is None or (text.strip()) == "":return None , 0.0
|
||||
arrs = self._split_camel_case(text).split(" ")
|
||||
langs = []
|
||||
for t in arrs:
|
||||
if len(t.strip()) <= 3:continue
|
||||
language, score = self.langid.classify(t)
|
||||
langs.append({"lang":language})
|
||||
if len(langs) == 0:return None , 0.0
|
||||
return Counter([item['lang'] for item in langs]).most_common(1)[0][0],1.0
|
||||
|
||||
def _lang_classify(self, cleans_text):
|
||||
language, score = self.langid.classify(cleans_text)
|
||||
# fix: Huggingface is np.float32
|
||||
if score is not None and isinstance(score, np.generic) and hasattr(score,"item"):
|
||||
score = score.item()
|
||||
score = round(score , 3)
|
||||
return language, score
|
||||
|
||||
def _get_filters_string(self):
|
||||
filters = self.Langfilters
|
||||
return "-".join(filters).lower().strip() if filters is not None else ""
|
||||
|
||||
def _parse_language(self, words , segment):
|
||||
LANG_JA = "ja"
|
||||
LANG_ZH = "zh"
|
||||
LANG_ZH_JA = f'{LANG_ZH}|{LANG_JA}'
|
||||
LANG_JA_ZH = f'{LANG_JA}|{LANG_ZH}'
|
||||
language = LANG_ZH
|
||||
regex_pattern = re.compile(r'([^\w\s]+)')
|
||||
lines = regex_pattern.split(segment)
|
||||
lines_max = len(lines)
|
||||
LANG_EOS =self._lang_eos
|
||||
for index, text in enumerate(lines):
|
||||
if len(text) == 0:continue
|
||||
EOS = index >= (lines_max - 1)
|
||||
nextId = index + 1
|
||||
nextText = lines[nextId] if not EOS else ""
|
||||
nextPunc = len(re.sub(regex_pattern,'',re.sub(r'\n+','',nextText)).strip()) == 0
|
||||
textPunc = len(re.sub(regex_pattern,'',re.sub(r'\n+','',text)).strip()) == 0
|
||||
if not EOS and (textPunc == True or ( len(nextText.strip()) >= 0 and nextPunc == True)):
|
||||
lines[nextId] = f'{text}{nextText}'
|
||||
continue
|
||||
number_tags = re.compile(r'(⑥\d{6,}⑥)')
|
||||
cleans_text = re.sub(number_tags, '' ,text)
|
||||
cleans_text = re.sub(r'\d+', '' ,cleans_text)
|
||||
cleans_text = self._cleans_text(cleans_text)
|
||||
# fix:Langid's recognition of short sentences is inaccurate, and it is spliced longer.
|
||||
if not EOS and len(cleans_text) <= 2:
|
||||
lines[nextId] = f'{text}{nextText}'
|
||||
continue
|
||||
language,score = self._lang_classify(cleans_text)
|
||||
prev_language , prev_text = self._get_prev_data(words)
|
||||
if language != LANG_ZH and all('\u4e00' <= c <= '\u9fff' for c in re.sub(r'\s','',cleans_text)):language,score = LANG_ZH,1
|
||||
if len(cleans_text) <= 5 and self._is_chinese(cleans_text):
|
||||
filters_string = self._get_filters_string()
|
||||
if score < self.LangPriorityThreshold and len(filters_string) > 0:
|
||||
index_ja , index_zh = filters_string.find(LANG_JA) , filters_string.find(LANG_ZH)
|
||||
if index_ja != -1 and index_ja < index_zh:language = LANG_JA
|
||||
elif index_zh != -1 and index_zh < index_ja:language = LANG_ZH
|
||||
if self._is_japanese_kana(cleans_text):language = LANG_JA
|
||||
elif len(cleans_text) > 2 and score > 0.90:pass
|
||||
elif EOS and LANG_EOS:language = LANG_ZH if len(cleans_text) <= 1 else language
|
||||
else:
|
||||
LANG_UNKNOWN = LANG_ZH_JA if language == LANG_ZH or (len(cleans_text) <=2 and prev_language == LANG_ZH) else LANG_JA_ZH
|
||||
match_end,match_char = self._match_ending(text, -1)
|
||||
referen = prev_language in LANG_UNKNOWN or LANG_UNKNOWN in prev_language if prev_language else False
|
||||
if match_char in "。.": language = prev_language if referen and len(words) > 0 else language
|
||||
else:language = f"{LANG_UNKNOWN}|…"
|
||||
text,*_ = re.subn(number_tags , self._restore_number , text )
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
# ----------------------------------------------------------
|
||||
# 【SSML】中文数字处理：Chinese Number Processing (SSML support)
|
||||
# 这里默认都是中文，用于处理 SSML 中文标签。当然可以支持任意语言，例如：
|
||||
# The default here is Chinese, which is used to process SSML Chinese tags. Of course, any language can be supported, for example:
|
||||
# 中文电话号码：<telephone>1234567</telephone>
|
||||
# 中文数字号码：<number>1234567</number>
|
||||
def _process_symbol_SSML(self, words,data):
|
||||
tag , match = data
|
||||
language = SSML = match[1]
|
||||
text = match[2]
|
||||
score = 1.0
|
||||
if SSML == "telephone":
|
||||
# 中文-电话号码
|
||||
language = "zh"
|
||||
text = self.LangSSML.to_chinese_telephone(text)
|
||||
elif SSML == "number":
|
||||
# 中文-数字读法
|
||||
language = "zh"
|
||||
text = self.LangSSML.to_chinese_number(text)
|
||||
elif SSML == "currency":
|
||||
# 中文-按金额发音
|
||||
language = "zh"
|
||||
text = self.LangSSML.to_chinese_currency(text)
|
||||
elif SSML == "date":
|
||||
# 中文-按金额发音
|
||||
language = "zh"
|
||||
text = self.LangSSML.to_chinese_date(text)
|
||||
self._addwords(words,language,text,score,SSML)
|
||||
|
||||
# ----------------------------------------------------------
|
||||
def _restore_number(self, matche):
|
||||
value = matche.group(0)
|
||||
text_cache = self._text_cache
|
||||
if value in text_cache:
|
||||
process , data = text_cache[value]
|
||||
tag , match = data
|
||||
value = match
|
||||
return value
|
||||
|
||||
def _pattern_symbols(self, item , text):
|
||||
if text is None:return text
|
||||
tag , pattern , process = item
|
||||
matches = pattern.findall(text)
|
||||
if len(matches) == 1 and "".join(matches[0]) == text:
|
||||
return text
|
||||
for i , match in enumerate(matches):
|
||||
key = f"⑥{tag}{i:06d}⑥"
|
||||
text = re.sub(pattern , key , text , count=1)
|
||||
self._text_cache[key] = (process , (tag , match))
|
||||
return text
|
||||
|
||||
def _process_symbol(self, words,data):
|
||||
tag , match = data
|
||||
language = match[1]
|
||||
text = match[2]
|
||||
score = 1.0
|
||||
filters = self._get_filters_string()
|
||||
if language not in filters:
|
||||
self._process_symbol_SSML(words,data)
|
||||
else:
|
||||
self._addwords(words,language,text,score,True)
|
||||
|
||||
def _process_english(self, words,data):
|
||||
tag , match = data
|
||||
text = match[0]
|
||||
filters = self._get_filters_string()
|
||||
priority_language = filters[:2]
|
||||
# Preview feature, other language segmentation processing
|
||||
enablePreview = self.EnablePreview
|
||||
if enablePreview == True:
|
||||
# Experimental: Other language support
|
||||
regex_pattern = re.compile(r'(.*?[。.?？!！]+[\n]{,1})')
|
||||
lines = regex_pattern.split(text)
|
||||
for index , text in enumerate(lines):
|
||||
if len(text.strip()) == 0:continue
|
||||
cleans_text = self._cleans_text(text)
|
||||
language,score = self._lang_classify(cleans_text)
|
||||
if language not in filters:
|
||||
language,score = self._mean_processing(cleans_text)
|
||||
if language is None or score <= 0.0:continue
|
||||
elif language in filters:pass # pass
|
||||
elif score >= 0.95:continue # High score, but not in the filter, excluded.
|
||||
elif score <= 0.15 and filters[:2] == "fr":language = priority_language
|
||||
else:language = "en"
|
||||
self._addwords(words,language,text,score)
|
||||
else:
|
||||
# Default is English
|
||||
language, score = "en", 1.0
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_Russian(self, words,data):
|
||||
tag , match = data
|
||||
text = match[0]
|
||||
language = "ru"
|
||||
score = 1.0
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_Thai(self, words,data):
|
||||
tag , match = data
|
||||
text = match[0]
|
||||
language = "th"
|
||||
score = 1.0
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_korean(self, words,data):
|
||||
tag , match = data
|
||||
text = match[0]
|
||||
language = "ko"
|
||||
score = 1.0
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_quotes(self, words,data):
|
||||
tag , match = data
|
||||
text = "".join(match)
|
||||
childs = self.PARSE_TAG.findall(text)
|
||||
if len(childs) > 0:
|
||||
self._process_tags(words , text , False)
|
||||
else:
|
||||
cleans_text = self._cleans_text(match[1])
|
||||
if len(cleans_text) <= 5:
|
||||
self._parse_language(words,text)
|
||||
else:
|
||||
language,score = self._lang_classify(cleans_text)
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_pinyin(self, words,data):
|
||||
tag , match = data
|
||||
text = match
|
||||
language = "zh"
|
||||
score = 1.0
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_number(self, words,data): # "$0" process only
|
||||
"""
|
||||
Numbers alone cannot accurately identify language.
|
||||
Because numbers are universal in all languages.
|
||||
So it won't be executed here, just for testing.
|
||||
"""
|
||||
tag , match = data
|
||||
language = words[0]["lang"] if len(words) > 0 else "zh"
|
||||
text = match
|
||||
score = 0.0
|
||||
self._addwords(words,language,text,score)
|
||||
|
||||
def _process_tags(self, words , text , root_tag):
|
||||
text_cache = self._text_cache
|
||||
segments = re.split(self.PARSE_TAG, text)
|
||||
segments_len = len(segments) - 1
|
||||
for index , text in enumerate(segments):
|
||||
if root_tag:self._lang_eos = index >= segments_len
|
||||
if self.PARSE_TAG.match(text):
|
||||
process , data = text_cache[text]
|
||||
if process:process(words , data)
|
||||
else:
|
||||
self._parse_language(words , text)
|
||||
return words
|
||||
|
||||
def _merge_results(self, words):
|
||||
new_word = []
|
||||
for index , cur_data in enumerate(words):
|
||||
if "symbol" in cur_data:del cur_data["symbol"]
|
||||
if index == 0:new_word.append(cur_data)
|
||||
else:
|
||||
pre_data = new_word[-1]
|
||||
if cur_data["lang"] == pre_data["lang"]:
|
||||
pre_data["text"] = f'{pre_data["text"]}{cur_data["text"]}'
|
||||
else:new_word.append(cur_data)
|
||||
return new_word
|
||||
|
||||
def _parse_symbols(self, text):
|
||||
TAG_NUM = "00" # "00" => default channels , "$0" => testing channel
|
||||
TAG_S1,TAG_S2,TAG_P1,TAG_P2,TAG_EN,TAG_KO,TAG_RU,TAG_TH = "$1" ,"$2" ,"$3" ,"$4" ,"$5" ,"$6" ,"$7","$8"
|
||||
TAG_BASE = re.compile(fr'(([【《（(“‘"\']*[LANGUAGE]+[\W\s]*)+)')
|
||||
# Get custom language filter
|
||||
filters = self.Langfilters
|
||||
filters = filters if filters is not None else ""
|
||||
# =======================================================================================================
|
||||
# Experimental: Other language support.Thử nghiệm: Hỗ trợ ngôn ngữ khác.Expérimental : prise en charge d’autres langues.
|
||||
# 相关语言字符如有缺失，熟悉相关语言的朋友，可以提交把缺失的发音符号补全。
|
||||
# If relevant language characters are missing, friends who are familiar with the relevant languages can submit a submission to complete the missing pronunciation symbols.
|
||||
# S'il manque des caractères linguistiques pertinents, les amis qui connaissent les langues concernées peuvent soumettre une soumission pour compléter les symboles de prononciation manquants.
|
||||
# Nếu thiếu ký tự ngôn ngữ liên quan, những người bạn quen thuộc với ngôn ngữ liên quan có thể gửi bài để hoàn thành các ký hiệu phát âm còn thiếu.
|
||||
# -------------------------------------------------------------------------------------------------------
|
||||
# Preview feature, other language support
|
||||
enablePreview = self.EnablePreview
|
||||
if "fr" in filters or \
|
||||
"vi" in filters:enablePreview = True
|
||||
self.EnablePreview = enablePreview
|
||||
# 实验性：法语字符支持。Prise en charge des caractères français
|
||||
RE_FR = "" if not enablePreview else "àáâãäåæçèéêëìíîïðñòóôõöùúûüýþÿ"
|
||||
# 实验性：越南语字符支持。Hỗ trợ ký tự tiếng Việt
|
||||
RE_VI = "" if not enablePreview else "đơưăáàảãạắằẳẵặấầẩẫậéèẻẽẹếềểễệíìỉĩịóòỏõọốồổỗộớờởỡợúùủũụứừửữựôâêơưỷỹ"
|
||||
# -------------------------------------------------------------------------------------------------------
|
||||
# Basic options:
|
||||
process_list = [
|
||||
( TAG_S1 , re.compile(self.SYMBOLS_PATTERN) , self._process_symbol ), # Symbol Tag
|
||||
( TAG_KO , re.compile(re.sub(r'LANGUAGE',f'\uac00-\ud7a3',TAG_BASE.pattern)) , self._process_korean ), # Korean words
|
||||
( TAG_TH , re.compile(re.sub(r'LANGUAGE',f'\u0E00-\u0E7F',TAG_BASE.pattern)) , self._process_Thai ), # Thai words support.
|
||||
( TAG_RU , re.compile(re.sub(r'LANGUAGE',f'А-Яа-яЁё',TAG_BASE.pattern)) , self._process_Russian ), # Russian words support.
|
||||
( TAG_NUM , re.compile(r'(\W*\d+\W+\d*\W*\d*)') , self._process_number ), # Number words, Universal in all languages, Ignore it.
|
||||
( TAG_EN , re.compile(re.sub(r'LANGUAGE',f'a-zA-Z{RE_FR}{RE_VI}',TAG_BASE.pattern)) , self._process_english ), # English words + Other language support.
|
||||
( TAG_P1 , re.compile(r'(["\'])(.*?)(\1)') , self._process_quotes ), # Regular quotes
|
||||
( TAG_P2 , re.compile(r'([\n]*[【《（(“‘])([^【《（(“‘’”)）》】]{3,})([’”)）》】][\W\s]*[\n]{,1})') , self._process_quotes ), # Special quotes, There are left and right.
|
||||
]
|
||||
# Extended options: Default False
|
||||
if self.keepPinyin == True:process_list.insert(1 ,
|
||||
( TAG_S2 , re.compile(r'([\(（{](?:\s*\w*\d\w*\s*)+[}）\)])') , self._process_pinyin ), # Chinese Pinyin Tag.
|
||||
)
|
||||
# -------------------------------------------------------------------------------------------------------
|
||||
words = []
|
||||
lines = re.findall(r'.*\n*', re.sub(self.PARSE_TAG, '' ,text))
|
||||
for index , text in enumerate(lines):
|
||||
if len(text.strip()) == 0:continue
|
||||
self._lang_eos = False
|
||||
self._text_cache = {}
|
||||
for item in process_list:
|
||||
text = self._pattern_symbols(item , text)
|
||||
cur_word = self._process_tags([] , text , True)
|
||||
if len(cur_word) == 0:continue
|
||||
cur_data = cur_word[0] if len(cur_word) > 0 else None
|
||||
pre_data = words[-1] if len(words) > 0 else None
|
||||
if cur_data and pre_data and cur_data["lang"] == pre_data["lang"] \
|
||||
and cur_data["symbol"] == False and pre_data["symbol"] :
|
||||
cur_data["text"] = f'{pre_data["text"]}{cur_data["text"]}'
|
||||
words.pop()
|
||||
words += cur_word
|
||||
if self.isLangMerge == True:words = self._merge_results(words)
|
||||
lang_count = self._lang_count
|
||||
if lang_count and len(lang_count) > 0:
|
||||
lang_count = dict(sorted(lang_count.items(), key=lambda x: x[1], reverse=True))
|
||||
lang_count = list(lang_count.items())
|
||||
self._lang_count = lang_count
|
||||
return words
|
||||
|
||||
def setfilters(self, filters):
|
||||
# 当过滤器更改时，清除缓存
|
||||
# 필터가 변경되면 캐시를 지웁니다.
|
||||
# フィルタが変更されると、キャッシュがクリアされます
|
||||
# When the filter changes, clear the cache
|
||||
if self.Langfilters != filters:
|
||||
self._clears()
|
||||
self.Langfilters = filters
|
||||
|
||||
def getfilters(self):
|
||||
return self.Langfilters
|
||||
|
||||
def setPriorityThreshold(self, threshold:float):
|
||||
self.LangPriorityThreshold = threshold
|
||||
|
||||
def getPriorityThreshold(self):
|
||||
return self.LangPriorityThreshold
|
||||
|
||||
def getCounts(self):
|
||||
lang_count = self._lang_count
|
||||
if lang_count is not None:return lang_count
|
||||
text_langs = self._text_langs
|
||||
if text_langs is None or len(text_langs) == 0:return [("zh",0)]
|
||||
lang_counts = defaultdict(int)
|
||||
for d in text_langs:lang_counts[d['lang']] += int(len(d['text'])*2) if d['lang'] == "zh" else len(d['text'])
|
||||
lang_counts = dict(sorted(lang_counts.items(), key=lambda x: x[1], reverse=True))
|
||||
lang_counts = list(lang_counts.items())
|
||||
self._lang_count = lang_counts
|
||||
return lang_counts
|
||||
|
||||
def getTexts(self, text:str):
|
||||
if text is None or len(text.strip()) == 0:
|
||||
self._clears()
|
||||
return []
|
||||
# lasts
|
||||
text_langs = self._text_langs
|
||||
if self._text_lasts == text and text_langs is not None:return text_langs
|
||||
# parse
|
||||
self._text_waits = []
|
||||
self._lang_count = None
|
||||
self._text_lasts = text
|
||||
text = self._parse_symbols(text)
|
||||
self._text_langs = text
|
||||
return text
|
||||
|
||||
def classify(self, text:str):
|
||||
return self.getTexts(text)
|
||||
|
||||
def printList(langlist):
|
||||
"""
|
||||
功能：打印数组结果
|
||||
기능: 어레이 결과 인쇄
|
||||
機能:配列結果を印刷
|
||||
Function: Print array results
|
||||
"""
|
||||
print("\n===================【打印结果】===================")
|
||||
if langlist is None or len(langlist) == 0:
|
||||
print("无内容结果,No content result")
|
||||
return
|
||||
for line in langlist:
|
||||
print(line)
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
# -----------------------------------
|
||||
# 更新日志：新版本分词更加精准。
|
||||
# Changelog: The new version of the word segmentation is more accurate.
|
||||
# チェンジログ:新しいバージョンの単語セグメンテーションはより正確です。
|
||||
# Changelog: 분할이라는 단어의 새로운 버전이 더 정확합니다.
|
||||
# -----------------------------------
|
||||
|
||||
# 输入示例1：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
|
||||
# text = "“昨日は雨が降った，音楽、映画。。。”你今天学习日语了吗？春は桜の季節です。语种分词是语音合成必不可少的环节。言語分詞は音声合成に欠かせない環節である！"
|
||||
|
||||
# 输入示例2：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
|
||||
# text = "欢迎来玩。東京，は日本の首都です。欢迎来玩. 太好了!"
|
||||
|
||||
# 输入示例3：（包含日文，中文）Input Example 1: (including Japanese, Chinese)
|
||||
# text = "明日、私たちは海辺にバカンスに行きます。你会说日语吗：“中国語、話せますか” 你的日语真好啊！"
|
||||
|
||||
|
||||
# 输入示例4：（包含日文，中文，韩语，英文）Input Example 4: (including Japanese, Chinese, Korean, English)
|
||||
# text = "你的名字叫<ja>佐々木？<ja>吗？韩语中的안녕 오빠读什么呢？あなたの体育の先生は誰ですか? 此次发布会带来了四款iPhone 15系列机型和三款Apple Watch等一系列新品，这次的iPad Air采用了LCD屏幕"
|
||||
|
||||
|
||||
# 试验性支持："fr"法语 , "vi"越南语 , "ru"俄语 , "th"泰语。Experimental: Other language support.
|
||||
langsegment = LangSegment()
|
||||
langsegment.setfilters(["fr", "vi" , "ja", "zh", "ko", "en" , "ru" , "th"])
|
||||
text = """
|
||||
我喜欢在雨天里听音乐。
|
||||
I enjoy listening to music on rainy days.
|
||||
雨の日に音楽を聴くのが好きです。
|
||||
비 오는 날에 음악을 듣는 것을 즐깁니다。
|
||||
J'aime écouter de la musique les jours de pluie.
|
||||
Tôi thích nghe nhạc vào những ngày mưa.
|
||||
Мне нравится слушать музыку в дождливую погоду.
|
||||
ฉันชอบฟังเพลงในวันที่ฝนตก
|
||||
"""
|
||||
|
||||
|
||||
|
||||
# 进行分词：（接入TTS项目仅需一行代码调用）Segmentation: (Only one line of code is required to access the TTS project)
|
||||
langlist = langsegment.getTexts(text)
|
||||
printList(langlist)
|
||||
|
||||
|
||||
# 语种统计:Language statistics:
|
||||
print("\n===================【语种统计】===================")
|
||||
# 获取所有语种数组结果，根据内容字数降序排列
|
||||
# Get the array results in all languages, sorted in descending order according to the number of content words
|
||||
langCounts = langsegment.getCounts()
|
||||
print(langCounts , "\n")
|
||||
|
||||
# 根据结果获取内容的主要语种 (语言，字数含标点)
|
||||
# Get the main language of content based on the results (language, word count including punctuation)
|
||||
lang , count = langCounts[0]
|
||||
print(f"输入内容的主要语言为 = {lang} ，字数 = {count}")
|
||||
print("==================================================\n")
|
||||
|
||||
|
||||
# 分词输出：lang=语言，text=内容。Word output: lang = language, text = content
|
||||
# ===================【打印结果】===================
|
||||
# {'lang': 'zh', 'text': '你的名字叫'}
|
||||
# {'lang': 'ja', 'text': '佐々木？'}
|
||||
# {'lang': 'zh', 'text': '吗？韩语中的'}
|
||||
# {'lang': 'ko', 'text': '안녕 오빠'}
|
||||
# {'lang': 'zh', 'text': '读什么呢？'}
|
||||
# {'lang': 'ja', 'text': 'あなたの体育の先生は誰ですか?'}
|
||||
# {'lang': 'zh', 'text': ' 此次发布会带来了四款'}
|
||||
# {'lang': 'en', 'text': 'i Phone '}
|
||||
# {'lang': 'zh', 'text': '15系列机型和三款'}
|
||||
# {'lang': 'en', 'text': 'Apple Watch '}
|
||||
# {'lang': 'zh', 'text': '等一系列新品，这次的'}
|
||||
# {'lang': 'en', 'text': 'i Pad Air '}
|
||||
# {'lang': 'zh', 'text': '采用了'}
|
||||
# {'lang': 'en', 'text': 'L C D '}
|
||||
# {'lang': 'zh', 'text': '屏幕'}
|
||||
# ===================【语种统计】===================
|
||||
|
||||
# ===================【语种统计】===================
|
||||
# [('zh', 51), ('ja', 19), ('en', 18), ('ko', 5)]
|
||||
|
||||
# 输入内容的主要语言为 = zh ，字数 = 51
|
||||
# ==================================================
|
||||
# The main language of the input content is = zh, word count = 51
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,9 @@
|
||||
from .LangSegment import LangSegment
|
||||
|
||||
|
||||
# release
|
||||
__version__ = '0.3.5'
|
||||
|
||||
|
||||
# develop
|
||||
__develop__ = 'dev-0.0.1'
|
||||
@@ -0,0 +1,327 @@
|
||||
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# Digital processing from GPT_SoVITS num.py （thanks）
|
||||
"""
|
||||
Rules to verbalize numbers into Chinese characters.
|
||||
https://zh.wikipedia.org/wiki/中文数字#現代中文
|
||||
"""
|
||||
|
||||
import re
|
||||
from collections import OrderedDict
|
||||
from typing import List
|
||||
|
||||
DIGITS = {str(i): tran for i, tran in enumerate('零一二三四五六七八九')}
|
||||
UNITS = OrderedDict({
|
||||
1: '十',
|
||||
2: '百',
|
||||
3: '千',
|
||||
4: '万',
|
||||
8: '亿',
|
||||
})
|
||||
|
||||
COM_QUANTIFIERS = '(处|台|架|枚|趟|幅|平|方|堵|间|床|株|批|项|例|列|篇|栋|注|亩|封|艘|把|目|套|段|人|所|朵|匹|张|座|回|场|尾|条|个|首|阙|阵|网|炮|顶|丘|棵|只|支|袭|辆|挑|担|颗|壳|窠|曲|墙|群|腔|砣|座|客|贯|扎|捆|刀|令|打|手|罗|坡|山|岭|江|溪|钟|队|单|双|对|出|口|头|脚|板|跳|枝|件|贴|针|线|管|名|位|身|堂|课|本|页|家|户|层|丝|毫|厘|分|钱|两|斤|担|铢|石|钧|锱|忽|(千|毫|微)克|毫|厘|(公)分|分|寸|尺|丈|里|寻|常|铺|程|(千|分|厘|毫|微)米|米|撮|勺|合|升|斗|石|盘|碗|碟|叠|桶|笼|盆|盒|杯|钟|斛|锅|簋|篮|盘|桶|罐|瓶|壶|卮|盏|箩|箱|煲|啖|袋|钵|年|月|日|季|刻|时|周|天|秒|分|小时|旬|纪|岁|世|更|夜|春|夏|秋|冬|代|伏|辈|丸|泡|粒|颗|幢|堆|条|根|支|道|面|片|张|颗|块|元|(亿|千万|百万|万|千|百)|(亿|千万|百万|万|千|百|美|)元|(亿|千万|百万|万|千|百|十|)吨|(亿|千万|百万|万|千|百|)块|角|毛|分)'
|
||||
|
||||
# 分数表达式
|
||||
RE_FRAC = re.compile(r'(-?)(\d+)/(\d+)')
|
||||
|
||||
|
||||
def replace_frac(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
sign = match.group(1)
|
||||
nominator = match.group(2)
|
||||
denominator = match.group(3)
|
||||
sign: str = "负" if sign else ""
|
||||
nominator: str = num2str(nominator)
|
||||
denominator: str = num2str(denominator)
|
||||
result = f"{sign}{denominator}分之{nominator}"
|
||||
return result
|
||||
|
||||
|
||||
# 百分数表达式
|
||||
RE_PERCENTAGE = re.compile(r'(-?)(\d+(\.\d+)?)%')
|
||||
|
||||
|
||||
def replace_percentage(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
sign = match.group(1)
|
||||
percent = match.group(2)
|
||||
sign: str = "负" if sign else ""
|
||||
percent: str = num2str(percent)
|
||||
result = f"{sign}百分之{percent}"
|
||||
return result
|
||||
|
||||
|
||||
# 整数表达式
|
||||
# 带负号的整数 -10
|
||||
RE_INTEGER = re.compile(r'(-)' r'(\d+)')
|
||||
|
||||
|
||||
def replace_negative_num(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
sign = match.group(1)
|
||||
number = match.group(2)
|
||||
sign: str = "负" if sign else ""
|
||||
number: str = num2str(number)
|
||||
result = f"{sign}{number}"
|
||||
return result
|
||||
|
||||
|
||||
# 编号-无符号整形
|
||||
# 00078
|
||||
RE_DEFAULT_NUM = re.compile(r'\d{3}\d*')
|
||||
|
||||
|
||||
def replace_default_num(match):
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
number = match.group(0)
|
||||
return verbalize_digit(number, alt_one=True)
|
||||
|
||||
|
||||
# 加减乘除
|
||||
# RE_ASMD = re.compile(
|
||||
# r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))([\+\-\×÷=])((-?)((\d+)(\.\d+)?)|(\.(\d+)))')
|
||||
RE_ASMD = re.compile(
|
||||
r'((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))([\+\-\×÷=])((-?)((\d+)(\.\d+)?[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|(\.\d+[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*)|([A-Za-z][⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]*))')
|
||||
|
||||
asmd_map = {
|
||||
'+': '加',
|
||||
'-': '减',
|
||||
'×': '乘',
|
||||
'÷': '除',
|
||||
'=': '等于'
|
||||
}
|
||||
|
||||
def replace_asmd(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
result = match.group(1) + asmd_map[match.group(8)] + match.group(9)
|
||||
return result
|
||||
|
||||
|
||||
# 次方专项
|
||||
RE_POWER = re.compile(r'[⁰¹²³⁴⁵⁶⁷⁸⁹ˣʸⁿ]+')
|
||||
|
||||
power_map = {
|
||||
'⁰': '0',
|
||||
'¹': '1',
|
||||
'²': '2',
|
||||
'³': '3',
|
||||
'⁴': '4',
|
||||
'⁵': '5',
|
||||
'⁶': '6',
|
||||
'⁷': '7',
|
||||
'⁸': '8',
|
||||
'⁹': '9',
|
||||
'ˣ': 'x',
|
||||
'ʸ': 'y',
|
||||
'ⁿ': 'n'
|
||||
}
|
||||
|
||||
def replace_power(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
power_num = ""
|
||||
for m in match.group(0):
|
||||
power_num += power_map[m]
|
||||
result = "的" + power_num + "次方"
|
||||
return result
|
||||
|
||||
|
||||
# 数字表达式
|
||||
# 纯小数
|
||||
RE_DECIMAL_NUM = re.compile(r'(-?)((\d+)(\.\d+))' r'|(\.(\d+))')
|
||||
# 正整数 + 量词
|
||||
RE_POSITIVE_QUANTIFIERS = re.compile(r"(\d+)([多余几\+])?" + COM_QUANTIFIERS)
|
||||
RE_NUMBER = re.compile(r'(-?)((\d+)(\.\d+)?)' r'|(\.(\d+))')
|
||||
|
||||
|
||||
def replace_positive_quantifier(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
number = match.group(1)
|
||||
match_2 = match.group(2)
|
||||
if match_2 == "+":
|
||||
match_2 = "多"
|
||||
match_2: str = match_2 if match_2 else ""
|
||||
quantifiers: str = match.group(3)
|
||||
number: str = num2str(number)
|
||||
result = f"{number}{match_2}{quantifiers}"
|
||||
return result
|
||||
|
||||
|
||||
def replace_number(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
sign = match.group(1)
|
||||
number = match.group(2)
|
||||
pure_decimal = match.group(5)
|
||||
if pure_decimal:
|
||||
result = num2str(pure_decimal)
|
||||
else:
|
||||
sign: str = "负" if sign else ""
|
||||
number: str = num2str(number)
|
||||
result = f"{sign}{number}"
|
||||
return result
|
||||
|
||||
|
||||
# 范围表达式
|
||||
# match.group(1) and match.group(8) are copy from RE_NUMBER
|
||||
|
||||
RE_RANGE = re.compile(
|
||||
r"""
|
||||
(?<![\d\+\-\×÷=]) # 使用反向前瞻以确保数字范围之前没有其他数字和操作符
|
||||
((-?)((\d+)(\.\d+)?)) # 匹配范围起始的负数或正数（整数或小数）
|
||||
[-~] # 匹配范围分隔符
|
||||
((-?)((\d+)(\.\d+)?)) # 匹配范围结束的负数或正数（整数或小数）
|
||||
(?![\d\+\-\×÷=]) # 使用正向前瞻以确保数字范围之后没有其他数字和操作符
|
||||
""", re.VERBOSE)
|
||||
|
||||
|
||||
def replace_range(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
first, second = match.group(1), match.group(6)
|
||||
first = RE_NUMBER.sub(replace_number, first)
|
||||
second = RE_NUMBER.sub(replace_number, second)
|
||||
result = f"{first}到{second}"
|
||||
return result
|
||||
|
||||
|
||||
# ~至表达式
|
||||
RE_TO_RANGE = re.compile(
|
||||
r'((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)[~]((-?)((\d+)(\.\d+)?)|(\.(\d+)))(%|°C|℃|度|摄氏度|cm2|cm²|cm3|cm³|cm|db|ds|kg|km|m2|m²|m³|m3|ml|m|mm|s)')
|
||||
|
||||
def replace_to_range(match) -> str:
|
||||
"""
|
||||
Args:
|
||||
match (re.Match)
|
||||
Returns:
|
||||
str
|
||||
"""
|
||||
result = match.group(0).replace('~', '至')
|
||||
return result
|
||||
|
||||
|
||||
def _get_value(value_string: str, use_zero: bool=True) -> List[str]:
|
||||
stripped = value_string.lstrip('0')
|
||||
if len(stripped) == 0:
|
||||
return []
|
||||
elif len(stripped) == 1:
|
||||
if use_zero and len(stripped) < len(value_string):
|
||||
return [DIGITS['0'], DIGITS[stripped]]
|
||||
else:
|
||||
return [DIGITS[stripped]]
|
||||
else:
|
||||
largest_unit = next(
|
||||
power for power in reversed(UNITS.keys()) if power < len(stripped))
|
||||
first_part = value_string[:-largest_unit]
|
||||
second_part = value_string[-largest_unit:]
|
||||
return _get_value(first_part) + [UNITS[largest_unit]] + _get_value(
|
||||
second_part)
|
||||
|
||||
|
||||
def verbalize_cardinal(value_string: str) -> str:
|
||||
if not value_string:
|
||||
return ''
|
||||
|
||||
# 000 -> '零' , 0 -> '零'
|
||||
value_string = value_string.lstrip('0')
|
||||
if len(value_string) == 0:
|
||||
return DIGITS['0']
|
||||
|
||||
result_symbols = _get_value(value_string)
|
||||
# verbalized number starting with '一十*' is abbreviated as `十*`
|
||||
if len(result_symbols) >= 2 and result_symbols[0] == DIGITS[
|
||||
'1'] and result_symbols[1] == UNITS[1]:
|
||||
result_symbols = result_symbols[1:]
|
||||
return ''.join(result_symbols)
|
||||
|
||||
|
||||
def verbalize_digit(value_string: str, alt_one=False) -> str:
|
||||
result_symbols = [DIGITS[digit] for digit in value_string]
|
||||
result = ''.join(result_symbols)
|
||||
if alt_one:
|
||||
result = result.replace("一", "幺")
|
||||
return result
|
||||
|
||||
|
||||
def num2str(value_string: str) -> str:
|
||||
integer_decimal = value_string.split('.')
|
||||
if len(integer_decimal) == 1:
|
||||
integer = integer_decimal[0]
|
||||
decimal = ''
|
||||
elif len(integer_decimal) == 2:
|
||||
integer, decimal = integer_decimal
|
||||
else:
|
||||
raise ValueError(
|
||||
f"The value string: '${value_string}' has more than one point in it."
|
||||
)
|
||||
|
||||
result = verbalize_cardinal(integer)
|
||||
|
||||
decimal = decimal.rstrip('0')
|
||||
if decimal:
|
||||
# '.22' is verbalized as '零点二二'
|
||||
# '3.20' is verbalized as '三点二
|
||||
result = result if result else "零"
|
||||
result += '点' + verbalize_digit(decimal)
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
text = ""
|
||||
text = num2str(text)
|
||||
print(text)
|
||||
pass
|
||||
@@ -0,0 +1,475 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Tuple, List, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.utils import BaseOutput, is_torch_version
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
|
||||
|
||||
|
||||
from .attention import LinearTransformerBlock, t2i_modulate
|
||||
from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder
|
||||
|
||||
|
||||
def cross_norm(hidden_states, controlnet_input):
|
||||
# input N x T x c
|
||||
mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
|
||||
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)
|
||||
controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states
|
||||
return controlnet_input
|
||||
|
||||
|
||||
# Copied from transformers.models.mixtral.modeling_mixtral.MixtralRotaryEmbedding with Mixtral->Qwen2
|
||||
class Qwen2RotaryEmbedding(nn.Module):
|
||||
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
|
||||
super().__init__()
|
||||
|
||||
self.dim = dim
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.base = base
|
||||
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
# Build here to make `torch.jit.trace` work.
|
||||
self._set_cos_sin_cache(
|
||||
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
|
||||
)
|
||||
|
||||
def _set_cos_sin_cache(self, seq_len, device, dtype):
|
||||
self.max_seq_len_cached = seq_len
|
||||
t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
|
||||
|
||||
freqs = torch.outer(t, self.inv_freq)
|
||||
# Different from paper, but it uses a different permutation in order to obtain the same calculation
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
|
||||
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
|
||||
|
||||
def forward(self, x, seq_len=None):
|
||||
# x: [bs, num_attention_heads, seq_len, head_size]
|
||||
if seq_len > self.max_seq_len_cached:
|
||||
self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
|
||||
|
||||
return (
|
||||
self.cos_cached[:seq_len].to(dtype=x.dtype),
|
||||
self.sin_cached[:seq_len].to(dtype=x.dtype),
|
||||
)
|
||||
|
||||
|
||||
class T2IFinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer of Sana.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
|
||||
super().__init__()
|
||||
self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
|
||||
self.out_channels = out_channels
|
||||
self.patch_size = patch_size
|
||||
|
||||
def unpatchfy(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
width: int,
|
||||
):
|
||||
# 4 unpatchify
|
||||
new_height, new_width = 1, hidden_states.size(1)
|
||||
hidden_states = hidden_states.reshape(
|
||||
shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
|
||||
).contiguous()
|
||||
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
|
||||
output = hidden_states.reshape(
|
||||
shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
|
||||
).contiguous()
|
||||
if width > new_width:
|
||||
output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
|
||||
elif width < new_width:
|
||||
output = output[:, :, :, :width]
|
||||
return output
|
||||
|
||||
def forward(self, x, t, output_length):
|
||||
shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)
|
||||
x = t2i_modulate(self.norm_final(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
# unpatchify
|
||||
output = self.unpatchfy(x, output_length)
|
||||
return output
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""2D Image to Patch Embedding"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
height=16,
|
||||
width=4096,
|
||||
patch_size=(16, 1),
|
||||
in_channels=8,
|
||||
embed_dim=1152,
|
||||
bias=True,
|
||||
):
|
||||
super().__init__()
|
||||
patch_size_h, patch_size_w = patch_size
|
||||
self.early_conv_layers = nn.Sequential(
|
||||
nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),
|
||||
torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True),
|
||||
nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)
|
||||
)
|
||||
self.patch_size = patch_size
|
||||
self.height, self.width = height // patch_size_h, width // patch_size_w
|
||||
self.base_size = self.width
|
||||
|
||||
def forward(self, latent):
|
||||
# early convolutions, N x C x H x W -> N x 256 * sqrt(patch_size) x H/patch_size x W/patch_size
|
||||
latent = self.early_conv_layers(latent)
|
||||
latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC
|
||||
return latent
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transformer2DModelOutput(BaseOutput):
|
||||
|
||||
sample: torch.FloatTensor
|
||||
proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
|
||||
|
||||
|
||||
class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: Optional[int] = 8,
|
||||
num_layers: int = 28,
|
||||
inner_dim: int = 1536,
|
||||
attention_head_dim: int = 64,
|
||||
num_attention_heads: int = 24,
|
||||
mlp_ratio: float = 4.0,
|
||||
out_channels: int = 8,
|
||||
max_position: int = 32768,
|
||||
rope_theta: float = 1000000.0,
|
||||
speaker_embedding_dim: int = 512,
|
||||
text_embedding_dim: int = 768,
|
||||
ssl_encoder_depths: List[int] = [9, 9],
|
||||
ssl_names: List[str] = ["mert", "m-hubert"],
|
||||
ssl_latent_dims: List[int] = [1024, 768],
|
||||
lyric_encoder_vocab_size: int = 6681,
|
||||
lyric_hidden_size: int = 1024,
|
||||
patch_size: List[int] = [16, 1],
|
||||
max_height: int = 16,
|
||||
max_width: int = 4096,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
self.inner_dim = inner_dim
|
||||
self.out_channels = out_channels
|
||||
self.max_position = max_position
|
||||
self.patch_size = patch_size
|
||||
|
||||
self.rope_theta = rope_theta
|
||||
|
||||
self.rotary_emb = Qwen2RotaryEmbedding(
|
||||
dim=self.attention_head_dim,
|
||||
max_position_embeddings=self.max_position,
|
||||
base=self.rope_theta,
|
||||
)
|
||||
|
||||
# 2. Define input layers
|
||||
self.in_channels = in_channels
|
||||
|
||||
# 3. Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
LinearTransformerBlock(
|
||||
dim=self.inner_dim,
|
||||
num_attention_heads=self.num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
mlp_ratio=mlp_ratio,
|
||||
add_cross_attention=True,
|
||||
add_cross_attention_dim=self.inner_dim,
|
||||
)
|
||||
for i in range(self.config.num_layers)
|
||||
]
|
||||
)
|
||||
self.num_layers = num_layers
|
||||
|
||||
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)
|
||||
self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True))
|
||||
|
||||
# speaker
|
||||
self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)
|
||||
|
||||
# genre
|
||||
self.genre_embedder = nn.Linear(text_embedding_dim, self.inner_dim)
|
||||
|
||||
# lyric
|
||||
self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)
|
||||
self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)
|
||||
self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
|
||||
|
||||
projector_dim = 2 * self.inner_dim
|
||||
|
||||
self.projectors = nn.ModuleList([
|
||||
nn.Sequential(
|
||||
nn.Linear(self.inner_dim, projector_dim),
|
||||
nn.SiLU(),
|
||||
nn.Linear(projector_dim, projector_dim),
|
||||
nn.SiLU(),
|
||||
nn.Linear(projector_dim, ssl_dim),
|
||||
) for ssl_dim in ssl_latent_dims
|
||||
])
|
||||
|
||||
self.ssl_latent_dims = ssl_latent_dims
|
||||
self.ssl_encoder_depths = ssl_encoder_depths
|
||||
|
||||
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
|
||||
self.ssl_names = ssl_names
|
||||
|
||||
self.proj_in = PatchEmbed(
|
||||
height=max_height,
|
||||
width=max_width,
|
||||
patch_size=patch_size,
|
||||
embed_dim=self.inner_dim,
|
||||
bias=True,
|
||||
)
|
||||
|
||||
self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
|
||||
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
|
||||
"""
|
||||
Sets the attention processor to use [feed forward
|
||||
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
|
||||
|
||||
Parameters:
|
||||
chunk_size (`int`, *optional*):
|
||||
The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
|
||||
over each tensor of dim=`dim`.
|
||||
dim (`int`, *optional*, defaults to `0`):
|
||||
The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
|
||||
or dim=1 (sequence length).
|
||||
"""
|
||||
if dim not in [0, 1]:
|
||||
raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
|
||||
|
||||
# By default chunk size is 1
|
||||
chunk_size = chunk_size or 1
|
||||
|
||||
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
|
||||
if hasattr(module, "set_chunk_feed_forward"):
|
||||
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_feed_forward(child, chunk_size, dim)
|
||||
|
||||
for module in self.children():
|
||||
fn_recursive_feed_forward(module, chunk_size, dim)
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def forward_lyric_encoder(
|
||||
self,
|
||||
lyric_token_idx: Optional[torch.LongTensor] = None,
|
||||
lyric_mask: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
# N x T x D
|
||||
lyric_embs = self.lyric_embs(lyric_token_idx)
|
||||
prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
|
||||
prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
|
||||
return prompt_prenet_out
|
||||
|
||||
def encode(
|
||||
self,
|
||||
encoder_text_hidden_states: Optional[torch.Tensor] = None,
|
||||
text_attention_mask: Optional[torch.LongTensor] = None,
|
||||
speaker_embeds: Optional[torch.FloatTensor] = None,
|
||||
lyric_token_idx: Optional[torch.LongTensor] = None,
|
||||
lyric_mask: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
|
||||
bs = encoder_text_hidden_states.shape[0]
|
||||
device = encoder_text_hidden_states.device
|
||||
|
||||
# speaker embedding
|
||||
encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)
|
||||
speaker_mask = torch.ones(bs, 1, device=device)
|
||||
|
||||
# genre embedding
|
||||
encoder_text_hidden_states = self.genre_embedder(encoder_text_hidden_states)
|
||||
|
||||
# lyric
|
||||
encoder_lyric_hidden_states = self.forward_lyric_encoder(
|
||||
lyric_token_idx=lyric_token_idx,
|
||||
lyric_mask=lyric_mask,
|
||||
)
|
||||
|
||||
encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
|
||||
encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
|
||||
return encoder_hidden_states, encoder_hidden_mask
|
||||
|
||||
def decode(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
encoder_hidden_mask: torch.Tensor,
|
||||
timestep: Optional[torch.Tensor],
|
||||
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
|
||||
output_length: int = 0,
|
||||
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
|
||||
controlnet_scale: Union[float, torch.Tensor] = 1.0,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
|
||||
embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
|
||||
temb = self.t_block(embedded_timestep)
|
||||
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
|
||||
# controlnet logic
|
||||
if block_controlnet_hidden_states is not None:
|
||||
control_condi = cross_norm(hidden_states, block_controlnet_hidden_states)
|
||||
hidden_states = hidden_states + control_condi * controlnet_scale
|
||||
|
||||
inner_hidden_states = []
|
||||
|
||||
rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
|
||||
encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
|
||||
|
||||
for index_block, block in enumerate(self.transformer_blocks):
|
||||
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_hidden_mask,
|
||||
rotary_freqs_cis=rotary_freqs_cis,
|
||||
rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
|
||||
temb=temb,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
hidden_states = block(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_hidden_mask,
|
||||
rotary_freqs_cis=rotary_freqs_cis,
|
||||
rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
|
||||
temb=temb,
|
||||
)
|
||||
|
||||
for ssl_encoder_depth in self.ssl_encoder_depths:
|
||||
if index_block == ssl_encoder_depth:
|
||||
inner_hidden_states.append(hidden_states)
|
||||
|
||||
proj_losses = []
|
||||
if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:
|
||||
|
||||
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):
|
||||
if ssl_hidden_state is None:
|
||||
continue
|
||||
# 1. N x T x D1 -> N x D x D2
|
||||
est_ssl_hidden_state = projector(inner_hidden_state)
|
||||
# 3. projection loss
|
||||
bs = inner_hidden_state.shape[0]
|
||||
proj_loss = 0.0
|
||||
for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):
|
||||
# 2. interpolate
|
||||
z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)
|
||||
|
||||
z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
|
||||
z = torch.nn.functional.normalize(z, dim=-1)
|
||||
# T x d -> T x 1 -> 1
|
||||
target = torch.ones(z.shape[0], device=z.device)
|
||||
proj_loss += self.cosine_loss(z, z_tilde, target)
|
||||
proj_losses.append((ssl_name, proj_loss / bs))
|
||||
|
||||
output = self.final_layer(hidden_states, embedded_timestep, output_length)
|
||||
if not return_dict:
|
||||
return (output, proj_losses)
|
||||
|
||||
return Transformer2DModelOutput(sample=output, proj_losses=proj_losses)
|
||||
|
||||
# @torch.compile
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
attention_mask: torch.Tensor,
|
||||
encoder_text_hidden_states: Optional[torch.Tensor] = None,
|
||||
text_attention_mask: Optional[torch.LongTensor] = None,
|
||||
speaker_embeds: Optional[torch.FloatTensor] = None,
|
||||
lyric_token_idx: Optional[torch.LongTensor] = None,
|
||||
lyric_mask: Optional[torch.LongTensor] = None,
|
||||
timestep: Optional[torch.Tensor] = None,
|
||||
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
|
||||
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
|
||||
controlnet_scale: Union[float, torch.Tensor] = 1.0,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
encoder_hidden_states, encoder_hidden_mask = self.encode(
|
||||
encoder_text_hidden_states=encoder_text_hidden_states,
|
||||
text_attention_mask=text_attention_mask,
|
||||
speaker_embeds=speaker_embeds,
|
||||
lyric_token_idx=lyric_token_idx,
|
||||
lyric_mask=lyric_mask,
|
||||
)
|
||||
|
||||
output_length = hidden_states.shape[-1]
|
||||
|
||||
output = self.decode(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_hidden_mask=encoder_hidden_mask,
|
||||
timestep=timestep,
|
||||
ssl_hidden_states=ssl_hidden_states,
|
||||
output_length=output_length,
|
||||
block_controlnet_hidden_states=block_controlnet_hidden_states,
|
||||
controlnet_scale=controlnet_scale,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,319 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.utils import logging
|
||||
from diffusers.models.normalization import RMSNorm
|
||||
|
||||
|
||||
try:
|
||||
# from .dcformer import DCMHAttention
|
||||
from .customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
|
||||
except ImportError:
|
||||
# from dcformer import DCMHAttention
|
||||
from customer_attention_processor import Attention, CustomLiteLAProcessor2_0, CustomerAttnProcessor2_0
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def val2list(x: list or tuple or any, repeat_time=1) -> list: # type: ignore
|
||||
"""Repeat `val` for `repeat_time` times and return the list or val if list/tuple."""
|
||||
if isinstance(x, (list, tuple)):
|
||||
return list(x)
|
||||
return [x for _ in range(repeat_time)]
|
||||
|
||||
|
||||
def val2tuple(x: list or tuple or any, min_len: int = 1, idx_repeat: int = -1) -> tuple: # type: ignore
|
||||
"""Return tuple with min_len by repeating element at idx_repeat."""
|
||||
# convert to list first
|
||||
x = val2list(x)
|
||||
|
||||
# repeat elements if necessary
|
||||
if len(x) > 0:
|
||||
x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
|
||||
|
||||
return tuple(x)
|
||||
|
||||
|
||||
def t2i_modulate(x, shift, scale):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
|
||||
if isinstance(kernel_size, tuple):
|
||||
return tuple([get_same_padding(ks) for ks in kernel_size])
|
||||
else:
|
||||
assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
|
||||
return kernel_size // 2
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
padding: Union[int, None] = None,
|
||||
use_bias=False,
|
||||
norm=None,
|
||||
act=None,
|
||||
):
|
||||
super().__init__()
|
||||
if padding is None:
|
||||
padding = get_same_padding(kernel_size)
|
||||
padding *= dilation
|
||||
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
self.kernel_size = kernel_size
|
||||
self.stride = stride
|
||||
self.dilation = dilation
|
||||
self.groups = groups
|
||||
self.padding = padding
|
||||
self.use_bias = use_bias
|
||||
|
||||
self.conv = nn.Conv1d(
|
||||
in_dim,
|
||||
out_dim,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=use_bias,
|
||||
)
|
||||
if norm is not None:
|
||||
self.norm = RMSNorm(out_dim, elementwise_affine=False)
|
||||
else:
|
||||
self.norm = None
|
||||
if act is not None:
|
||||
self.act = nn.SiLU(inplace=True)
|
||||
else:
|
||||
self.act = None
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = self.conv(x)
|
||||
if self.norm:
|
||||
x = self.norm(x)
|
||||
if self.act:
|
||||
x = self.act(x)
|
||||
return x
|
||||
|
||||
|
||||
class GLUMBConv(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_features: int,
|
||||
hidden_features: int,
|
||||
out_feature=None,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding: Union[int, None] = None,
|
||||
use_bias=False,
|
||||
norm=(None, None, None),
|
||||
act=("silu", "silu", None),
|
||||
dilation=1,
|
||||
):
|
||||
out_feature = out_feature or in_features
|
||||
super().__init__()
|
||||
use_bias = val2tuple(use_bias, 3)
|
||||
norm = val2tuple(norm, 3)
|
||||
act = val2tuple(act, 3)
|
||||
|
||||
self.glu_act = nn.SiLU(inplace=False)
|
||||
self.inverted_conv = ConvLayer(
|
||||
in_features,
|
||||
hidden_features * 2,
|
||||
1,
|
||||
use_bias=use_bias[0],
|
||||
norm=norm[0],
|
||||
act=act[0],
|
||||
)
|
||||
self.depth_conv = ConvLayer(
|
||||
hidden_features * 2,
|
||||
hidden_features * 2,
|
||||
kernel_size,
|
||||
stride=stride,
|
||||
groups=hidden_features * 2,
|
||||
padding=padding,
|
||||
use_bias=use_bias[1],
|
||||
norm=norm[1],
|
||||
act=None,
|
||||
dilation=dilation,
|
||||
)
|
||||
self.point_conv = ConvLayer(
|
||||
hidden_features,
|
||||
out_feature,
|
||||
1,
|
||||
use_bias=use_bias[2],
|
||||
norm=norm[2],
|
||||
act=act[2],
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x = x.transpose(1, 2)
|
||||
x = self.inverted_conv(x)
|
||||
x = self.depth_conv(x)
|
||||
|
||||
x, gate = torch.chunk(x, 2, dim=1)
|
||||
gate = self.glu_act(gate)
|
||||
x = x * gate
|
||||
|
||||
x = self.point_conv(x)
|
||||
x = x.transpose(1, 2)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class LinearTransformerBlock(nn.Module):
|
||||
"""
|
||||
A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
use_adaln_single=True,
|
||||
cross_attention_dim=None,
|
||||
added_kv_proj_dim=None,
|
||||
context_pre_only=False,
|
||||
mlp_ratio=4.0,
|
||||
add_cross_attention=False,
|
||||
add_cross_attention_dim=None,
|
||||
qk_norm=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.norm1 = RMSNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.attn = Attention(
|
||||
query_dim=dim,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
added_kv_proj_dim=added_kv_proj_dim,
|
||||
dim_head=attention_head_dim,
|
||||
heads=num_attention_heads,
|
||||
out_dim=dim,
|
||||
bias=True,
|
||||
qk_norm=qk_norm,
|
||||
processor=CustomLiteLAProcessor2_0(),
|
||||
)
|
||||
|
||||
self.add_cross_attention = add_cross_attention
|
||||
self.context_pre_only = context_pre_only
|
||||
|
||||
if add_cross_attention and add_cross_attention_dim is not None:
|
||||
self.cross_attn = Attention(
|
||||
query_dim=dim,
|
||||
cross_attention_dim=add_cross_attention_dim,
|
||||
added_kv_proj_dim=add_cross_attention_dim,
|
||||
dim_head=attention_head_dim,
|
||||
heads=num_attention_heads,
|
||||
out_dim=dim,
|
||||
context_pre_only=context_pre_only,
|
||||
bias=True,
|
||||
qk_norm=qk_norm,
|
||||
processor=CustomerAttnProcessor2_0(),
|
||||
)
|
||||
|
||||
self.norm2 = RMSNorm(dim, 1e-06, elementwise_affine=False)
|
||||
|
||||
self.ff = GLUMBConv(
|
||||
in_features=dim,
|
||||
hidden_features=int(dim * mlp_ratio),
|
||||
use_bias=(True, True, False),
|
||||
norm=(None, None, None),
|
||||
act=("silu", "silu", None),
|
||||
)
|
||||
self.use_adaln_single = use_adaln_single
|
||||
if use_adaln_single:
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
encoder_hidden_states: torch.FloatTensor = None,
|
||||
attention_mask: torch.FloatTensor = None,
|
||||
encoder_attention_mask: torch.FloatTensor = None,
|
||||
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
|
||||
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
|
||||
temb: torch.FloatTensor = None,
|
||||
):
|
||||
|
||||
N = hidden_states.shape[0]
|
||||
|
||||
# step 1: AdaLN single
|
||||
if self.use_adaln_single:
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None] + temb.reshape(N, 6, -1)
|
||||
).chunk(6, dim=1)
|
||||
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
if self.use_adaln_single:
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
|
||||
|
||||
# step 2: attention
|
||||
if not self.add_cross_attention:
|
||||
attn_output, encoder_hidden_states = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
rotary_freqs_cis=rotary_freqs_cis,
|
||||
rotary_freqs_cis_cross=rotary_freqs_cis_cross,
|
||||
)
|
||||
else:
|
||||
attn_output, _ = self.attn(
|
||||
hidden_states=norm_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=None,
|
||||
encoder_attention_mask=None,
|
||||
rotary_freqs_cis=rotary_freqs_cis,
|
||||
rotary_freqs_cis_cross=None,
|
||||
)
|
||||
|
||||
if self.use_adaln_single:
|
||||
attn_output = gate_msa * attn_output
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
if self.add_cross_attention:
|
||||
attn_output = self.cross_attn(
|
||||
hidden_states=hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
rotary_freqs_cis=rotary_freqs_cis,
|
||||
rotary_freqs_cis_cross=rotary_freqs_cis_cross,
|
||||
)
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# step 3: add norm
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
if self.use_adaln_single:
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
|
||||
|
||||
# step 4: feed forward
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
if self.use_adaln_single:
|
||||
ff_output = gate_mlp * ff_output
|
||||
|
||||
hidden_states = hidden_states + ff_output
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"_class_name": "Transformer2DModel",
|
||||
"_diffusers_version": "0.27.2",
|
||||
"in_channels": 8,
|
||||
"num_layers": 24,
|
||||
"inner_dim": 2560,
|
||||
"attention_head_dim": 128,
|
||||
"num_attention_heads": 20,
|
||||
"mlp_ratio": 2.5,
|
||||
"out_channels": 8,
|
||||
"max_position": 32768,
|
||||
"rope_theta": 1000000.0,
|
||||
"speaker_embedding_dim": 512,
|
||||
"text_embedding_dim": 768,
|
||||
"ssl_encoder_depths": [8, 8],
|
||||
"ssl_names": ["mert", "m-hubert"],
|
||||
"ssl_latent_dims": [1024, 768],
|
||||
"patch_size": [16, 1],
|
||||
"max_height": 16,
|
||||
"max_width": 32768,
|
||||
"lyric_encoder_vocab_size": 6693,
|
||||
"lyric_hidden_size": 1024
|
||||
}
|
||||
@@ -0,0 +1,339 @@
|
||||
# Copyright 2024 The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from typing import Optional, Union, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from diffusers.utils import logging
|
||||
from diffusers.models.attention_processor import Attention
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
class CustomLiteLAProcessor2_0:
|
||||
"""Attention processor used typically in processing the SD3-like self-attention projections. add rms norm for query and key and apply RoPE"""
|
||||
|
||||
def __init__(self):
|
||||
self.kernel_func = nn.ReLU(inplace=False)
|
||||
self.eps = 1e-15
|
||||
self.pad_val = 1.0
|
||||
|
||||
def apply_rotary_emb(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
|
||||
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
|
||||
reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
|
||||
tensors contain rotary embeddings and are returned as real tensors.
|
||||
|
||||
Args:
|
||||
x (`torch.Tensor`):
|
||||
Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
|
||||
freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
|
||||
"""
|
||||
cos, sin = freqs_cis # [S, D]
|
||||
cos = cos[None, None]
|
||||
sin = sin[None, None]
|
||||
cos, sin = cos.to(x.device), sin.to(x.device)
|
||||
|
||||
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
|
||||
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
|
||||
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
|
||||
|
||||
return out
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states: torch.FloatTensor,
|
||||
encoder_hidden_states: torch.FloatTensor = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
|
||||
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> torch.FloatTensor:
|
||||
hidden_states_len = hidden_states.shape[1]
|
||||
|
||||
input_ndim = hidden_states.ndim
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
if encoder_hidden_states is not None:
|
||||
context_input_ndim = encoder_hidden_states.ndim
|
||||
if context_input_ndim == 4:
|
||||
batch_size, channel, height, width = encoder_hidden_states.shape
|
||||
encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size = hidden_states.shape[0]
|
||||
|
||||
# `sample` projections.
|
||||
dtype = hidden_states.dtype
|
||||
query = attn.to_q(hidden_states)
|
||||
key = attn.to_k(hidden_states)
|
||||
value = attn.to_v(hidden_states)
|
||||
|
||||
# `context` projections.
|
||||
has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
|
||||
if encoder_hidden_states is not None and has_encoder_hidden_state_proj:
|
||||
encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
|
||||
encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
|
||||
encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
|
||||
|
||||
# attention
|
||||
if not attn.is_cross_attention:
|
||||
query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)
|
||||
key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
|
||||
value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
|
||||
else:
|
||||
query = hidden_states
|
||||
key = encoder_hidden_states
|
||||
value = encoder_hidden_states
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
|
||||
key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)
|
||||
value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
|
||||
|
||||
# RoPE需要 [B, H, S, D] 输入
|
||||
# 此时 query是 [B, H, D, S], 需要转成 [B, H, S, D] 才能应用RoPE
|
||||
query = query.permute(0, 1, 3, 2) # [B, H, S, D] (从 [B, H, D, S])
|
||||
|
||||
# Apply query and key normalization if needed
|
||||
if attn.norm_q is not None:
|
||||
query = attn.norm_q(query)
|
||||
if attn.norm_k is not None:
|
||||
key = attn.norm_k(key)
|
||||
|
||||
# Apply RoPE if needed
|
||||
if rotary_freqs_cis is not None:
|
||||
query = self.apply_rotary_emb(query, rotary_freqs_cis)
|
||||
if not attn.is_cross_attention:
|
||||
key = self.apply_rotary_emb(key, rotary_freqs_cis)
|
||||
elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
|
||||
key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
|
||||
|
||||
# 此时 query是 [B, H, S, D]，需要还原成 [B, H, D, S]
|
||||
query = query.permute(0, 1, 3, 2) # [B, H, D, S]
|
||||
|
||||
if attention_mask is not None:
|
||||
# attention_mask: [B, S] -> [B, 1, S, 1]
|
||||
attention_mask = attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S, 1]
|
||||
query = query * attention_mask.permute(0, 1, 3, 2) # [B, H, S, D] * [B, 1, S, 1]
|
||||
if not attn.is_cross_attention:
|
||||
key = key * attention_mask # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
|
||||
value = value * attention_mask.permute(0, 1, 3, 2) # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
|
||||
|
||||
if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
|
||||
encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype) # [B, 1, S_enc, 1]
|
||||
# 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc]
|
||||
key = key * encoder_attention_mask # [B, h, S_enc, D] * [B, 1, S_enc, 1]
|
||||
value = value * encoder_attention_mask.permute(0, 1, 3, 2) # [B, h, D, S_enc] * [B, 1, 1, S_enc]
|
||||
|
||||
query = self.kernel_func(query)
|
||||
key = self.kernel_func(key)
|
||||
|
||||
query, key, value = query.float(), key.float(), value.float()
|
||||
|
||||
value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)
|
||||
|
||||
vk = torch.matmul(value, key)
|
||||
|
||||
hidden_states = torch.matmul(vk, query)
|
||||
|
||||
if hidden_states.dtype in [torch.float16, torch.bfloat16]:
|
||||
hidden_states = hidden_states.float()
|
||||
|
||||
hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
|
||||
|
||||
hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)
|
||||
|
||||
hidden_states = hidden_states.to(dtype)
|
||||
if encoder_hidden_states is not None:
|
||||
encoder_hidden_states = encoder_hidden_states.to(dtype)
|
||||
|
||||
# Split the attention outputs.
|
||||
if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:
|
||||
hidden_states, encoder_hidden_states = (
|
||||
hidden_states[:, : hidden_states_len],
|
||||
hidden_states[:, hidden_states_len:],
|
||||
)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):
|
||||
encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
if encoder_hidden_states is not None and context_input_ndim == 4:
|
||||
encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if torch.get_autocast_gpu_dtype() == torch.float16:
|
||||
hidden_states = hidden_states.clip(-65504, 65504)
|
||||
if encoder_hidden_states is not None:
|
||||
encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
|
||||
|
||||
return hidden_states, encoder_hidden_states
|
||||
|
||||
|
||||
class CustomerAttnProcessor2_0:
|
||||
r"""
|
||||
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
if not hasattr(F, "scaled_dot_product_attention"):
|
||||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||||
|
||||
def apply_rotary_emb(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
|
||||
to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
|
||||
reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
|
||||
tensors contain rotary embeddings and are returned as real tensors.
|
||||
|
||||
Args:
|
||||
x (`torch.Tensor`):
|
||||
Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
|
||||
freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
|
||||
"""
|
||||
cos, sin = freqs_cis # [S, D]
|
||||
cos = cos[None, None]
|
||||
sin = sin[None, None]
|
||||
cos, sin = cos.to(x.device), sin.to(x.device)
|
||||
|
||||
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
|
||||
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
|
||||
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
|
||||
|
||||
return out
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
hidden_states: torch.FloatTensor,
|
||||
encoder_hidden_states: torch.FloatTensor = None,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
|
||||
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
|
||||
*args,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
|
||||
residual = hidden_states
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
|
||||
has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
|
||||
|
||||
if attn.group_norm is not None:
|
||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
elif attn.norm_cross:
|
||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||||
|
||||
key = attn.to_k(encoder_hidden_states)
|
||||
value = attn.to_v(encoder_hidden_states)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
if attn.norm_q is not None:
|
||||
query = attn.norm_q(query)
|
||||
if attn.norm_k is not None:
|
||||
key = attn.norm_k(key)
|
||||
|
||||
# Apply RoPE if needed
|
||||
if rotary_freqs_cis is not None:
|
||||
query = self.apply_rotary_emb(query, rotary_freqs_cis)
|
||||
if not attn.is_cross_attention:
|
||||
key = self.apply_rotary_emb(key, rotary_freqs_cis)
|
||||
elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
|
||||
key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
|
||||
|
||||
if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
|
||||
# attention_mask: N x S1
|
||||
# encoder_attention_mask: N x S2
|
||||
# cross attention 整合attention_mask和encoder_attention_mask
|
||||
combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
|
||||
attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)
|
||||
attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)
|
||||
|
||||
elif not attn.is_cross_attention and attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
# scaled_dot_product_attention expects attention_mask shape to be
|
||||
# (batch, heads, source_length, target_length)
|
||||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||||
|
||||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||||
hidden_states = F.scaled_dot_product_attention(
|
||||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||||
)
|
||||
|
||||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
hidden_states = hidden_states.to(query.dtype)
|
||||
|
||||
# linear proj
|
||||
hidden_states = attn.to_out[0](hidden_states)
|
||||
# dropout
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||||
|
||||
if attn.residual_connection:
|
||||
hidden_states = hidden_states + residual
|
||||
|
||||
hidden_states = hidden_states / attn.rescale_output_factor
|
||||
|
||||
return hidden_states
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,66 @@
|
||||
import re
|
||||
from opencc import OpenCC
|
||||
|
||||
|
||||
t2s_converter = OpenCC('t2s')
|
||||
s2t_converter = OpenCC('s2t')
|
||||
|
||||
|
||||
EMOJI_PATTERN = re.compile(
|
||||
"["
|
||||
"\U0001F600-\U0001F64F" # Emoticons
|
||||
"]+", flags=re.UNICODE
|
||||
)
|
||||
|
||||
# 创建一个翻译表，用于替换和移除字符
|
||||
TRANSLATION_TABLE = str.maketrans({
|
||||
'-': ' ', # 将 '-' 替换为空格
|
||||
',': None,
|
||||
'.': None,
|
||||
'，': None,
|
||||
'。': None,
|
||||
'!': None,
|
||||
'！': None,
|
||||
'?': None,
|
||||
'？': None,
|
||||
'…': None,
|
||||
';': None,
|
||||
'；': None,
|
||||
':': None,
|
||||
'：': None,
|
||||
'\u3000': ' ', # 将全角空格替换为空格
|
||||
})
|
||||
|
||||
# 替换括号中的内容，包括中括号和小括号
|
||||
BACKSLASH_PATTERN = re.compile(r'\(.*?\)|\[.*?\]')
|
||||
|
||||
SPACE_PATTERN = re.compile('(?<!^)\s+(?!$)')
|
||||
|
||||
|
||||
def normalize_text(text, language, strip=True):
|
||||
"""
|
||||
对文本进行标准化处理，去除标点符号，转为小写（如果适用）
|
||||
"""
|
||||
# Step 1: 替换 '-' 为 ' ' 并移除标点符号
|
||||
text = text.translate(TRANSLATION_TABLE)
|
||||
|
||||
# Step 2: 移除表情符号
|
||||
text = EMOJI_PATTERN.sub('', text)
|
||||
|
||||
# Step 3: 连续空白字符替换为单个空格，首位除外
|
||||
text = SPACE_PATTERN.sub(' ', text)
|
||||
|
||||
# Step 4: 去除首尾空白字符（如果需要）
|
||||
if strip:
|
||||
text = text.strip()
|
||||
|
||||
# Step 5: 转为小写
|
||||
text = text.lower()
|
||||
|
||||
# Step 6: 多语言转换
|
||||
if language == "zh":
|
||||
text = t2s_converter.convert(text)
|
||||
if language == "yue":
|
||||
text = s2t_converter.convert(text)
|
||||
# 其他语言根据需要添加
|
||||
return text
|
||||
@@ -0,0 +1,883 @@
|
||||
import os
|
||||
import re
|
||||
import textwrap
|
||||
from functools import cached_property
|
||||
|
||||
import pypinyin
|
||||
import torch
|
||||
from hangul_romanize import Transliter
|
||||
from hangul_romanize.rule import academic
|
||||
from num2words import num2words
|
||||
from spacy.lang.ar import Arabic
|
||||
from spacy.lang.en import English
|
||||
from spacy.lang.es import Spanish
|
||||
from spacy.lang.ja import Japanese
|
||||
from spacy.lang.zh import Chinese
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
from .zh_num2words import TextNorm as zh_num2words
|
||||
from typing import Dict, List, Optional, Set, Union
|
||||
|
||||
|
||||
#copy from https://github.com/coqui-ai/TTS/blob/dbf1a08a0d4e47fdad6172e433eeb34bc6b13b4e/TTS/tts/layers/xtts/tokenizer.py
|
||||
def get_spacy_lang(lang):
|
||||
if lang == "zh":
|
||||
return Chinese()
|
||||
elif lang == "ja":
|
||||
return Japanese()
|
||||
elif lang == "ar":
|
||||
return Arabic()
|
||||
elif lang == "es":
|
||||
return Spanish()
|
||||
else:
|
||||
# For most languages, Enlish does the job
|
||||
return English()
|
||||
|
||||
|
||||
def split_sentence(text, lang, text_split_length=250):
|
||||
"""Preprocess the input text"""
|
||||
text_splits = []
|
||||
if text_split_length is not None and len(text) >= text_split_length:
|
||||
text_splits.append("")
|
||||
nlp = get_spacy_lang(lang)
|
||||
nlp.add_pipe("sentencizer")
|
||||
doc = nlp(text)
|
||||
for sentence in doc.sents:
|
||||
if len(text_splits[-1]) + len(str(sentence)) <= text_split_length:
|
||||
# if the last sentence + the current sentence is less than the text_split_length
|
||||
# then add the current sentence to the last sentence
|
||||
text_splits[-1] += " " + str(sentence)
|
||||
text_splits[-1] = text_splits[-1].lstrip()
|
||||
elif len(str(sentence)) > text_split_length:
|
||||
# if the current sentence is greater than the text_split_length
|
||||
for line in textwrap.wrap(
|
||||
str(sentence),
|
||||
width=text_split_length,
|
||||
drop_whitespace=True,
|
||||
break_on_hyphens=False,
|
||||
tabsize=1,
|
||||
):
|
||||
text_splits.append(str(line))
|
||||
else:
|
||||
text_splits.append(str(sentence))
|
||||
|
||||
if len(text_splits) > 1:
|
||||
if text_splits[0] == "":
|
||||
del text_splits[0]
|
||||
else:
|
||||
text_splits = [text.lstrip()]
|
||||
|
||||
return text_splits
|
||||
|
||||
|
||||
_whitespace_re = re.compile(r"\s+")
|
||||
|
||||
# List of (regular expression, replacement) pairs for abbreviations:
|
||||
_abbreviations = {
|
||||
"en": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("mrs", "misess"),
|
||||
("mr", "mister"),
|
||||
("dr", "doctor"),
|
||||
("st", "saint"),
|
||||
("co", "company"),
|
||||
("jr", "junior"),
|
||||
("maj", "major"),
|
||||
("gen", "general"),
|
||||
("drs", "doctors"),
|
||||
("rev", "reverend"),
|
||||
("lt", "lieutenant"),
|
||||
("hon", "honorable"),
|
||||
("sgt", "sergeant"),
|
||||
("capt", "captain"),
|
||||
("esq", "esquire"),
|
||||
("ltd", "limited"),
|
||||
("col", "colonel"),
|
||||
("ft", "fort"),
|
||||
]
|
||||
],
|
||||
"es": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("sra", "señora"),
|
||||
("sr", "señor"),
|
||||
("dr", "doctor"),
|
||||
("dra", "doctora"),
|
||||
("st", "santo"),
|
||||
("co", "compañía"),
|
||||
("jr", "junior"),
|
||||
("ltd", "limitada"),
|
||||
]
|
||||
],
|
||||
"fr": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("mme", "madame"),
|
||||
("mr", "monsieur"),
|
||||
("dr", "docteur"),
|
||||
("st", "saint"),
|
||||
("co", "compagnie"),
|
||||
("jr", "junior"),
|
||||
("ltd", "limitée"),
|
||||
]
|
||||
],
|
||||
"de": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("fr", "frau"),
|
||||
("dr", "doktor"),
|
||||
("st", "sankt"),
|
||||
("co", "firma"),
|
||||
("jr", "junior"),
|
||||
]
|
||||
],
|
||||
"pt": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("sra", "senhora"),
|
||||
("sr", "senhor"),
|
||||
("dr", "doutor"),
|
||||
("dra", "doutora"),
|
||||
("st", "santo"),
|
||||
("co", "companhia"),
|
||||
("jr", "júnior"),
|
||||
("ltd", "limitada"),
|
||||
]
|
||||
],
|
||||
"it": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
# ("sig.ra", "signora"),
|
||||
("sig", "signore"),
|
||||
("dr", "dottore"),
|
||||
("st", "santo"),
|
||||
("co", "compagnia"),
|
||||
("jr", "junior"),
|
||||
("ltd", "limitata"),
|
||||
]
|
||||
],
|
||||
"pl": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("p", "pani"),
|
||||
("m", "pan"),
|
||||
("dr", "doktor"),
|
||||
("sw", "święty"),
|
||||
("jr", "junior"),
|
||||
]
|
||||
],
|
||||
"ar": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
# There are not many common abbreviations in Arabic as in English.
|
||||
]
|
||||
],
|
||||
"zh": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
# Chinese doesn't typically use abbreviations in the same way as Latin-based scripts.
|
||||
]
|
||||
],
|
||||
"cs": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("dr", "doktor"), # doctor
|
||||
("ing", "inženýr"), # engineer
|
||||
("p", "pan"), # Could also map to pani for woman but no easy way to do it
|
||||
# Other abbreviations would be specialized and not as common.
|
||||
]
|
||||
],
|
||||
"ru": [
|
||||
(re.compile("\\b%s\\b" % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("г-жа", "госпожа"), # Mrs.
|
||||
("г-н", "господин"), # Mr.
|
||||
("д-р", "доктор"), # doctor
|
||||
# Other abbreviations are less common or specialized.
|
||||
]
|
||||
],
|
||||
"nl": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("dhr", "de heer"), # Mr.
|
||||
("mevr", "mevrouw"), # Mrs.
|
||||
("dr", "dokter"), # doctor
|
||||
("jhr", "jonkheer"), # young lord or nobleman
|
||||
# Dutch uses more abbreviations, but these are the most common ones.
|
||||
]
|
||||
],
|
||||
"tr": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("b", "bay"), # Mr.
|
||||
("byk", "büyük"), # büyük
|
||||
("dr", "doktor"), # doctor
|
||||
# Add other Turkish abbreviations here if needed.
|
||||
]
|
||||
],
|
||||
"hu": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("dr", "doktor"), # doctor
|
||||
("b", "bácsi"), # Mr.
|
||||
("nőv", "nővér"), # nurse
|
||||
# Add other Hungarian abbreviations here if needed.
|
||||
]
|
||||
],
|
||||
"ko": [
|
||||
(re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
# Korean doesn't typically use abbreviations in the same way as Latin-based scripts.
|
||||
]
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def expand_abbreviations_multilingual(text, lang="en"):
|
||||
for regex, replacement in _abbreviations[lang]:
|
||||
text = re.sub(regex, replacement, text)
|
||||
return text
|
||||
|
||||
|
||||
_symbols_multilingual = {
|
||||
"en": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " and "),
|
||||
("@", " at "),
|
||||
("%", " percent "),
|
||||
("#", " hash "),
|
||||
("$", " dollar "),
|
||||
("£", " pound "),
|
||||
("°", " degree "),
|
||||
]
|
||||
],
|
||||
"es": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " y "),
|
||||
("@", " arroba "),
|
||||
("%", " por ciento "),
|
||||
("#", " numeral "),
|
||||
("$", " dolar "),
|
||||
("£", " libra "),
|
||||
("°", " grados "),
|
||||
]
|
||||
],
|
||||
"fr": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " et "),
|
||||
("@", " arobase "),
|
||||
("%", " pour cent "),
|
||||
("#", " dièse "),
|
||||
("$", " dollar "),
|
||||
("£", " livre "),
|
||||
("°", " degrés "),
|
||||
]
|
||||
],
|
||||
"de": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " und "),
|
||||
("@", " at "),
|
||||
("%", " prozent "),
|
||||
("#", " raute "),
|
||||
("$", " dollar "),
|
||||
("£", " pfund "),
|
||||
("°", " grad "),
|
||||
]
|
||||
],
|
||||
"pt": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " e "),
|
||||
("@", " arroba "),
|
||||
("%", " por cento "),
|
||||
("#", " cardinal "),
|
||||
("$", " dólar "),
|
||||
("£", " libra "),
|
||||
("°", " graus "),
|
||||
]
|
||||
],
|
||||
"it": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " e "),
|
||||
("@", " chiocciola "),
|
||||
("%", " per cento "),
|
||||
("#", " cancelletto "),
|
||||
("$", " dollaro "),
|
||||
("£", " sterlina "),
|
||||
("°", " gradi "),
|
||||
]
|
||||
],
|
||||
"pl": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " i "),
|
||||
("@", " małpa "),
|
||||
("%", " procent "),
|
||||
("#", " krzyżyk "),
|
||||
("$", " dolar "),
|
||||
("£", " funt "),
|
||||
("°", " stopnie "),
|
||||
]
|
||||
],
|
||||
"ar": [
|
||||
# Arabic
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " و "),
|
||||
("@", " على "),
|
||||
("%", " في المئة "),
|
||||
("#", " رقم "),
|
||||
("$", " دولار "),
|
||||
("£", " جنيه "),
|
||||
("°", " درجة "),
|
||||
]
|
||||
],
|
||||
"zh": [
|
||||
# Chinese
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " 和 "),
|
||||
("@", " 在 "),
|
||||
("%", " 百分之 "),
|
||||
("#", " 号 "),
|
||||
("$", " 美元 "),
|
||||
("£", " 英镑 "),
|
||||
("°", " 度 "),
|
||||
]
|
||||
],
|
||||
"cs": [
|
||||
# Czech
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " a "),
|
||||
("@", " na "),
|
||||
("%", " procento "),
|
||||
("#", " křížek "),
|
||||
("$", " dolar "),
|
||||
("£", " libra "),
|
||||
("°", " stupně "),
|
||||
]
|
||||
],
|
||||
"ru": [
|
||||
# Russian
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " и "),
|
||||
("@", " собака "),
|
||||
("%", " процентов "),
|
||||
("#", " номер "),
|
||||
("$", " доллар "),
|
||||
("£", " фунт "),
|
||||
("°", " градус "),
|
||||
]
|
||||
],
|
||||
"nl": [
|
||||
# Dutch
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " en "),
|
||||
("@", " bij "),
|
||||
("%", " procent "),
|
||||
("#", " hekje "),
|
||||
("$", " dollar "),
|
||||
("£", " pond "),
|
||||
("°", " graden "),
|
||||
]
|
||||
],
|
||||
"tr": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " ve "),
|
||||
("@", " at "),
|
||||
("%", " yüzde "),
|
||||
("#", " diyez "),
|
||||
("$", " dolar "),
|
||||
("£", " sterlin "),
|
||||
("°", " derece "),
|
||||
]
|
||||
],
|
||||
"hu": [
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " és "),
|
||||
("@", " kukac "),
|
||||
("%", " százalék "),
|
||||
("#", " kettőskereszt "),
|
||||
("$", " dollár "),
|
||||
("£", " font "),
|
||||
("°", " fok "),
|
||||
]
|
||||
],
|
||||
"ko": [
|
||||
# Korean
|
||||
(re.compile(r"%s" % re.escape(x[0]), re.IGNORECASE), x[1])
|
||||
for x in [
|
||||
("&", " 그리고 "),
|
||||
("@", " 에 "),
|
||||
("%", " 퍼센트 "),
|
||||
("#", " 번호 "),
|
||||
("$", " 달러 "),
|
||||
("£", " 파운드 "),
|
||||
("°", " 도 "),
|
||||
]
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def expand_symbols_multilingual(text, lang="en"):
|
||||
for regex, replacement in _symbols_multilingual[lang]:
|
||||
text = re.sub(regex, replacement, text)
|
||||
text = text.replace(" ", " ") # Ensure there are no double spaces
|
||||
return text.strip()
|
||||
|
||||
|
||||
_ordinal_re = {
|
||||
"en": re.compile(r"([0-9]+)(st|nd|rd|th)"),
|
||||
"es": re.compile(r"([0-9]+)(º|ª|er|o|a|os|as)"),
|
||||
"fr": re.compile(r"([0-9]+)(º|ª|er|re|e|ème)"),
|
||||
"de": re.compile(r"([0-9]+)(st|nd|rd|th|º|ª|\.(?=\s|$))"),
|
||||
"pt": re.compile(r"([0-9]+)(º|ª|o|a|os|as)"),
|
||||
"it": re.compile(r"([0-9]+)(º|°|ª|o|a|i|e)"),
|
||||
"pl": re.compile(r"([0-9]+)(º|ª|st|nd|rd|th)"),
|
||||
"ar": re.compile(r"([0-9]+)(ون|ين|ث|ر|ى)"),
|
||||
"cs": re.compile(r"([0-9]+)\.(?=\s|$)"), # In Czech, a dot is often used after the number to indicate ordinals.
|
||||
"ru": re.compile(r"([0-9]+)(-й|-я|-е|-ое|-ье|-го)"),
|
||||
"nl": re.compile(r"([0-9]+)(de|ste|e)"),
|
||||
"tr": re.compile(r"([0-9]+)(\.|inci|nci|uncu|üncü|\.)"),
|
||||
"hu": re.compile(r"([0-9]+)(\.|adik|edik|odik|edik|ödik|ödike|ik)"),
|
||||
"ko": re.compile(r"([0-9]+)(번째|번|차|째)"),
|
||||
}
|
||||
_number_re = re.compile(r"[0-9]+")
|
||||
_currency_re = {
|
||||
"USD": re.compile(r"((\$[0-9\.\,]*[0-9]+)|([0-9\.\,]*[0-9]+\$))"),
|
||||
"GBP": re.compile(r"((£[0-9\.\,]*[0-9]+)|([0-9\.\,]*[0-9]+£))"),
|
||||
"EUR": re.compile(r"(([0-9\.\,]*[0-9]+€)|((€[0-9\.\,]*[0-9]+)))"),
|
||||
}
|
||||
|
||||
_comma_number_re = re.compile(r"\b\d{1,3}(,\d{3})*(\.\d+)?\b")
|
||||
_dot_number_re = re.compile(r"\b\d{1,3}(.\d{3})*(\,\d+)?\b")
|
||||
_decimal_number_re = re.compile(r"([0-9]+[.,][0-9]+)")
|
||||
|
||||
|
||||
def _remove_commas(m):
|
||||
text = m.group(0)
|
||||
if "," in text:
|
||||
text = text.replace(",", "")
|
||||
return text
|
||||
|
||||
|
||||
def _remove_dots(m):
|
||||
text = m.group(0)
|
||||
if "." in text:
|
||||
text = text.replace(".", "")
|
||||
return text
|
||||
|
||||
|
||||
def _expand_decimal_point(m, lang="en"):
|
||||
amount = m.group(1).replace(",", ".")
|
||||
return num2words(float(amount), lang=lang if lang != "cs" else "cz")
|
||||
|
||||
|
||||
def _expand_currency(m, lang="en", currency="USD"):
|
||||
amount = float((re.sub(r"[^\d.]", "", m.group(0).replace(",", "."))))
|
||||
full_amount = num2words(amount, to="currency", currency=currency, lang=lang if lang != "cs" else "cz")
|
||||
|
||||
and_equivalents = {
|
||||
"en": ", ",
|
||||
"es": " con ",
|
||||
"fr": " et ",
|
||||
"de": " und ",
|
||||
"pt": " e ",
|
||||
"it": " e ",
|
||||
"pl": ", ",
|
||||
"cs": ", ",
|
||||
"ru": ", ",
|
||||
"nl": ", ",
|
||||
"ar": ", ",
|
||||
"tr": ", ",
|
||||
"hu": ", ",
|
||||
"ko": ", ",
|
||||
}
|
||||
|
||||
if amount.is_integer():
|
||||
last_and = full_amount.rfind(and_equivalents[lang])
|
||||
if last_and != -1:
|
||||
full_amount = full_amount[:last_and]
|
||||
|
||||
return full_amount
|
||||
|
||||
|
||||
def _expand_ordinal(m, lang="en"):
|
||||
return num2words(int(m.group(1)), ordinal=True, lang=lang if lang != "cs" else "cz")
|
||||
|
||||
|
||||
def _expand_number(m, lang="en"):
|
||||
return num2words(int(m.group(0)), lang=lang if lang != "cs" else "cz")
|
||||
|
||||
|
||||
def expand_numbers_multilingual(text, lang="en"):
|
||||
if lang == "zh":
|
||||
text = zh_num2words()(text)
|
||||
else:
|
||||
if lang in ["en", "ru"]:
|
||||
text = re.sub(_comma_number_re, _remove_commas, text)
|
||||
else:
|
||||
text = re.sub(_dot_number_re, _remove_dots, text)
|
||||
try:
|
||||
text = re.sub(_currency_re["GBP"], lambda m: _expand_currency(m, lang, "GBP"), text)
|
||||
text = re.sub(_currency_re["USD"], lambda m: _expand_currency(m, lang, "USD"), text)
|
||||
text = re.sub(_currency_re["EUR"], lambda m: _expand_currency(m, lang, "EUR"), text)
|
||||
except:
|
||||
pass
|
||||
if lang != "tr":
|
||||
text = re.sub(_decimal_number_re, lambda m: _expand_decimal_point(m, lang), text)
|
||||
text = re.sub(_ordinal_re[lang], lambda m: _expand_ordinal(m, lang), text)
|
||||
text = re.sub(_number_re, lambda m: _expand_number(m, lang), text)
|
||||
return text
|
||||
|
||||
|
||||
def lowercase(text):
|
||||
return text.lower()
|
||||
|
||||
|
||||
def collapse_whitespace(text):
|
||||
return re.sub(_whitespace_re, " ", text)
|
||||
|
||||
|
||||
def multilingual_cleaners(text, lang):
|
||||
text = text.replace('"', "")
|
||||
if lang == "tr":
|
||||
text = text.replace("İ", "i")
|
||||
text = text.replace("Ö", "ö")
|
||||
text = text.replace("Ü", "ü")
|
||||
text = lowercase(text)
|
||||
try:
|
||||
text = expand_numbers_multilingual(text, lang)
|
||||
except:
|
||||
pass
|
||||
try:
|
||||
text = expand_abbreviations_multilingual(text, lang)
|
||||
except:
|
||||
pass
|
||||
try:
|
||||
text = expand_symbols_multilingual(text, lang=lang)
|
||||
except:
|
||||
pass
|
||||
text = collapse_whitespace(text)
|
||||
return text
|
||||
|
||||
|
||||
def basic_cleaners(text):
|
||||
"""Basic pipeline that lowercases and collapses whitespace without transliteration."""
|
||||
text = lowercase(text)
|
||||
text = collapse_whitespace(text)
|
||||
return text
|
||||
|
||||
|
||||
def chinese_transliterate(text):
|
||||
return "".join(
|
||||
[p[0] for p in pypinyin.pinyin(text, style=pypinyin.Style.TONE3, heteronym=False, neutral_tone_with_five=True)]
|
||||
)
|
||||
|
||||
|
||||
def japanese_cleaners(text, katsu):
|
||||
text = katsu.romaji(text)
|
||||
text = lowercase(text)
|
||||
return text
|
||||
|
||||
|
||||
def korean_transliterate(text):
|
||||
r = Transliter(academic)
|
||||
return r.translit(text)
|
||||
|
||||
|
||||
DEFAULT_VOCAB_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vocab.json")
|
||||
|
||||
|
||||
class VoiceBpeTokenizer:
|
||||
def __init__(self, vocab_file=DEFAULT_VOCAB_FILE):
|
||||
self.tokenizer = None
|
||||
if vocab_file is not None:
|
||||
self.tokenizer = Tokenizer.from_file(vocab_file)
|
||||
self.char_limits = {
|
||||
"en": 10000,
|
||||
"de": 253,
|
||||
"fr": 273,
|
||||
"es": 239,
|
||||
"it": 213,
|
||||
"pt": 203,
|
||||
"pl": 224,
|
||||
"zh": 82,
|
||||
"ar": 166,
|
||||
"cs": 186,
|
||||
"ru": 182,
|
||||
"nl": 251,
|
||||
"tr": 226,
|
||||
"ja": 71,
|
||||
"hu": 224,
|
||||
"ko": 95,
|
||||
}
|
||||
|
||||
@cached_property
|
||||
def katsu(self):
|
||||
import cutlet
|
||||
|
||||
return cutlet.Cutlet()
|
||||
|
||||
def check_input_length(self, txt, lang):
|
||||
lang = lang.split("-")[0] # remove the region
|
||||
limit = self.char_limits.get(lang, 250)
|
||||
# if len(txt) > limit:
|
||||
# print(
|
||||
# f"[!] Warning: The text length exceeds the character limit of {limit} for language '{lang}', this might cause truncated audio."
|
||||
# )
|
||||
|
||||
def preprocess_text(self, txt, lang):
|
||||
if lang in {"ar", "cs", "de", "en", "es", "fr", "hu", "it", "nl", "pl", "pt", "ru", "tr", "zh", "ko"}:
|
||||
txt = multilingual_cleaners(txt, lang)
|
||||
if lang == "zh":
|
||||
txt = chinese_transliterate(txt)
|
||||
if lang == "ko":
|
||||
txt = korean_transliterate(txt)
|
||||
elif lang == "ja":
|
||||
txt = japanese_cleaners(txt, self.katsu)
|
||||
elif lang == "hi":
|
||||
# @manmay will implement this
|
||||
txt = basic_cleaners(txt)
|
||||
else:
|
||||
raise NotImplementedError(f"Language '{lang}' is not supported.")
|
||||
return txt
|
||||
|
||||
def encode(self, txt, lang):
|
||||
lang = lang.split("-")[0] # remove the region
|
||||
self.check_input_length(txt, lang)
|
||||
txt = self.preprocess_text(txt, lang)
|
||||
lang = "zh-cn" if lang == "zh" else lang
|
||||
txt = f"[{lang}]{txt}"
|
||||
txt = txt.replace(" ", "[SPACE]")
|
||||
return self.tokenizer.encode(txt).ids
|
||||
|
||||
def decode(self, seq, skip_special_tokens=False):
|
||||
if isinstance(seq, torch.Tensor):
|
||||
seq = seq.cpu().numpy()
|
||||
txt = self.tokenizer.decode(seq, skip_special_tokens=False).replace(" ", "")
|
||||
txt = txt.replace("[SPACE]", " ")
|
||||
txt = txt.replace("[STOP]", "")
|
||||
# txt = txt.replace("[UNK]", "")
|
||||
return txt
|
||||
|
||||
|
||||
#copy from https://github.com/huggingface/transformers/blob/main/src/transformers/tokenization_utils_base.py#L3936
|
||||
def batch_decode(
|
||||
self,
|
||||
sequences: Union[List[int], List[List[int]], "np.ndarray", "torch.Tensor", "tf.Tensor"],
|
||||
skip_special_tokens: bool = False,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Convert a list of lists of token ids into a list of strings by calling decode.
|
||||
|
||||
Args:
|
||||
sequences (`Union[List[int], List[List[int]], np.ndarray, torch.Tensor, tf.Tensor]`):
|
||||
List of tokenized input ids. Can be obtained using the `__call__` method.
|
||||
skip_special_tokens (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not to remove special tokens in the decoding.
|
||||
kwargs (additional keyword arguments, *optional*):
|
||||
Will be passed to the underlying model specific decode method.
|
||||
|
||||
Returns:
|
||||
`List[str]`: The list of decoded sentences.
|
||||
"""
|
||||
return [
|
||||
self.decode(seq)
|
||||
for seq in sequences
|
||||
]
|
||||
|
||||
#https://github.com/coqui-ai/TTS/blob/dev/TTS/tts/layers/xtts/trainer/dataset.py#L202
|
||||
# def pad(self):
|
||||
|
||||
def __len__(self):
|
||||
return self.tokenizer.get_vocab_size()
|
||||
|
||||
def get_number_tokens(self):
|
||||
return max(self.tokenizer.get_vocab().values()) + 1
|
||||
|
||||
|
||||
def test_expand_numbers_multilingual():
|
||||
test_cases = [
|
||||
# English
|
||||
("In 12.5 seconds.", "In twelve point five seconds.", "en"),
|
||||
("There were 50 soldiers.", "There were fifty soldiers.", "en"),
|
||||
("This is a 1st test", "This is a first test", "en"),
|
||||
("That will be $20 sir.", "That will be twenty dollars sir.", "en"),
|
||||
("That will be 20€ sir.", "That will be twenty euro sir.", "en"),
|
||||
("That will be 20.15€ sir.", "That will be twenty euro, fifteen cents sir.", "en"),
|
||||
("That's 100,000.5.", "That's one hundred thousand point five.", "en"),
|
||||
# French
|
||||
("En 12,5 secondes.", "En douze virgule cinq secondes.", "fr"),
|
||||
("Il y avait 50 soldats.", "Il y avait cinquante soldats.", "fr"),
|
||||
("Ceci est un 1er test", "Ceci est un premier test", "fr"),
|
||||
("Cela vous fera $20 monsieur.", "Cela vous fera vingt dollars monsieur.", "fr"),
|
||||
("Cela vous fera 20€ monsieur.", "Cela vous fera vingt euros monsieur.", "fr"),
|
||||
("Cela vous fera 20,15€ monsieur.", "Cela vous fera vingt euros et quinze centimes monsieur.", "fr"),
|
||||
("Ce sera 100.000,5.", "Ce sera cent mille virgule cinq.", "fr"),
|
||||
# German
|
||||
("In 12,5 Sekunden.", "In zwölf Komma fünf Sekunden.", "de"),
|
||||
("Es gab 50 Soldaten.", "Es gab fünfzig Soldaten.", "de"),
|
||||
("Dies ist ein 1. Test", "Dies ist ein erste Test", "de"), # Issue with gender
|
||||
("Das macht $20 Herr.", "Das macht zwanzig Dollar Herr.", "de"),
|
||||
("Das macht 20€ Herr.", "Das macht zwanzig Euro Herr.", "de"),
|
||||
("Das macht 20,15€ Herr.", "Das macht zwanzig Euro und fünfzehn Cent Herr.", "de"),
|
||||
# Spanish
|
||||
("En 12,5 segundos.", "En doce punto cinco segundos.", "es"),
|
||||
("Había 50 soldados.", "Había cincuenta soldados.", "es"),
|
||||
("Este es un 1er test", "Este es un primero test", "es"),
|
||||
("Eso le costará $20 señor.", "Eso le costará veinte dólares señor.", "es"),
|
||||
("Eso le costará 20€ señor.", "Eso le costará veinte euros señor.", "es"),
|
||||
("Eso le costará 20,15€ señor.", "Eso le costará veinte euros con quince céntimos señor.", "es"),
|
||||
# Italian
|
||||
("In 12,5 secondi.", "In dodici virgola cinque secondi.", "it"),
|
||||
("C'erano 50 soldati.", "C'erano cinquanta soldati.", "it"),
|
||||
("Questo è un 1° test", "Questo è un primo test", "it"),
|
||||
("Ti costerà $20 signore.", "Ti costerà venti dollari signore.", "it"),
|
||||
("Ti costerà 20€ signore.", "Ti costerà venti euro signore.", "it"),
|
||||
("Ti costerà 20,15€ signore.", "Ti costerà venti euro e quindici centesimi signore.", "it"),
|
||||
# Portuguese
|
||||
("Em 12,5 segundos.", "Em doze vírgula cinco segundos.", "pt"),
|
||||
("Havia 50 soldados.", "Havia cinquenta soldados.", "pt"),
|
||||
("Este é um 1º teste", "Este é um primeiro teste", "pt"),
|
||||
("Isso custará $20 senhor.", "Isso custará vinte dólares senhor.", "pt"),
|
||||
("Isso custará 20€ senhor.", "Isso custará vinte euros senhor.", "pt"),
|
||||
(
|
||||
"Isso custará 20,15€ senhor.",
|
||||
"Isso custará vinte euros e quinze cêntimos senhor.",
|
||||
"pt",
|
||||
), # "cêntimos" should be "centavos" num2words issue
|
||||
# Polish
|
||||
("W 12,5 sekundy.", "W dwanaście przecinek pięć sekundy.", "pl"),
|
||||
("Było 50 żołnierzy.", "Było pięćdziesiąt żołnierzy.", "pl"),
|
||||
("To będzie kosztować 20€ panie.", "To będzie kosztować dwadzieścia euro panie.", "pl"),
|
||||
("To będzie kosztować 20,15€ panie.", "To będzie kosztować dwadzieścia euro, piętnaście centów panie.", "pl"),
|
||||
# Arabic
|
||||
("في الـ 12,5 ثانية.", "في الـ اثنا عشر , خمسون ثانية.", "ar"),
|
||||
("كان هناك 50 جنديًا.", "كان هناك خمسون جنديًا.", "ar"),
|
||||
# ("ستكون النتيجة $20 يا سيد.", 'ستكون النتيجة عشرون دولار يا سيد.', 'ar'), # $ and € are mising from num2words
|
||||
# ("ستكون النتيجة 20€ يا سيد.", 'ستكون النتيجة عشرون يورو يا سيد.', 'ar'),
|
||||
# Czech
|
||||
("Za 12,5 vteřiny.", "Za dvanáct celá pět vteřiny.", "cs"),
|
||||
("Bylo tam 50 vojáků.", "Bylo tam padesát vojáků.", "cs"),
|
||||
("To bude stát 20€ pane.", "To bude stát dvacet euro pane.", "cs"),
|
||||
("To bude 20.15€ pane.", "To bude dvacet euro, patnáct centů pane.", "cs"),
|
||||
# Russian
|
||||
("Через 12.5 секунды.", "Через двенадцать запятая пять секунды.", "ru"),
|
||||
("Там было 50 солдат.", "Там было пятьдесят солдат.", "ru"),
|
||||
("Это будет 20.15€ сэр.", "Это будет двадцать евро, пятнадцать центов сэр.", "ru"),
|
||||
("Это будет стоить 20€ господин.", "Это будет стоить двадцать евро господин.", "ru"),
|
||||
# Dutch
|
||||
("In 12,5 seconden.", "In twaalf komma vijf seconden.", "nl"),
|
||||
("Er waren 50 soldaten.", "Er waren vijftig soldaten.", "nl"),
|
||||
("Dat wordt dan $20 meneer.", "Dat wordt dan twintig dollar meneer.", "nl"),
|
||||
("Dat wordt dan 20€ meneer.", "Dat wordt dan twintig euro meneer.", "nl"),
|
||||
# Chinese (Simplified)
|
||||
("在12.5秒内", "在十二点五秒内", "zh"),
|
||||
("有50名士兵", "有五十名士兵", "zh"),
|
||||
# ("那将是$20先生", '那将是二十美元先生', 'zh'), currency doesn't work
|
||||
# ("那将是20€先生", '那将是二十欧元先生', 'zh'),
|
||||
# Turkish
|
||||
# ("12,5 saniye içinde.", 'On iki virgül beş saniye içinde.', 'tr'), # decimal doesn't work for TR
|
||||
("50 asker vardı.", "elli asker vardı.", "tr"),
|
||||
("Bu 1. test", "Bu birinci test", "tr"),
|
||||
# ("Bu 100.000,5.", 'Bu yüz bin virgül beş.', 'tr'),
|
||||
# Hungarian
|
||||
("12,5 másodperc alatt.", "tizenkettő egész öt tized másodperc alatt.", "hu"),
|
||||
("50 katona volt.", "ötven katona volt.", "hu"),
|
||||
("Ez az 1. teszt", "Ez az első teszt", "hu"),
|
||||
# Korean
|
||||
("12.5 초 안에.", "십이 점 다섯 초 안에.", "ko"),
|
||||
("50 명의 병사가 있었다.", "오십 명의 병사가 있었다.", "ko"),
|
||||
("이것은 1 번째 테스트입니다", "이것은 첫 번째 테스트입니다", "ko"),
|
||||
]
|
||||
for a, b, lang in test_cases:
|
||||
out = expand_numbers_multilingual(a, lang=lang)
|
||||
assert out == b, f"'{out}' vs '{b}'"
|
||||
|
||||
|
||||
def test_abbreviations_multilingual():
|
||||
test_cases = [
|
||||
# English
|
||||
("Hello Mr. Smith.", "Hello mister Smith.", "en"),
|
||||
("Dr. Jones is here.", "doctor Jones is here.", "en"),
|
||||
# Spanish
|
||||
("Hola Sr. Garcia.", "Hola señor Garcia.", "es"),
|
||||
("La Dra. Martinez es muy buena.", "La doctora Martinez es muy buena.", "es"),
|
||||
# French
|
||||
("Bonjour Mr. Dupond.", "Bonjour monsieur Dupond.", "fr"),
|
||||
("Mme. Moreau est absente aujourd'hui.", "madame Moreau est absente aujourd'hui.", "fr"),
|
||||
# German
|
||||
("Frau Dr. Müller ist sehr klug.", "Frau doktor Müller ist sehr klug.", "de"),
|
||||
# Portuguese
|
||||
("Olá Sr. Silva.", "Olá senhor Silva.", "pt"),
|
||||
("Dra. Costa, você está disponível?", "doutora Costa, você está disponível?", "pt"),
|
||||
# Italian
|
||||
("Buongiorno, Sig. Rossi.", "Buongiorno, signore Rossi.", "it"),
|
||||
# ("Sig.ra Bianchi, posso aiutarti?", 'signora Bianchi, posso aiutarti?', 'it'), # Issue with matching that pattern
|
||||
# Polish
|
||||
("Dzień dobry, P. Kowalski.", "Dzień dobry, pani Kowalski.", "pl"),
|
||||
("M. Nowak, czy mogę zadać pytanie?", "pan Nowak, czy mogę zadać pytanie?", "pl"),
|
||||
# Czech
|
||||
("P. Novák", "pan Novák", "cs"),
|
||||
("Dr. Vojtěch", "doktor Vojtěch", "cs"),
|
||||
# Dutch
|
||||
("Dhr. Jansen", "de heer Jansen", "nl"),
|
||||
("Mevr. de Vries", "mevrouw de Vries", "nl"),
|
||||
# Russian
|
||||
("Здравствуйте Г-н Иванов.", "Здравствуйте господин Иванов.", "ru"),
|
||||
("Д-р Смирнов здесь, чтобы увидеть вас.", "доктор Смирнов здесь, чтобы увидеть вас.", "ru"),
|
||||
# Turkish
|
||||
("Merhaba B. Yılmaz.", "Merhaba bay Yılmaz.", "tr"),
|
||||
("Dr. Ayşe burada.", "doktor Ayşe burada.", "tr"),
|
||||
# Hungarian
|
||||
("Dr. Szabó itt van.", "doktor Szabó itt van.", "hu"),
|
||||
]
|
||||
|
||||
for a, b, lang in test_cases:
|
||||
out = expand_abbreviations_multilingual(a, lang=lang)
|
||||
assert out == b, f"'{out}' vs '{b}'"
|
||||
|
||||
|
||||
def test_symbols_multilingual():
|
||||
test_cases = [
|
||||
("I have 14% battery", "I have 14 percent battery", "en"),
|
||||
("Te veo @ la fiesta", "Te veo arroba la fiesta", "es"),
|
||||
("J'ai 14° de fièvre", "J'ai 14 degrés de fièvre", "fr"),
|
||||
("Die Rechnung beträgt £ 20", "Die Rechnung beträgt pfund 20", "de"),
|
||||
("O meu email é ana&joao@gmail.com", "O meu email é ana e joao arroba gmail.com", "pt"),
|
||||
("linguaggio di programmazione C#", "linguaggio di programmazione C cancelletto", "it"),
|
||||
("Moja temperatura to 36.6°", "Moja temperatura to 36.6 stopnie", "pl"),
|
||||
("Mám 14% baterie", "Mám 14 procento baterie", "cs"),
|
||||
("Těším se na tebe @ party", "Těším se na tebe na party", "cs"),
|
||||
("У меня 14% заряда", "У меня 14 процентов заряда", "ru"),
|
||||
("Я буду @ дома", "Я буду собака дома", "ru"),
|
||||
("Ik heb 14% batterij", "Ik heb 14 procent batterij", "nl"),
|
||||
("Ik zie je @ het feest", "Ik zie je bij het feest", "nl"),
|
||||
("لدي 14% في البطارية", "لدي 14 في المئة في البطارية", "ar"),
|
||||
("我的电量为 14%", "我的电量为 14 百分之", "zh"),
|
||||
("Pilim %14 dolu.", "Pilim yüzde 14 dolu.", "tr"),
|
||||
("Az akkumulátorom töltöttsége 14%", "Az akkumulátorom töltöttsége 14 százalék", "hu"),
|
||||
("배터리 잔량이 14%입니다.", "배터리 잔량이 14 퍼센트입니다.", "ko"),
|
||||
]
|
||||
|
||||
for a, b, lang in test_cases:
|
||||
out = expand_symbols_multilingual(a, lang=lang)
|
||||
assert out == b, f"'{out}' vs '{b}'"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_expand_numbers_multilingual()
|
||||
test_abbreviations_multilingual()
|
||||
test_symbols_multilingual()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,141 @@
|
||||
import os
|
||||
import torch
|
||||
from diffusers import AutoencoderDC
|
||||
import torchaudio
|
||||
import torchvision.transforms as transforms
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.loaders import FromOriginalModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
|
||||
|
||||
try:
|
||||
from .music_vocoder import ADaMoSHiFiGANV1
|
||||
except ImportError:
|
||||
from music_vocoder import ADaMoSHiFiGANV1
|
||||
|
||||
|
||||
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
DEFAULT_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_dcae_f8c8")
|
||||
VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder")
|
||||
|
||||
|
||||
class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
||||
@register_to_config
|
||||
def __init__(self, source_sample_rate=None, dcae_checkpoint_path=DEFAULT_PRETRAINED_PATH, vocoder_checkpoint_path=VOCODER_PRETRAINED_PATH):
|
||||
super(MusicDCAE, self).__init__()
|
||||
|
||||
self.dcae = AutoencoderDC.from_pretrained(dcae_checkpoint_path)
|
||||
self.vocoder = ADaMoSHiFiGANV1.from_pretrained(vocoder_checkpoint_path)
|
||||
|
||||
if source_sample_rate is None:
|
||||
source_sample_rate = 48000
|
||||
|
||||
self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
|
||||
|
||||
self.transform = transforms.Compose([
|
||||
transforms.Normalize(0.5, 0.5),
|
||||
])
|
||||
self.min_mel_value = -11.0
|
||||
self.max_mel_value = 3.0
|
||||
self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
|
||||
self.mel_chunk_size = 1024
|
||||
self.time_dimention_multiple = 8
|
||||
self.latent_chunk_size = self.mel_chunk_size // self.time_dimention_multiple
|
||||
self.scale_factor = 0.1786
|
||||
self.shift_factor = -1.9091
|
||||
|
||||
def load_audio(self, audio_path):
|
||||
audio, sr = torchaudio.load(audio_path)
|
||||
return audio, sr
|
||||
|
||||
def forward_mel(self, audios):
|
||||
mels = []
|
||||
for i in range(len(audios)):
|
||||
image = self.vocoder.mel_transform(audios[i])
|
||||
mels.append(image)
|
||||
mels = torch.stack(mels)
|
||||
return mels
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, audios, audio_lengths=None, sr=None):
|
||||
if audio_lengths is None:
|
||||
audio_lengths = torch.tensor([audios.shape[2]] * audios.shape[0])
|
||||
audio_lengths = audio_lengths.to(audios.device)
|
||||
|
||||
# audios: N x 2 x T, 48kHz
|
||||
device = audios.device
|
||||
dtype = audios.dtype
|
||||
|
||||
if sr is None:
|
||||
sr = 48000
|
||||
resampler = self.resampler
|
||||
else:
|
||||
resampler = torchaudio.transforms.Resample(sr, 44100).to(device).to(dtype)
|
||||
|
||||
audio = resampler(audios)
|
||||
|
||||
max_audio_len = audio.shape[-1]
|
||||
if max_audio_len % (8 * 512) != 0:
|
||||
audio = torch.nn.functional.pad(audio, (0, 8 * 512 - max_audio_len % (8 * 512)))
|
||||
|
||||
mels = self.forward_mel(audio)
|
||||
mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
|
||||
mels = self.transform(mels)
|
||||
latents = []
|
||||
for mel in mels:
|
||||
latent = self.dcae.encoder(mel.unsqueeze(0))
|
||||
latents.append(latent)
|
||||
latents = torch.cat(latents, dim=0)
|
||||
latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
|
||||
latents = (latents - self.shift_factor) * self.scale_factor
|
||||
return latents, latent_lengths
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, latents, audio_lengths=None, sr=None):
|
||||
latents = latents / self.scale_factor + self.shift_factor
|
||||
|
||||
pred_wavs = []
|
||||
|
||||
for latent in latents:
|
||||
mels = self.dcae.decoder(latent.unsqueeze(0))
|
||||
mels = mels * 0.5 + 0.5
|
||||
mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
|
||||
wav = self.vocoder.decode(mels[0]).squeeze(1)
|
||||
|
||||
if sr is not None:
|
||||
resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
|
||||
wav = resampler(wav)
|
||||
else:
|
||||
sr = 44100
|
||||
pred_wavs.append(wav)
|
||||
|
||||
if audio_lengths is not None:
|
||||
pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
|
||||
return sr, pred_wavs
|
||||
|
||||
def forward(self, audios, audio_lengths=None, sr=None):
|
||||
latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
|
||||
sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
|
||||
return sr, pred_wavs, latents, latent_lengths
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
audio, sr = torchaudio.load("test.wav")
|
||||
audio_lengths = torch.tensor([audio.shape[1]])
|
||||
audios = audio.unsqueeze(0)
|
||||
|
||||
# test encode only
|
||||
model = MusicDCAE()
|
||||
# latents, latent_lengths = model.encode(audios, audio_lengths)
|
||||
# print("latents shape: ", latents.shape)
|
||||
# print("latent_lengths: ", latent_lengths)
|
||||
|
||||
# test encode and decode
|
||||
sr, pred_wavs, latents, latent_lengths = model(audios, audio_lengths, sr)
|
||||
print("reconstructed wavs: ", pred_wavs[0].shape)
|
||||
print("latents shape: ", latents.shape)
|
||||
print("latent_lengths: ", latent_lengths)
|
||||
print("sr: ", sr)
|
||||
torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr)
|
||||
print("test_reconstructed.flac")
|
||||
Executable
+107
@@ -0,0 +1,107 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from torchaudio.transforms import MelScale
|
||||
|
||||
|
||||
class LinearSpectrogram(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_fft=2048,
|
||||
win_length=2048,
|
||||
hop_length=512,
|
||||
center=False,
|
||||
mode="pow2_sqrt",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.n_fft = n_fft
|
||||
self.win_length = win_length
|
||||
self.hop_length = hop_length
|
||||
self.center = center
|
||||
self.mode = mode
|
||||
|
||||
self.register_buffer("window", torch.hann_window(win_length))
|
||||
|
||||
def forward(self, y: Tensor) -> Tensor:
|
||||
if y.ndim == 3:
|
||||
y = y.squeeze(1)
|
||||
|
||||
y = torch.nn.functional.pad(
|
||||
y.unsqueeze(1),
|
||||
(
|
||||
(self.win_length - self.hop_length) // 2,
|
||||
(self.win_length - self.hop_length + 1) // 2,
|
||||
),
|
||||
mode="reflect",
|
||||
).squeeze(1)
|
||||
dtype = y.dtype
|
||||
spec = torch.stft(
|
||||
y.float(),
|
||||
self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
win_length=self.win_length,
|
||||
window=self.window,
|
||||
center=self.center,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
spec = torch.view_as_real(spec)
|
||||
|
||||
if self.mode == "pow2_sqrt":
|
||||
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
|
||||
spec = spec.to(dtype)
|
||||
return spec
|
||||
|
||||
|
||||
class LogMelSpectrogram(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate=44100,
|
||||
n_fft=2048,
|
||||
win_length=2048,
|
||||
hop_length=512,
|
||||
n_mels=128,
|
||||
center=False,
|
||||
f_min=0.0,
|
||||
f_max=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.n_fft = n_fft
|
||||
self.win_length = win_length
|
||||
self.hop_length = hop_length
|
||||
self.center = center
|
||||
self.n_mels = n_mels
|
||||
self.f_min = f_min
|
||||
self.f_max = f_max or sample_rate // 2
|
||||
|
||||
self.spectrogram = LinearSpectrogram(n_fft, win_length, hop_length, center)
|
||||
self.mel_scale = MelScale(
|
||||
self.n_mels,
|
||||
self.sample_rate,
|
||||
self.f_min,
|
||||
self.f_max,
|
||||
self.n_fft // 2 + 1,
|
||||
"slaney",
|
||||
"slaney",
|
||||
)
|
||||
|
||||
def compress(self, x: Tensor) -> Tensor:
|
||||
return torch.log(torch.clamp(x, min=1e-5))
|
||||
|
||||
def decompress(self, x: Tensor) -> Tensor:
|
||||
return torch.exp(x)
|
||||
|
||||
def forward(self, x: Tensor, return_linear: bool = False) -> Tensor:
|
||||
linear = self.spectrogram(x)
|
||||
x = self.mel_scale(linear)
|
||||
x = self.compress(x)
|
||||
# print(x.shape)
|
||||
if return_linear:
|
||||
return x, self.compress(linear)
|
||||
|
||||
return x
|
||||
Executable
+576
@@ -0,0 +1,576 @@
|
||||
import librosa
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from functools import partial
|
||||
from math import prod
|
||||
from typing import Callable, Tuple, List
|
||||
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import Conv1d
|
||||
from torch.nn.utils import weight_norm
|
||||
from torch.nn.utils.parametrize import remove_parametrizations as remove_weight_norm
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.loaders import FromOriginalModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
|
||||
|
||||
try:
|
||||
from music_log_mel import LogMelSpectrogram
|
||||
except ImportError:
|
||||
from .music_log_mel import LogMelSpectrogram
|
||||
|
||||
|
||||
def drop_path(
|
||||
x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
|
||||
):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||
|
||||
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
|
||||
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
|
||||
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
|
||||
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
|
||||
'survival rate' as the argument.
|
||||
|
||||
""" # noqa: E501
|
||||
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (
|
||||
x.ndim - 1
|
||||
) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
||||
if keep_prob > 0.0 and scale_by_keep:
|
||||
random_tensor.div_(keep_prob)
|
||||
return x * random_tensor
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" # noqa: E501
|
||||
|
||||
def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
self.scale_by_keep = scale_by_keep
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"drop_prob={round(self.drop_prob,3):0.3f}"
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
|
||||
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
|
||||
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
|
||||
with shape (batch_size, channels, height, width).
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||
self.eps = eps
|
||||
self.data_format = data_format
|
||||
if self.data_format not in ["channels_last", "channels_first"]:
|
||||
raise NotImplementedError
|
||||
self.normalized_shape = (normalized_shape,)
|
||||
|
||||
def forward(self, x):
|
||||
if self.data_format == "channels_last":
|
||||
return F.layer_norm(
|
||||
x, self.normalized_shape, self.weight, self.bias, self.eps
|
||||
)
|
||||
elif self.data_format == "channels_first":
|
||||
u = x.mean(1, keepdim=True)
|
||||
s = (x - u).pow(2).mean(1, keepdim=True)
|
||||
x = (x - u) / torch.sqrt(s + self.eps)
|
||||
x = self.weight[:, None] * x + self.bias[:, None]
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
r"""ConvNeXt Block. There are two equivalent implementations:
|
||||
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
|
||||
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
|
||||
We use (2) as we find it slightly faster in PyTorch
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
|
||||
kernel_size (int): Kernel size for depthwise conv. Default: 7.
|
||||
dilation (int): Dilation for depthwise conv. Default: 1.
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
drop_path: float = 0.0,
|
||||
layer_scale_init_value: float = 1e-6,
|
||||
mlp_ratio: float = 4.0,
|
||||
kernel_size: int = 7,
|
||||
dilation: int = 1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dwconv = nn.Conv1d(
|
||||
dim,
|
||||
dim,
|
||||
kernel_size=kernel_size,
|
||||
padding=int(dilation * (kernel_size - 1) / 2),
|
||||
groups=dim,
|
||||
) # depthwise conv
|
||||
self.norm = LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(
|
||||
dim, int(mlp_ratio * dim)
|
||||
) # pointwise/1x1 convs, implemented with linear layers
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value *
|
||||
torch.ones((dim)), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
self.drop_path = DropPath(
|
||||
drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
def forward(self, x, apply_residual: bool = True):
|
||||
input = x
|
||||
|
||||
x = self.dwconv(x)
|
||||
x = x.permute(0, 2, 1) # (N, C, L) -> (N, L, C)
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
|
||||
x = x.permute(0, 2, 1) # (N, L, C) -> (N, C, L)
|
||||
x = self.drop_path(x)
|
||||
|
||||
if apply_residual:
|
||||
x = input + x
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class ParallelConvNeXtBlock(nn.Module):
|
||||
def __init__(self, kernel_sizes: List[int], *args, **kwargs):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
ConvNeXtBlock(kernel_size=kernel_size, *args, **kwargs)
|
||||
for kernel_size in kernel_sizes
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.stack(
|
||||
[block(x, apply_residual=False) for block in self.blocks] + [x],
|
||||
dim=1,
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
class ConvNeXtEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels=3,
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0.0,
|
||||
layer_scale_init_value=1e-6,
|
||||
kernel_sizes: Tuple[int] = (7,),
|
||||
):
|
||||
super().__init__()
|
||||
assert len(depths) == len(dims)
|
||||
|
||||
self.channel_layers = nn.ModuleList()
|
||||
stem = nn.Sequential(
|
||||
nn.Conv1d(
|
||||
input_channels,
|
||||
dims[0],
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
padding_mode="replicate",
|
||||
),
|
||||
LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
|
||||
)
|
||||
self.channel_layers.append(stem)
|
||||
|
||||
for i in range(len(depths) - 1):
|
||||
mid_layer = nn.Sequential(
|
||||
LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
|
||||
nn.Conv1d(dims[i], dims[i + 1], kernel_size=1),
|
||||
)
|
||||
self.channel_layers.append(mid_layer)
|
||||
|
||||
block_fn = (
|
||||
partial(ConvNeXtBlock, kernel_size=kernel_sizes[0])
|
||||
if len(kernel_sizes) == 1
|
||||
else partial(ParallelConvNeXtBlock, kernel_sizes=kernel_sizes)
|
||||
)
|
||||
|
||||
self.stages = nn.ModuleList()
|
||||
drop_path_rates = [
|
||||
x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
|
||||
]
|
||||
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = nn.Sequential(
|
||||
*[
|
||||
block_fn(
|
||||
dim=dims[i],
|
||||
drop_path=drop_path_rates[cur + j],
|
||||
layer_scale_init_value=layer_scale_init_value,
|
||||
)
|
||||
for j in range(depths[i])
|
||||
]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
self.norm = LayerNorm(dims[-1], eps=1e-6, data_format="channels_first")
|
||||
self.apply(self._init_weights)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
for channel_layer, stage in zip(self.channel_layers, self.stages):
|
||||
x = channel_layer(x)
|
||||
x = stage(x)
|
||||
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
def init_weights(m, mean=0.0, std=0.01):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
m.weight.data.normal_(mean, std)
|
||||
|
||||
|
||||
def get_padding(kernel_size, dilation=1):
|
||||
return (kernel_size * dilation - dilation) // 2
|
||||
|
||||
|
||||
class ResBlock1(torch.nn.Module):
|
||||
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
|
||||
super().__init__()
|
||||
|
||||
self.convs1 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[0],
|
||||
padding=get_padding(kernel_size, dilation[0]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[1],
|
||||
padding=get_padding(kernel_size, dilation[1]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[2],
|
||||
padding=get_padding(kernel_size, dilation[2]),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs1.apply(init_weights)
|
||||
|
||||
self.convs2 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs2.apply(init_weights)
|
||||
|
||||
def forward(self, x):
|
||||
for c1, c2 in zip(self.convs1, self.convs2):
|
||||
xt = F.silu(x)
|
||||
xt = c1(xt)
|
||||
xt = F.silu(xt)
|
||||
xt = c2(xt)
|
||||
x = xt + x
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for conv in self.convs1:
|
||||
remove_weight_norm(conv)
|
||||
for conv in self.convs2:
|
||||
remove_weight_norm(conv)
|
||||
|
||||
|
||||
class HiFiGANGenerator(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
hop_length: int = 512,
|
||||
upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
|
||||
upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
|
||||
resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
|
||||
resblock_dilation_sizes: Tuple[Tuple[int]] = (
|
||||
(1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
||||
num_mels: int = 128,
|
||||
upsample_initial_channel: int = 512,
|
||||
use_template: bool = True,
|
||||
pre_conv_kernel_size: int = 7,
|
||||
post_conv_kernel_size: int = 7,
|
||||
post_activation: Callable = partial(nn.SiLU, inplace=True),
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
assert (
|
||||
prod(upsample_rates) == hop_length
|
||||
), f"hop_length must be {prod(upsample_rates)}"
|
||||
|
||||
self.conv_pre = weight_norm(
|
||||
nn.Conv1d(
|
||||
num_mels,
|
||||
upsample_initial_channel,
|
||||
pre_conv_kernel_size,
|
||||
1,
|
||||
padding=get_padding(pre_conv_kernel_size),
|
||||
)
|
||||
)
|
||||
|
||||
self.num_upsamples = len(upsample_rates)
|
||||
self.num_kernels = len(resblock_kernel_sizes)
|
||||
|
||||
self.noise_convs = nn.ModuleList()
|
||||
self.use_template = use_template
|
||||
self.ups = nn.ModuleList()
|
||||
|
||||
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
|
||||
c_cur = upsample_initial_channel // (2 ** (i + 1))
|
||||
self.ups.append(
|
||||
weight_norm(
|
||||
nn.ConvTranspose1d(
|
||||
upsample_initial_channel // (2**i),
|
||||
upsample_initial_channel // (2 ** (i + 1)),
|
||||
k,
|
||||
u,
|
||||
padding=(k - u) // 2,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
if not use_template:
|
||||
continue
|
||||
|
||||
if i + 1 < len(upsample_rates):
|
||||
stride_f0 = np.prod(upsample_rates[i + 1:])
|
||||
self.noise_convs.append(
|
||||
Conv1d(
|
||||
1,
|
||||
c_cur,
|
||||
kernel_size=stride_f0 * 2,
|
||||
stride=stride_f0,
|
||||
padding=stride_f0 // 2,
|
||||
)
|
||||
)
|
||||
else:
|
||||
self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
|
||||
|
||||
self.resblocks = nn.ModuleList()
|
||||
for i in range(len(self.ups)):
|
||||
ch = upsample_initial_channel // (2 ** (i + 1))
|
||||
for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
|
||||
self.resblocks.append(ResBlock1(ch, k, d))
|
||||
|
||||
self.activation_post = post_activation()
|
||||
self.conv_post = weight_norm(
|
||||
nn.Conv1d(
|
||||
ch,
|
||||
1,
|
||||
post_conv_kernel_size,
|
||||
1,
|
||||
padding=get_padding(post_conv_kernel_size),
|
||||
)
|
||||
)
|
||||
self.ups.apply(init_weights)
|
||||
self.conv_post.apply(init_weights)
|
||||
|
||||
def forward(self, x, template=None):
|
||||
x = self.conv_pre(x)
|
||||
|
||||
for i in range(self.num_upsamples):
|
||||
x = F.silu(x, inplace=True)
|
||||
x = self.ups[i](x)
|
||||
|
||||
if self.use_template:
|
||||
x = x + self.noise_convs[i](template)
|
||||
|
||||
xs = None
|
||||
|
||||
for j in range(self.num_kernels):
|
||||
if xs is None:
|
||||
xs = self.resblocks[i * self.num_kernels + j](x)
|
||||
else:
|
||||
xs += self.resblocks[i * self.num_kernels + j](x)
|
||||
|
||||
x = xs / self.num_kernels
|
||||
|
||||
x = self.activation_post(x)
|
||||
x = self.conv_post(x)
|
||||
x = torch.tanh(x)
|
||||
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for up in self.ups:
|
||||
remove_weight_norm(up)
|
||||
for block in self.resblocks:
|
||||
block.remove_weight_norm()
|
||||
remove_weight_norm(self.conv_pre)
|
||||
remove_weight_norm(self.conv_post)
|
||||
|
||||
|
||||
class ADaMoSHiFiGANV1(ModelMixin, ConfigMixin, FromOriginalModelMixin):
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int = 128,
|
||||
depths: List[int] = [3, 3, 9, 3],
|
||||
dims: List[int] = [128, 256, 384, 512],
|
||||
drop_path_rate: float = 0.0,
|
||||
kernel_sizes: Tuple[int] = (7,),
|
||||
upsample_rates: Tuple[int] = (4, 4, 2, 2, 2, 2, 2),
|
||||
upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
|
||||
resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
|
||||
resblock_dilation_sizes: Tuple[Tuple[int]] = (
|
||||
(1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
||||
num_mels: int = 512,
|
||||
upsample_initial_channel: int = 1024,
|
||||
use_template: bool = False,
|
||||
pre_conv_kernel_size: int = 13,
|
||||
post_conv_kernel_size: int = 13,
|
||||
sampling_rate: int = 44100,
|
||||
n_fft: int = 2048,
|
||||
win_length: int = 2048,
|
||||
hop_length: int = 512,
|
||||
f_min: int = 40,
|
||||
f_max: int = 16000,
|
||||
n_mels: int = 128,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.backbone = ConvNeXtEncoder(
|
||||
input_channels=input_channels,
|
||||
depths=depths,
|
||||
dims=dims,
|
||||
drop_path_rate=drop_path_rate,
|
||||
kernel_sizes=kernel_sizes,
|
||||
)
|
||||
|
||||
self.head = HiFiGANGenerator(
|
||||
hop_length=hop_length,
|
||||
upsample_rates=upsample_rates,
|
||||
upsample_kernel_sizes=upsample_kernel_sizes,
|
||||
resblock_kernel_sizes=resblock_kernel_sizes,
|
||||
resblock_dilation_sizes=resblock_dilation_sizes,
|
||||
num_mels=num_mels,
|
||||
upsample_initial_channel=upsample_initial_channel,
|
||||
use_template=use_template,
|
||||
pre_conv_kernel_size=pre_conv_kernel_size,
|
||||
post_conv_kernel_size=post_conv_kernel_size,
|
||||
)
|
||||
self.sampling_rate = sampling_rate
|
||||
self.mel_transform = LogMelSpectrogram(
|
||||
sample_rate=sampling_rate,
|
||||
n_fft=n_fft,
|
||||
win_length=win_length,
|
||||
hop_length=hop_length,
|
||||
f_min=f_min,
|
||||
f_max=f_max,
|
||||
n_mels=n_mels,
|
||||
)
|
||||
self.eval()
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, mel):
|
||||
y = self.backbone(mel)
|
||||
y = self.head(y)
|
||||
return y
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, x):
|
||||
return self.mel_transform(x)
|
||||
|
||||
def forward(self, mel):
|
||||
y = self.backbone(mel)
|
||||
y = self.head(y)
|
||||
return y
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import soundfile as sf
|
||||
|
||||
x = "test_audio.flac"
|
||||
model = ADaMoSHiFiGANV1.from_pretrained("./checkpoints/music_vocoder", local_files_only=True)
|
||||
|
||||
wav, sr = librosa.load(x, sr=44100, mono=True)
|
||||
wav = torch.from_numpy(wav).float()[None]
|
||||
mel = model.encode(wav)
|
||||
|
||||
wav = model.decode(mel)[0].mT
|
||||
sf.write("test_audio_vocoder_rec.flac", wav.cpu().numpy(), 44100)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,394 @@
|
||||
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.utils import BaseOutput, logging
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's `step` function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
"""
|
||||
|
||||
prev_sample: torch.FloatTensor
|
||||
|
||||
|
||||
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
|
||||
"""
|
||||
Euler scheduler.
|
||||
|
||||
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
|
||||
methods the library implements for all schedulers such as loading and saving.
|
||||
|
||||
Args:
|
||||
num_train_timesteps (`int`, defaults to 1000):
|
||||
The number of diffusion steps to train the model.
|
||||
timestep_spacing (`str`, defaults to `"linspace"`):
|
||||
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
|
||||
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
|
||||
shift (`float`, defaults to 1.0):
|
||||
The shift value for the timestep schedule.
|
||||
"""
|
||||
|
||||
_compatibles = []
|
||||
order = 1
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_train_timesteps: int = 1000,
|
||||
shift: float = 1.0,
|
||||
use_dynamic_shifting=False,
|
||||
base_shift: Optional[float] = 0.5,
|
||||
max_shift: Optional[float] = 1.15,
|
||||
base_image_seq_len: Optional[int] = 256,
|
||||
max_image_seq_len: Optional[int] = 4096,
|
||||
):
|
||||
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
|
||||
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
|
||||
|
||||
sigmas = timesteps / num_train_timesteps
|
||||
if not use_dynamic_shifting:
|
||||
# when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
|
||||
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
|
||||
|
||||
self.timesteps = sigmas * num_train_timesteps
|
||||
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
|
||||
self.sigma_min = self.sigmas[-1].item()
|
||||
self.sigma_max = self.sigmas[0].item()
|
||||
|
||||
@property
|
||||
def step_index(self):
|
||||
"""
|
||||
The index counter for current timestep. It will increase 1 after each scheduler step.
|
||||
"""
|
||||
return self._step_index
|
||||
|
||||
@property
|
||||
def begin_index(self):
|
||||
"""
|
||||
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
|
||||
"""
|
||||
return self._begin_index
|
||||
|
||||
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
|
||||
def set_begin_index(self, begin_index: int = 0):
|
||||
"""
|
||||
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
|
||||
|
||||
Args:
|
||||
begin_index (`int`):
|
||||
The begin index for the scheduler.
|
||||
"""
|
||||
self._begin_index = begin_index
|
||||
|
||||
def scale_noise(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
timestep: Union[float, torch.FloatTensor],
|
||||
noise: Optional[torch.FloatTensor] = None,
|
||||
) -> torch.FloatTensor:
|
||||
"""
|
||||
Forward process in flow-matching
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor`):
|
||||
The input sample.
|
||||
timestep (`int`, *optional*):
|
||||
The current timestep in the diffusion chain.
|
||||
|
||||
Returns:
|
||||
`torch.FloatTensor`:
|
||||
A scaled input sample.
|
||||
"""
|
||||
# Make sure sigmas and timesteps have the same device and dtype as original_samples
|
||||
sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
|
||||
|
||||
if sample.device.type == "mps" and torch.is_floating_point(timestep):
|
||||
# mps does not support float64
|
||||
schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
|
||||
timestep = timestep.to(sample.device, dtype=torch.float32)
|
||||
else:
|
||||
schedule_timesteps = self.timesteps.to(sample.device)
|
||||
timestep = timestep.to(sample.device)
|
||||
|
||||
# self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
|
||||
if self.begin_index is None:
|
||||
step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
|
||||
elif self.step_index is not None:
|
||||
# add_noise is called after first denoising step (for inpainting)
|
||||
step_indices = [self.step_index] * timestep.shape[0]
|
||||
else:
|
||||
# add noise is called before first denoising step to create initial latent(img2img)
|
||||
step_indices = [self.begin_index] * timestep.shape[0]
|
||||
|
||||
sigma = sigmas[step_indices].flatten()
|
||||
while len(sigma.shape) < len(sample.shape):
|
||||
sigma = sigma.unsqueeze(-1)
|
||||
|
||||
sample = sigma * noise + (1.0 - sigma) * sample
|
||||
|
||||
return sample
|
||||
|
||||
def _sigma_to_t(self, sigma):
|
||||
return sigma * self.config.num_train_timesteps
|
||||
|
||||
def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
def set_timesteps(
|
||||
self,
|
||||
num_inference_steps: int = None,
|
||||
device: Union[str, torch.device] = None,
|
||||
sigmas: Optional[List[float]] = None,
|
||||
mu: Optional[float] = None,
|
||||
):
|
||||
"""
|
||||
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
|
||||
|
||||
Args:
|
||||
num_inference_steps (`int`):
|
||||
The number of diffusion steps used when generating samples with a pre-trained model.
|
||||
device (`str` or `torch.device`, *optional*):
|
||||
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
||||
"""
|
||||
|
||||
if self.config.use_dynamic_shifting and mu is None:
|
||||
raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`")
|
||||
|
||||
if sigmas is None:
|
||||
self.num_inference_steps = num_inference_steps
|
||||
timesteps = np.linspace(
|
||||
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
|
||||
)
|
||||
|
||||
sigmas = timesteps / self.config.num_train_timesteps
|
||||
|
||||
if self.config.use_dynamic_shifting:
|
||||
sigmas = self.time_shift(mu, 1.0, sigmas)
|
||||
else:
|
||||
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
|
||||
|
||||
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
|
||||
timesteps = sigmas * self.config.num_train_timesteps
|
||||
|
||||
self.timesteps = timesteps.to(device=device)
|
||||
self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
|
||||
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
||||
if schedule_timesteps is None:
|
||||
schedule_timesteps = self.timesteps
|
||||
|
||||
indices = (schedule_timesteps == timestep).nonzero()
|
||||
|
||||
# The sigma index that is taken for the **very** first `step`
|
||||
# is always the second index (or the last index if there is only 1)
|
||||
# This way we can ensure we don't accidentally skip a sigma in
|
||||
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
|
||||
pos = 1 if len(indices) > 1 else 0
|
||||
|
||||
return indices[pos].item()
|
||||
|
||||
def _init_step_index(self, timestep):
|
||||
if self.begin_index is None:
|
||||
if isinstance(timestep, torch.Tensor):
|
||||
timestep = timestep.to(self.timesteps.device)
|
||||
self._step_index = self.index_for_timestep(timestep)
|
||||
else:
|
||||
self._step_index = self._begin_index
|
||||
|
||||
def step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: Union[float, torch.FloatTensor],
|
||||
sample: torch.FloatTensor,
|
||||
s_churn: float = 0.0,
|
||||
s_tmin: float = 0.0,
|
||||
s_tmax: float = float("inf"),
|
||||
s_noise: float = 1.0,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
return_dict: bool = True,
|
||||
omega: Union[float, np.array] = 0.0
|
||||
) -> Union[FlowMatchEulerDiscreteSchedulerOutput, Tuple]:
|
||||
"""
|
||||
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
|
||||
Args:
|
||||
model_output (`torch.FloatTensor`):
|
||||
The direct output from learned diffusion model.
|
||||
timestep (`float`):
|
||||
The current discrete timestep in the diffusion chain.
|
||||
sample (`torch.FloatTensor`):
|
||||
A current instance of a sample created by the diffusion process.
|
||||
s_churn (`float`):
|
||||
s_tmin (`float`):
|
||||
s_tmax (`float`):
|
||||
s_noise (`float`, defaults to 1.0):
|
||||
Scaling factor for noise added to the sample.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
A random number generator.
|
||||
return_dict (`bool`):
|
||||
Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
[`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`:
|
||||
If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is
|
||||
returned, otherwise a tuple is returned where the first element is the sample tensor.
|
||||
"""
|
||||
|
||||
def logistic_function(x, L=0.9, U=1.1, x_0=0.0, k=1):
|
||||
# L = Lower bound
|
||||
# U = Upper bound
|
||||
# x_0 = Midpoint (x corresponding to y = 1.0)
|
||||
# k = Steepness, can adjust based on preference
|
||||
|
||||
if isinstance(x, torch.Tensor):
|
||||
device_ = x.device
|
||||
x = x.to(torch.float).cpu().numpy()
|
||||
|
||||
new_x = L + (U - L) / (1 + np.exp(-k * (x - x_0)))
|
||||
|
||||
if isinstance(new_x, np.ndarray):
|
||||
new_x = torch.from_numpy(new_x).to(device_)
|
||||
return new_x
|
||||
|
||||
self.omega_bef_rescale = omega
|
||||
omega = logistic_function(omega, k=0.1)
|
||||
self.omega_aft_rescale = omega
|
||||
|
||||
if (
|
||||
isinstance(timestep, int)
|
||||
or isinstance(timestep, torch.IntTensor)
|
||||
or isinstance(timestep, torch.LongTensor)
|
||||
):
|
||||
raise ValueError(
|
||||
(
|
||||
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
|
||||
" `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
|
||||
" one of the `scheduler.timesteps` as a timestep."
|
||||
),
|
||||
)
|
||||
|
||||
if self.step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
|
||||
# Upcast to avoid precision issues when computing prev_sample
|
||||
sample = sample.to(torch.float32)
|
||||
|
||||
sigma = self.sigmas[self.step_index]
|
||||
sigma_next = self.sigmas[self.step_index + 1]
|
||||
|
||||
## --
|
||||
## mean shift 1
|
||||
dx = (sigma_next - sigma) * model_output
|
||||
m = dx.mean()
|
||||
# print(dx.shape) # torch.Size([1, 16, 128, 128])
|
||||
# print(f'm: {m}') # m: -0.0014209747314453125
|
||||
# raise NotImplementedError
|
||||
dx_ = (dx - m) * omega + m
|
||||
prev_sample = sample + dx_
|
||||
|
||||
# ## --
|
||||
# ## mean shift 2
|
||||
# m = model_output.mean()
|
||||
# model_output_ = (model_output - m) * omega + m
|
||||
# prev_sample = sample + (sigma_next - sigma) * model_output_
|
||||
|
||||
# ## --
|
||||
# ## original
|
||||
# prev_sample = sample + (sigma_next - sigma) * model_output * omega
|
||||
|
||||
# ## --
|
||||
# ## spatial mean 1
|
||||
# dx = (sigma_next - sigma) * model_output
|
||||
# m = dx.mean(dim=(0, 1), keepdim=True)
|
||||
# # print(dx.shape) # torch.Size([1, 16, 128, 128])
|
||||
# # print(m.shape) # torch.Size([1, 1, 128, 128])
|
||||
# # raise NotImplementedError
|
||||
# dx_ = (dx - m) * omega + m
|
||||
# prev_sample = sample + dx_
|
||||
|
||||
# ## --
|
||||
# ## spatial mean 2
|
||||
# m = model_output.mean(dim=(0, 1), keepdim=True)
|
||||
# model_output_ = (model_output - m) * omega + m
|
||||
# prev_sample = sample + (sigma_next - sigma) * model_output_
|
||||
|
||||
# ## --
|
||||
# ## channel mean 1
|
||||
# m = model_output.mean(dim=(2, 3), keepdim=True)
|
||||
# # print(m.shape) # torch.Size([1, 16, 1, 1])
|
||||
# model_output_ = (model_output - m) * omega + m
|
||||
# prev_sample = sample + (sigma_next - sigma) * model_output_
|
||||
|
||||
# ## --
|
||||
# ## channel mean 2
|
||||
# dx = (sigma_next - sigma) * model_output
|
||||
# m = dx.mean(dim=(2, 3), keepdim=True)
|
||||
# # print(m.shape) # torch.Size([1, 16, 1, 1])
|
||||
# dx_ = (dx - m) * omega + m
|
||||
# prev_sample = sample + dx_
|
||||
|
||||
# ## --
|
||||
# ## keep sample mean
|
||||
# m_tgt = sample.mean()
|
||||
# prev_sample_ = sample + (sigma_next - sigma) * model_output * omega
|
||||
# m_src = prev_sample_.mean()
|
||||
# prev_sample = prev_sample_ - m_src + m_tgt
|
||||
|
||||
# ## --
|
||||
# ## test
|
||||
# # print(sample.mean())
|
||||
# prev_sample = sample + (sigma_next - sigma) * model_output * omega
|
||||
# # raise NotImplementedError
|
||||
|
||||
# Cast sample back to model compatible dtype
|
||||
prev_sample = prev_sample.to(model_output.dtype)
|
||||
|
||||
# upon completion increase step index by one
|
||||
self._step_index += 1
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
|
||||
|
||||
def __len__(self):
|
||||
return self.config.num_train_timesteps
|
||||
@@ -0,0 +1,348 @@
|
||||
# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.utils import BaseOutput, logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class FlowMatchHeunDiscreteSchedulerOutput(BaseOutput):
|
||||
"""
|
||||
Output class for the scheduler's `step` function output.
|
||||
|
||||
Args:
|
||||
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
||||
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
||||
denoising loop.
|
||||
"""
|
||||
|
||||
prev_sample: torch.FloatTensor
|
||||
|
||||
|
||||
class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
|
||||
"""
|
||||
Heun scheduler.
|
||||
|
||||
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
|
||||
methods the library implements for all schedulers such as loading and saving.
|
||||
|
||||
Args:
|
||||
num_train_timesteps (`int`, defaults to 1000):
|
||||
The number of diffusion steps to train the model.
|
||||
timestep_spacing (`str`, defaults to `"linspace"`):
|
||||
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
|
||||
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
|
||||
shift (`float`, defaults to 1.0):
|
||||
The shift value for the timestep schedule.
|
||||
"""
|
||||
|
||||
_compatibles = []
|
||||
order = 2
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_train_timesteps: int = 1000,
|
||||
shift: float = 1.0,
|
||||
):
|
||||
timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
|
||||
timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
|
||||
|
||||
sigmas = timesteps / num_train_timesteps
|
||||
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
|
||||
|
||||
self.timesteps = sigmas * num_train_timesteps
|
||||
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication
|
||||
self.sigma_min = self.sigmas[-1].item()
|
||||
self.sigma_max = self.sigmas[0].item()
|
||||
|
||||
@property
|
||||
def step_index(self):
|
||||
"""
|
||||
The index counter for current timestep. It will increase 1 after each scheduler step.
|
||||
"""
|
||||
return self._step_index
|
||||
|
||||
@property
|
||||
def begin_index(self):
|
||||
"""
|
||||
The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
|
||||
"""
|
||||
return self._begin_index
|
||||
|
||||
# Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
|
||||
def set_begin_index(self, begin_index: int = 0):
|
||||
"""
|
||||
Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
|
||||
|
||||
Args:
|
||||
begin_index (`int`):
|
||||
The begin index for the scheduler.
|
||||
"""
|
||||
self._begin_index = begin_index
|
||||
|
||||
def scale_noise(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
timestep: Union[float, torch.FloatTensor],
|
||||
noise: Optional[torch.FloatTensor] = None,
|
||||
) -> torch.FloatTensor:
|
||||
"""
|
||||
Forward process in flow-matching
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor`):
|
||||
The input sample.
|
||||
timestep (`int`, *optional*):
|
||||
The current timestep in the diffusion chain.
|
||||
|
||||
Returns:
|
||||
`torch.FloatTensor`:
|
||||
A scaled input sample.
|
||||
"""
|
||||
if self.step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
|
||||
sigma = self.sigmas[self.step_index]
|
||||
sample = sigma * noise + (1.0 - sigma) * sample
|
||||
|
||||
return sample
|
||||
|
||||
def _sigma_to_t(self, sigma):
|
||||
return sigma * self.config.num_train_timesteps
|
||||
|
||||
def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
|
||||
"""
|
||||
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
|
||||
|
||||
Args:
|
||||
num_inference_steps (`int`):
|
||||
The number of diffusion steps used when generating samples with a pre-trained model.
|
||||
device (`str` or `torch.device`, *optional*):
|
||||
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
||||
"""
|
||||
self.num_inference_steps = num_inference_steps
|
||||
|
||||
timesteps = np.linspace(
|
||||
self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
|
||||
)
|
||||
|
||||
sigmas = timesteps / self.config.num_train_timesteps
|
||||
sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas)
|
||||
sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
|
||||
|
||||
timesteps = sigmas * self.config.num_train_timesteps
|
||||
timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2)])
|
||||
self.timesteps = timesteps.to(device=device)
|
||||
|
||||
sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
|
||||
self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
|
||||
|
||||
# empty dt and derivative
|
||||
self.prev_derivative = None
|
||||
self.dt = None
|
||||
|
||||
self._step_index = None
|
||||
self._begin_index = None
|
||||
|
||||
def index_for_timestep(self, timestep, schedule_timesteps=None):
|
||||
if schedule_timesteps is None:
|
||||
schedule_timesteps = self.timesteps
|
||||
|
||||
indices = (schedule_timesteps == timestep).nonzero()
|
||||
|
||||
# The sigma index that is taken for the **very** first `step`
|
||||
# is always the second index (or the last index if there is only 1)
|
||||
# This way we can ensure we don't accidentally skip a sigma in
|
||||
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
|
||||
pos = 1 if len(indices) > 1 else 0
|
||||
|
||||
return indices[pos].item()
|
||||
|
||||
def _init_step_index(self, timestep):
|
||||
if self.begin_index is None:
|
||||
if isinstance(timestep, torch.Tensor):
|
||||
timestep = timestep.to(self.timesteps.device)
|
||||
self._step_index = self.index_for_timestep(timestep)
|
||||
else:
|
||||
self._step_index = self._begin_index
|
||||
|
||||
@property
|
||||
def state_in_first_order(self):
|
||||
return self.dt is None
|
||||
|
||||
def step(
|
||||
self,
|
||||
model_output: torch.FloatTensor,
|
||||
timestep: Union[float, torch.FloatTensor],
|
||||
sample: torch.FloatTensor,
|
||||
s_churn: float = 0.0,
|
||||
s_tmin: float = 0.0,
|
||||
s_tmax: float = float("inf"),
|
||||
s_noise: float = 1.0,
|
||||
generator: Optional[torch.Generator] = None,
|
||||
return_dict: bool = True,
|
||||
omega: Union[float, np.array] = 0.0
|
||||
) -> Union[FlowMatchHeunDiscreteSchedulerOutput, Tuple]:
|
||||
"""
|
||||
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
||||
process from the learned model outputs (most often the predicted noise).
|
||||
|
||||
Args:
|
||||
model_output (`torch.FloatTensor`):
|
||||
The direct output from learned diffusion model.
|
||||
timestep (`float`):
|
||||
The current discrete timestep in the diffusion chain.
|
||||
sample (`torch.FloatTensor`):
|
||||
A current instance of a sample created by the diffusion process.
|
||||
s_churn (`float`):
|
||||
s_tmin (`float`):
|
||||
s_tmax (`float`):
|
||||
s_noise (`float`, defaults to 1.0):
|
||||
Scaling factor for noise added to the sample.
|
||||
generator (`torch.Generator`, *optional*):
|
||||
A random number generator.
|
||||
return_dict (`bool`):
|
||||
Whether or not to return a [`~schedulers.scheduling_Heun_discrete.HeunDiscreteSchedulerOutput`] or
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
[`~schedulers.scheduling_Heun_discrete.HeunDiscreteSchedulerOutput`] or `tuple`:
|
||||
If return_dict is `True`, [`~schedulers.scheduling_Heun_discrete.HeunDiscreteSchedulerOutput`] is
|
||||
returned, otherwise a tuple is returned where the first element is the sample tensor.
|
||||
"""
|
||||
|
||||
def logistic_function(x, L=0.9, U=1.1, x_0=0.0, k=1):
|
||||
# L = Lower bound
|
||||
# U = Upper bound
|
||||
# x_0 = Midpoint (x corresponding to y = 1.0)
|
||||
# k = Steepness, can adjust based on preference
|
||||
|
||||
if isinstance(x, torch.Tensor):
|
||||
device_ = x.device
|
||||
x = x.to(torch.float).cpu().numpy()
|
||||
|
||||
new_x = L + (U - L) / (1 + np.exp(-k * (x - x_0)))
|
||||
|
||||
if isinstance(new_x, np.ndarray):
|
||||
new_x = torch.from_numpy(new_x).to(device_)
|
||||
return new_x
|
||||
|
||||
self.omega_bef_rescale = omega
|
||||
omega = logistic_function(omega, k=0.1)
|
||||
self.omega_aft_rescale = omega
|
||||
|
||||
if (
|
||||
isinstance(timestep, int)
|
||||
or isinstance(timestep, torch.IntTensor)
|
||||
or isinstance(timestep, torch.LongTensor)
|
||||
):
|
||||
raise ValueError(
|
||||
(
|
||||
"Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
|
||||
" `HeunDiscreteScheduler.step()` is not supported. Make sure to pass"
|
||||
" one of the `scheduler.timesteps` as a timestep."
|
||||
),
|
||||
)
|
||||
|
||||
if self.step_index is None:
|
||||
self._init_step_index(timestep)
|
||||
|
||||
# Upcast to avoid precision issues when computing prev_sample
|
||||
sample = sample.to(torch.float32)
|
||||
|
||||
if self.state_in_first_order:
|
||||
sigma = self.sigmas[self.step_index]
|
||||
sigma_next = self.sigmas[self.step_index + 1]
|
||||
else:
|
||||
# 2nd order / Heun's method
|
||||
sigma = self.sigmas[self.step_index - 1]
|
||||
sigma_next = self.sigmas[self.step_index]
|
||||
|
||||
gamma = min(s_churn / (len(self.sigmas) - 1), 2**0.5 - 1) if s_tmin <= sigma <= s_tmax else 0.0
|
||||
|
||||
sigma_hat = sigma * (gamma + 1)
|
||||
|
||||
if gamma > 0:
|
||||
noise = randn_tensor(
|
||||
model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
|
||||
)
|
||||
eps = noise * s_noise
|
||||
sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
|
||||
|
||||
if self.state_in_first_order:
|
||||
# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
|
||||
denoised = sample - model_output * sigma
|
||||
# 2. convert to an ODE derivative for 1st order
|
||||
derivative = (sample - denoised) / sigma_hat
|
||||
# 3. Delta timestep
|
||||
dt = sigma_next - sigma_hat
|
||||
|
||||
# store for 2nd order step
|
||||
self.prev_derivative = derivative
|
||||
self.dt = dt
|
||||
self.sample = sample
|
||||
else:
|
||||
# 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
|
||||
denoised = sample - model_output * sigma_next
|
||||
# 2. 2nd order / Heun's method
|
||||
derivative = (sample - denoised) / sigma_next
|
||||
derivative = 0.5 * (self.prev_derivative + derivative)
|
||||
|
||||
# 3. take prev timestep & sample
|
||||
dt = self.dt
|
||||
sample = self.sample
|
||||
|
||||
# free dt and derivative
|
||||
# Note, this puts the scheduler in "first order mode"
|
||||
self.prev_derivative = None
|
||||
self.dt = None
|
||||
self.sample = None
|
||||
|
||||
# original sample way
|
||||
# prev_sample = sample + derivative * dt
|
||||
|
||||
dx = derivative * dt
|
||||
m = dx.mean()
|
||||
dx_ = (dx - m) * omega + m
|
||||
prev_sample = sample + dx_
|
||||
|
||||
# Cast sample back to model compatible dtype
|
||||
prev_sample = prev_sample.to(model_output.dtype)
|
||||
|
||||
# upon completion increase step index by one
|
||||
self._step_index += 1
|
||||
|
||||
if not return_dict:
|
||||
return (prev_sample,)
|
||||
|
||||
return FlowMatchHeunDiscreteSchedulerOutput(prev_sample=prev_sample)
|
||||
|
||||
def __len__(self):
|
||||
return self.config.num_train_timesteps
|
||||
@@ -0,0 +1,555 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import random
|
||||
from torch.utils.data import Dataset
|
||||
from datasets import load_from_disk
|
||||
from loguru import logger
|
||||
import time
|
||||
import traceback
|
||||
import torchaudio
|
||||
from pathlib import Path
|
||||
import re
|
||||
from language_segmentation import LangSegment
|
||||
from models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer
|
||||
import warnings
|
||||
|
||||
warnings.simplefilter("ignore", category=FutureWarning)
|
||||
|
||||
DEFAULT_TRAIN_PATH = "./data/example_dataset"
|
||||
|
||||
def is_silent_audio(audio_tensor, silence_threshold=0.95):
|
||||
"""
|
||||
Determine if an audio is silent and should be discarded
|
||||
|
||||
Args:
|
||||
audio_tensor: torch.Tensor from torchaudio, shape (num_channels, num_samples)
|
||||
silence_threshold: Silence threshold ratio, default 0.95 means 95%
|
||||
|
||||
Returns:
|
||||
bool: True if audio should be discarded, False if it should be kept
|
||||
"""
|
||||
# Check if each sample point is zero across all channels
|
||||
silent_samples = torch.all(audio_tensor == 0, dim=0)
|
||||
|
||||
# Calculate silence ratio
|
||||
silent_ratio = torch.mean(silent_samples.float()).item()
|
||||
|
||||
return silent_ratio > silence_threshold
|
||||
|
||||
|
||||
# Supported languages for tokenization
|
||||
SUPPORT_LANGUAGES = {
|
||||
"en": 259, "de": 260, "fr": 262, "es": 284, "it": 285,
|
||||
"pt": 286, "pl": 294, "tr": 295, "ru": 267, "cs": 293,
|
||||
"nl": 297, "ar": 5022, "zh": 5023, "ja": 5412, "hu": 5753,
|
||||
"ko": 6152, "hi": 6680
|
||||
}
|
||||
|
||||
# Regex pattern for structure markers like [Verse], [Chorus], etc.
|
||||
structure_pattern = re.compile(r"\[.*?\]")
|
||||
|
||||
|
||||
class Text2MusicDataset(Dataset):
|
||||
"""
|
||||
Dataset for text-to-music generation that processes lyrics and audio files
|
||||
"""
|
||||
|
||||
def __init__(self, train=True, train_dataset_path=DEFAULT_TRAIN_PATH,
|
||||
max_duration=240.0, sample_size=None, shuffle=True,
|
||||
minibatch_size=1):
|
||||
"""
|
||||
Initialize the Text2Music dataset
|
||||
|
||||
Args:
|
||||
train: Whether this is a training dataset
|
||||
train_dataset_path: Path to the dataset
|
||||
max_duration: Maximum audio duration in seconds
|
||||
sample_size: Optional limit on number of samples to use
|
||||
shuffle: Whether to shuffle the dataset
|
||||
minibatch_size: Size of mini-batches
|
||||
"""
|
||||
self.train_dataset_path = train_dataset_path
|
||||
self.max_duration = max_duration
|
||||
self.minibatch_size = minibatch_size
|
||||
self.train = train
|
||||
|
||||
# Initialize language segmentation
|
||||
self.lang_segment = LangSegment()
|
||||
self.lang_segment.setfilters([
|
||||
'af', 'am', 'an', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'dz', 'el',
|
||||
'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fo', 'fr', 'ga', 'gl', 'gu', 'he', 'hi', 'hr', 'ht', 'hu', 'hy',
|
||||
'id', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lb', 'lo', 'lt', 'lv', 'mg',
|
||||
'mk', 'ml', 'mn', 'mr', 'ms', 'mt', 'nb', 'ne', 'nl', 'nn', 'no', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'qu',
|
||||
'ro', 'ru', 'rw', 'se', 'si', 'sk', 'sl', 'sq', 'sr', 'sv', 'sw', 'ta', 'te', 'th', 'tl', 'tr', 'ug', 'uk',
|
||||
'ur', 'vi', 'vo', 'wa', 'xh', 'zh', 'zu'
|
||||
])
|
||||
|
||||
# Initialize lyric tokenizer
|
||||
self.lyric_tokenizer = VoiceBpeTokenizer()
|
||||
|
||||
# Load dataset
|
||||
self.setup_full(train, shuffle, sample_size)
|
||||
logger.info(f"Dataset size: {len(self)} total {self.total_samples} samples")
|
||||
|
||||
def setup_full(self, train=True, shuffle=True, sample_size=None):
|
||||
"""
|
||||
Load and prepare the dataset
|
||||
|
||||
Args:
|
||||
train: Whether this is a training dataset
|
||||
shuffle: Whether to shuffle the dataset
|
||||
sample_size: Optional limit on number of samples to use
|
||||
"""
|
||||
pretrain_ds = load_from_disk(self.train_dataset_path)
|
||||
|
||||
if sample_size is not None:
|
||||
pretrain_ds = pretrain_ds.select(range(sample_size))
|
||||
|
||||
self.pretrain_ds = pretrain_ds
|
||||
self.total_samples = len(self.pretrain_ds)
|
||||
|
||||
def __len__(self):
|
||||
"""Return the number of batches in the dataset"""
|
||||
if self.total_samples % self.minibatch_size == 0:
|
||||
return self.total_samples // self.minibatch_size
|
||||
else:
|
||||
return self.total_samples // self.minibatch_size + 1
|
||||
|
||||
def get_lang(self, text):
|
||||
"""
|
||||
Detect the language of a text
|
||||
|
||||
Args:
|
||||
text: Input text
|
||||
|
||||
Returns:
|
||||
tuple: (primary_language, language_segments, language_counts)
|
||||
"""
|
||||
language = "en"
|
||||
langs = []
|
||||
try:
|
||||
langs = self.lang_segment.getTexts(text)
|
||||
langCounts = self.lang_segment.getCounts()
|
||||
language = langCounts[0][0]
|
||||
# If primary language is English but there's another language, use the second one
|
||||
if len(langCounts) > 1 and language == "en":
|
||||
language = langCounts[1][0]
|
||||
except Exception:
|
||||
language = "en"
|
||||
return language, langs, langCounts
|
||||
|
||||
def tokenize_lyrics(self, lyrics, debug=False, key=None):
|
||||
"""
|
||||
Tokenize lyrics into token indices
|
||||
|
||||
Args:
|
||||
lyrics: Lyrics text
|
||||
debug: Whether to print debug information
|
||||
key: Optional key identifier
|
||||
|
||||
Returns:
|
||||
list: Token indices
|
||||
"""
|
||||
lines = lyrics.split("\n")
|
||||
lyric_token_idx = [261] # Start token
|
||||
|
||||
# Detect language
|
||||
lang, langs, lang_counter = self.get_lang(lyrics)
|
||||
|
||||
# Determine most common language
|
||||
most_common_lang = "en"
|
||||
if len(lang_counter) > 0:
|
||||
most_common_lang = lang_counter[0][0]
|
||||
if most_common_lang == "":
|
||||
most_common_lang = "en"
|
||||
|
||||
if most_common_lang not in SUPPORT_LANGUAGES:
|
||||
raise ValueError(f"Unsupported language: {most_common_lang}")
|
||||
|
||||
# Process each language segment
|
||||
for lang_seg in langs:
|
||||
lang = lang_seg["lang"]
|
||||
text = lang_seg["text"]
|
||||
|
||||
# Normalize language codes
|
||||
if lang not in SUPPORT_LANGUAGES:
|
||||
lang = "en"
|
||||
if "zh" in lang:
|
||||
lang = "zh"
|
||||
if "spa" in lang:
|
||||
lang = "es"
|
||||
|
||||
# Process each line in the segment
|
||||
lines = text.split("\n")
|
||||
for line in lines:
|
||||
if not line.strip():
|
||||
lyric_token_idx += [2] # Line break token
|
||||
continue
|
||||
|
||||
try:
|
||||
# Handle structure markers like [Verse], [Chorus]
|
||||
if structure_pattern.match(line):
|
||||
token_idx = self.lyric_tokenizer.encode(line, "en")
|
||||
else:
|
||||
# Try tokenizing with most common language first
|
||||
token_idx = self.lyric_tokenizer.encode(line, most_common_lang)
|
||||
|
||||
# If debug mode, show tokenization results
|
||||
if debug:
|
||||
toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
|
||||
logger.info(f"debug using most_common_lang {line} --> {most_common_lang} --> {toks}")
|
||||
|
||||
# If tokenization contains unknown token (1), try with segment language
|
||||
if 1 in token_idx:
|
||||
token_idx = self.lyric_tokenizer.encode(line, lang)
|
||||
|
||||
if debug:
|
||||
toks = self.lyric_tokenizer.batch_decode([[tok_id] for tok_id in token_idx])
|
||||
logger.info(f"debug {line} --> {lang} --> {toks}")
|
||||
|
||||
# Add tokens and line break
|
||||
lyric_token_idx = lyric_token_idx + token_idx + [2]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Tokenize error: {e} for line: {line}, major_language: {lang}")
|
||||
|
||||
return lyric_token_idx
|
||||
|
||||
def tokenize_lyrics_map(self, item, debug=False):
|
||||
"""
|
||||
Process and tokenize lyrics in a dataset item
|
||||
|
||||
Args:
|
||||
item: Dataset item containing lyrics
|
||||
debug: Whether to print debug information
|
||||
|
||||
Returns:
|
||||
dict: Updated item with tokenized lyrics
|
||||
"""
|
||||
norm_lyrics = item["norm_lyrics"]
|
||||
|
||||
# Filter out prompts that match pattern "write a .* song that genre is"
|
||||
pattern = r"write a .* song that genre is"
|
||||
if re.search(pattern, norm_lyrics):
|
||||
norm_lyrics = ""
|
||||
item["lyric_token_idx"] = [0]
|
||||
item["norm_lyrics"] = norm_lyrics
|
||||
return item
|
||||
|
||||
key = item["keys"]
|
||||
|
||||
# Handle empty lyrics
|
||||
if not item["norm_lyrics"].strip():
|
||||
item["lyric_token_idx"] = [0]
|
||||
return item
|
||||
|
||||
# Tokenize lyrics
|
||||
item["lyric_token_idx"] = self.tokenize_lyrics(norm_lyrics, debug, key)
|
||||
return item
|
||||
|
||||
def get_speaker_emb_file(self, speaker_emb_path):
|
||||
"""
|
||||
Load speaker embedding file
|
||||
|
||||
Args:
|
||||
speaker_emb_path: Path to speaker embedding file
|
||||
|
||||
Returns:
|
||||
torch.Tensor or None: Speaker embedding
|
||||
"""
|
||||
data = None
|
||||
try:
|
||||
data = torch.load(speaker_emb_path, map_location="cpu")
|
||||
except Exception:
|
||||
pass
|
||||
return data
|
||||
|
||||
def get_audio(self, item):
|
||||
"""
|
||||
Load and preprocess audio file
|
||||
|
||||
Args:
|
||||
item: Dataset item containing filename
|
||||
|
||||
Returns:
|
||||
torch.Tensor or None: Processed audio tensor
|
||||
"""
|
||||
filename = item["filename"]
|
||||
sr = 48000
|
||||
try:
|
||||
audio, sr = torchaudio.load(filename)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load audio {item}: {e}")
|
||||
return None
|
||||
|
||||
if audio is None:
|
||||
logger.error(f"Failed to load audio {item}")
|
||||
return None
|
||||
|
||||
# Convert mono to stereo if needed
|
||||
if audio.shape[0] == 1:
|
||||
audio = torch.cat([audio, audio], dim=0)
|
||||
|
||||
# Take first two channels if more than stereo
|
||||
audio = audio[:2]
|
||||
|
||||
# Resample if needed
|
||||
if sr != 48000:
|
||||
audio = torchaudio.transforms.Resample(sr, 48000)(audio)
|
||||
|
||||
# Clip values to [-1.0, 1.0]
|
||||
audio = torch.clamp(audio, -1.0, 1.0)
|
||||
|
||||
# Pad to minimum 3 seconds if needed
|
||||
if audio.shape[-1] < 48000 * 3:
|
||||
audio = torch.nn.functional.pad(audio, (0, 48000 * 3 - audio.shape[-1]), 'constant', 0)
|
||||
|
||||
# Check if audio is silent
|
||||
if is_silent_audio(audio):
|
||||
logger.error(f"Silent audio {item}")
|
||||
return None
|
||||
|
||||
return audio
|
||||
|
||||
def process(self, item):
|
||||
"""
|
||||
Process a dataset item into model-ready features
|
||||
|
||||
Args:
|
||||
item: Dataset item
|
||||
|
||||
Returns:
|
||||
list: List of processed examples
|
||||
"""
|
||||
# Get audio
|
||||
audio = self.get_audio(item)
|
||||
if audio is None:
|
||||
return []
|
||||
|
||||
music_wavs = audio
|
||||
|
||||
# Get speaker embedding
|
||||
key = item["keys"]
|
||||
speaker_emb_path = item.get("speaker_emb_path")
|
||||
if not speaker_emb_path:
|
||||
speaker_emb = self.get_speaker_emb_file(speaker_emb_path)
|
||||
else:
|
||||
speaker_emb = torch.zeros(512)
|
||||
|
||||
# Process prompt/tags
|
||||
prompt = item["tags"]
|
||||
if len(prompt) == 0:
|
||||
prompt = ["music"]
|
||||
|
||||
# Shuffle tags and join with commas
|
||||
random.shuffle(prompt)
|
||||
prompt = ", ".join(prompt)
|
||||
|
||||
# Handle recaption data if available
|
||||
recaption = item.get("recaption", {})
|
||||
valid_recaption = []
|
||||
for k, v in recaption.items():
|
||||
if isinstance(v, str) and len(v) > 0:
|
||||
valid_recaption.append(v)
|
||||
|
||||
# Add original prompt to recaption options and randomly select one
|
||||
valid_recaption.append(prompt)
|
||||
prompt = random.choice(valid_recaption)
|
||||
prompt = prompt[:256] # Limit prompt length
|
||||
|
||||
# Process lyrics
|
||||
lyric_token_idx = item["lyric_token_idx"]
|
||||
lyric_token_idx = torch.tensor(lyric_token_idx).long()
|
||||
lyric_token_idx = lyric_token_idx[:4096] # Limit lyric context length
|
||||
lyric_mask = torch.ones(len(lyric_token_idx))
|
||||
|
||||
# Create lyric chunks for display
|
||||
candidate_lyric_chunk = []
|
||||
lyrics = item["norm_lyrics"]
|
||||
lyrics_lines = lyrics.split("\n")
|
||||
for lyric_line in lyrics_lines:
|
||||
candidate_lyric_chunk.append({
|
||||
"lyric": lyric_line,
|
||||
})
|
||||
|
||||
# Limit audio length
|
||||
longest_length = 24 * 10 * 48000 # 240 seconds
|
||||
music_wavs = music_wavs[:, :longest_length]
|
||||
vocal_wavs = torch.zeros_like(music_wavs)
|
||||
wav_len = music_wavs.shape[-1]
|
||||
|
||||
# Create example dictionary
|
||||
example = {
|
||||
"key": key,
|
||||
"vocal_wav": vocal_wavs,
|
||||
"target_wav": music_wavs,
|
||||
"wav_length": wav_len,
|
||||
"prompt": prompt,
|
||||
"speaker_emb": speaker_emb,
|
||||
"lyric_token_id": lyric_token_idx,
|
||||
"lyric_mask": lyric_mask,
|
||||
"structured_tag": {"recaption": recaption},
|
||||
"candidate_lyric_chunk": candidate_lyric_chunk,
|
||||
}
|
||||
return [example]
|
||||
|
||||
def get_full_features(self, idx):
|
||||
"""
|
||||
Get full features for a dataset index
|
||||
|
||||
Args:
|
||||
idx: Dataset index
|
||||
|
||||
Returns:
|
||||
dict: Dictionary of features
|
||||
"""
|
||||
examples = {
|
||||
"keys": [],
|
||||
"target_wavs": [],
|
||||
"vocal_wavs": [],
|
||||
"wav_lengths": [],
|
||||
"structured_tags": [],
|
||||
"prompts": [],
|
||||
"speaker_embs": [],
|
||||
"lyric_token_ids": [],
|
||||
"lyric_masks": [],
|
||||
"candidate_lyric_chunks": [],
|
||||
}
|
||||
|
||||
item = self.pretrain_ds[idx]
|
||||
item["idx"] = idx
|
||||
item = self.tokenize_lyrics_map(item)
|
||||
features = self.process(item)
|
||||
|
||||
if features:
|
||||
for feature in features:
|
||||
for k, v in feature.items():
|
||||
# Handle key mapping more explicitly
|
||||
target_key = k + "s" # Default plural form
|
||||
|
||||
# Special case handling for keys that don't follow simple plural pattern
|
||||
if k == "key":
|
||||
target_key = "keys"
|
||||
elif k == "wav_length":
|
||||
target_key = "wav_lengths"
|
||||
elif k == "candidate_lyric_chunk":
|
||||
target_key = "candidate_lyric_chunks"
|
||||
|
||||
if v is not None and target_key in examples:
|
||||
examples[target_key].append(v)
|
||||
|
||||
return examples
|
||||
|
||||
def pack_batch(self, batch):
|
||||
"""
|
||||
Pack a batch of examples
|
||||
|
||||
Args:
|
||||
batch: List of examples
|
||||
|
||||
Returns:
|
||||
dict: Packed batch
|
||||
"""
|
||||
packed_batch = {}
|
||||
for item in batch:
|
||||
for k, v in item.items():
|
||||
if k not in packed_batch:
|
||||
packed_batch[k] = v
|
||||
continue
|
||||
packed_batch[k] += v
|
||||
return packed_batch
|
||||
|
||||
def collate_fn(self, batch):
|
||||
"""
|
||||
Collate function for DataLoader
|
||||
|
||||
Args:
|
||||
batch: List of examples
|
||||
|
||||
Returns:
|
||||
dict: Collated batch with padded tensors
|
||||
"""
|
||||
batch = self.pack_batch(batch)
|
||||
output = {}
|
||||
|
||||
for k, v in batch.items():
|
||||
if k in ["keys", "structured_tags", "prompts", "candidate_lyric_chunks"]:
|
||||
# Pass through lists without modification
|
||||
padded_input_list = v
|
||||
elif k in ["wav_lengths"]:
|
||||
# Convert to LongTensor
|
||||
padded_input_list = torch.LongTensor(v)
|
||||
elif k in ["src_wavs", "target_wavs", "vocal_wavs"]:
|
||||
# Pad audio to max length
|
||||
max_length = max(seq.shape[1] for seq in v)
|
||||
padded_input_list = torch.stack([
|
||||
torch.nn.functional.pad(seq, (0, max_length - seq.shape[1]), 'constant', 0)
|
||||
for seq in v
|
||||
])
|
||||
elif k in ["clap_conditions"]:
|
||||
# Pad time dimension of embeddings
|
||||
max_length = max(seq.shape[0] for seq in v)
|
||||
v = [
|
||||
torch.nn.functional.pad(seq, (0, 0, 0, max_length - seq.shape[0]), 'constant', 0)
|
||||
for seq in v
|
||||
]
|
||||
padded_input_list = torch.stack(v)
|
||||
elif k == "speaker_embs":
|
||||
# Stack speaker embeddings
|
||||
padded_input_list = torch.stack(v)
|
||||
elif k in ["chunk_masks", "clap_attention_masks", "lyric_token_ids", "lyric_masks"]:
|
||||
# Pad sequence tensors
|
||||
max_length = max(len(seq) for seq in v)
|
||||
padded_input_list = torch.stack([
|
||||
torch.nn.functional.pad(seq, (0, max_length - len(seq)), 'constant', 0)
|
||||
for seq in v
|
||||
])
|
||||
|
||||
output[k] = padded_input_list
|
||||
|
||||
return output
|
||||
|
||||
def __getitem__(self, idx):
|
||||
"""
|
||||
Get item at index with error handling
|
||||
|
||||
Args:
|
||||
idx: Dataset index
|
||||
|
||||
Returns:
|
||||
dict: Example features
|
||||
"""
|
||||
try:
|
||||
example = self.get_full_features(idx)
|
||||
if len(example["keys"]) == 0:
|
||||
raise Exception(f"Empty example {idx=}")
|
||||
return example
|
||||
except Exception as e:
|
||||
# Log error and try a different random index
|
||||
logger.error(f"Error in getting item {idx}: {e}")
|
||||
traceback.print_exc()
|
||||
new_idx = random.choice(range(len(self)))
|
||||
return self.__getitem__(new_idx)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage
|
||||
dataset = Text2MusicDataset()
|
||||
print(f"Dataset size: {len(dataset)}")
|
||||
item = dataset[0]
|
||||
print(item)
|
||||
|
||||
for k, v in item.items():
|
||||
if len(v) > 0 and isinstance(v[0], torch.Tensor):
|
||||
print(k, [v[i].shape for i in range(len(v))])
|
||||
else:
|
||||
print(k, v)
|
||||
|
||||
item2 = dataset[1]
|
||||
batch = dataset.collate_fn([item, item2])
|
||||
for k, v in batch.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
print(k, end=" ")
|
||||
print(v.shape, v.min(), v.max())
|
||||
else:
|
||||
print(k, v)
|
||||
@@ -0,0 +1,613 @@
|
||||
import gradio as gr
|
||||
import librosa
|
||||
|
||||
|
||||
TAG_DEFAULT = "funk, pop, soul, rock, melodic, guitar, drums, bass, keyboard, percussion, 105 BPM, energetic, upbeat, groovy, vibrant, dynamic"
|
||||
LYRIC_DEFAULT = """[verse]
|
||||
Neon lights they flicker bright
|
||||
City hums in dead of night
|
||||
Rhythms pulse through concrete veins
|
||||
Lost in echoes of refrains
|
||||
|
||||
[verse]
|
||||
Bassline groovin' in my chest
|
||||
Heartbeats match the city's zest
|
||||
Electric whispers fill the air
|
||||
Synthesized dreams everywhere
|
||||
|
||||
[chorus]
|
||||
Turn it up and let it flow
|
||||
Feel the fire let it grow
|
||||
In this rhythm we belong
|
||||
Hear the night sing out our song
|
||||
|
||||
[verse]
|
||||
Guitar strings they start to weep
|
||||
Wake the soul from silent sleep
|
||||
Every note a story told
|
||||
In this night we’re bold and gold
|
||||
|
||||
[bridge]
|
||||
Voices blend in harmony
|
||||
Lost in pure cacophony
|
||||
Timeless echoes timeless cries
|
||||
Soulful shouts beneath the skies
|
||||
|
||||
[verse]
|
||||
Keyboard dances on the keys
|
||||
Melodies on evening breeze
|
||||
Catch the tune and hold it tight
|
||||
In this moment we take flight
|
||||
"""
|
||||
|
||||
|
||||
|
||||
|
||||
def create_output_ui(task_name="Text2Music"):
|
||||
# For many consumer-grade GPU devices, only one batch can be run
|
||||
output_audio1 = gr.Audio(type="filepath", label=f"{task_name} Generated Audio 1")
|
||||
# output_audio2 = gr.Audio(type="filepath", label="Generated Audio 2")
|
||||
with gr.Accordion(f"{task_name} Parameters", open=False):
|
||||
input_params_json = gr.JSON(label=f"{task_name} Parameters")
|
||||
# outputs = [output_audio1, output_audio2]
|
||||
outputs = [output_audio1]
|
||||
return outputs, input_params_json
|
||||
|
||||
|
||||
def dump_func(*args):
|
||||
print(args)
|
||||
return []
|
||||
|
||||
|
||||
def create_text2music_ui(
|
||||
gr,
|
||||
text2music_process_func,
|
||||
sample_data_func=None,
|
||||
):
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
with gr.Row(equal_height=True):
|
||||
# add markdown, tags and lyrics examples are from ai music generation community
|
||||
audio_duration = gr.Slider(-1, 240.0, step=0.00001, value=-1, label="Audio Duration", interactive=True, info="-1 means random duration (30 ~ 240).", scale=9)
|
||||
sample_bnt = gr.Button("Sample", variant="primary", scale=1)
|
||||
|
||||
prompt = gr.Textbox(lines=2, label="Tags", max_lines=4, value=TAG_DEFAULT, info="Support tags, descriptions, and scene. Use commas to separate different tags.\ntags and lyrics examples are from ai music generation community")
|
||||
lyrics = gr.Textbox(lines=9, label="Lyrics", max_lines=13, value=LYRIC_DEFAULT, info="Support lyric structure tags like [verse], [chorus], and [bridge] to separate different parts of the lyrics.\nUse [instrumental] or [inst] to generate instrumental music. Not support genre structure tag in lyrics")
|
||||
|
||||
with gr.Accordion("Basic Settings", open=False):
|
||||
infer_step = gr.Slider(minimum=1, maximum=1000, step=1, value=27, label="Infer Steps", interactive=True)
|
||||
guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=15.0, label="Guidance Scale", interactive=True, info="When guidance_scale_lyric > 1 and guidance_scale_text > 1, the guidance scale will not be applied.")
|
||||
guidance_scale_text = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Text", interactive=True, info="Guidance scale for text condition. It can only apply to cfg. set guidance_scale_text=5.0, guidance_scale_lyric=1.5 for start")
|
||||
guidance_scale_lyric = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="Guidance Scale Lyric", interactive=True)
|
||||
|
||||
manual_seeds = gr.Textbox(label="manual seeds (default None)", placeholder="1,2,3,4", value=None, info="Seed for the generation")
|
||||
|
||||
with gr.Accordion("Advanced Settings", open=False):
|
||||
scheduler_type = gr.Radio(["euler", "heun"], value="euler", label="Scheduler Type", elem_id="scheduler_type", info="Scheduler type for the generation. euler is recommended. heun will take more time.")
|
||||
cfg_type = gr.Radio(["cfg", "apg", "cfg_star"], value="apg", label="CFG Type", elem_id="cfg_type", info="CFG type for the generation. apg is recommended. cfg and cfg_star are almost the same.")
|
||||
use_erg_tag = gr.Checkbox(label="use ERG for tag", value=True, info="Use Entropy Rectifying Guidance for tag. It will multiple a temperature to the attention to make a weaker tag condition and make better diversity.")
|
||||
use_erg_lyric = gr.Checkbox(label="use ERG for lyric", value=True, info="The same but apply to lyric encoder's attention.")
|
||||
use_erg_diffusion = gr.Checkbox(label="use ERG for diffusion", value=True, info="The same but apply to diffusion model's attention.")
|
||||
|
||||
omega_scale = gr.Slider(minimum=-100.0, maximum=100.0, step=0.1, value=10.0, label="Granularity Scale", interactive=True, info="Granularity scale for the generation. Higher values can reduce artifacts")
|
||||
|
||||
guidance_interval = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Interval", interactive=True, info="Guidance interval for the generation. 0.5 means only apply guidance in the middle steps (0.25 * infer_steps to 0.75 * infer_steps)")
|
||||
guidance_interval_decay = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="Guidance Interval Decay", interactive=True, info="Guidance interval decay for the generation. Guidance scale will decay from guidance_scale to min_guidance_scale in the interval. 0.0 means no decay.")
|
||||
min_guidance_scale = gr.Slider(minimum=0.0, maximum=200.0, step=0.1, value=3.0, label="Min Guidance Scale", interactive=True, info="Min guidance scale for guidance interval decay's end scale")
|
||||
oss_steps = gr.Textbox(label="OSS Steps", placeholder="16, 29, 52, 96, 129, 158, 172, 183, 189, 200", value=None, info="Optimal Steps for the generation. But not test well")
|
||||
|
||||
text2music_bnt = gr.Button("Generate", variant="primary")
|
||||
|
||||
with gr.Column():
|
||||
outputs, input_params_json = create_output_ui()
|
||||
with gr.Tab("retake"):
|
||||
retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
|
||||
retake_seeds = gr.Textbox(label="retake seeds (default None)", placeholder="", value=None)
|
||||
retake_bnt = gr.Button("Retake", variant="primary")
|
||||
retake_outputs, retake_input_params_json = create_output_ui("Retake")
|
||||
|
||||
def retake_process_func(json_data, retake_variance, retake_seeds):
|
||||
return text2music_process_func(
|
||||
json_data["audio_duration"],
|
||||
json_data["prompt"],
|
||||
json_data["lyrics"],
|
||||
json_data["infer_step"],
|
||||
json_data["guidance_scale"],
|
||||
json_data["scheduler_type"],
|
||||
json_data["cfg_type"],
|
||||
json_data["omega_scale"],
|
||||
", ".join(map(str, json_data["actual_seeds"])),
|
||||
json_data["guidance_interval"],
|
||||
json_data["guidance_interval_decay"],
|
||||
json_data["min_guidance_scale"],
|
||||
json_data["use_erg_tag"],
|
||||
json_data["use_erg_lyric"],
|
||||
json_data["use_erg_diffusion"],
|
||||
", ".join(map(str, json_data["oss_steps"])),
|
||||
json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
|
||||
json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
|
||||
retake_seeds=retake_seeds,
|
||||
retake_variance=retake_variance,
|
||||
task="retake",
|
||||
)
|
||||
|
||||
retake_bnt.click(
|
||||
fn=retake_process_func,
|
||||
inputs=[
|
||||
input_params_json,
|
||||
retake_variance,
|
||||
retake_seeds,
|
||||
],
|
||||
outputs=retake_outputs + [retake_input_params_json],
|
||||
)
|
||||
with gr.Tab("repainting"):
|
||||
retake_variance = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.2, label="variance")
|
||||
retake_seeds = gr.Textbox(label="repaint seeds (default None)", placeholder="", value=None)
|
||||
repaint_start = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Repaint Start Time", interactive=True)
|
||||
repaint_end = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Repaint End Time", interactive=True)
|
||||
repaint_source = gr.Radio(["text2music", "last_repaint", "upload"], value="text2music", label="Repaint Source", elem_id="repaint_source")
|
||||
|
||||
repaint_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="repaint_source_audio_upload")
|
||||
repaint_source.change(
|
||||
fn=lambda x: gr.update(visible=x == "upload", elem_id="repaint_source_audio_upload"),
|
||||
inputs=[repaint_source],
|
||||
outputs=[repaint_source_audio_upload],
|
||||
)
|
||||
|
||||
repaint_bnt = gr.Button("Repaint", variant="primary")
|
||||
repaint_outputs, repaint_input_params_json = create_output_ui("Repaint")
|
||||
|
||||
def repaint_process_func(
|
||||
text2music_json_data,
|
||||
repaint_json_data,
|
||||
retake_variance,
|
||||
retake_seeds,
|
||||
repaint_start,
|
||||
repaint_end,
|
||||
repaint_source,
|
||||
repaint_source_audio_upload,
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
):
|
||||
if repaint_source == "upload":
|
||||
src_audio_path = repaint_source_audio_upload
|
||||
audio_duration = librosa.get_duration(filename=src_audio_path)
|
||||
json_data = {
|
||||
"audio_duration": audio_duration
|
||||
}
|
||||
elif repaint_source == "text2music":
|
||||
json_data = text2music_json_data
|
||||
src_audio_path = json_data["audio_path"]
|
||||
elif repaint_source == "last_repaint":
|
||||
json_data = repaint_json_data
|
||||
src_audio_path = json_data["audio_path"]
|
||||
|
||||
return text2music_process_func(
|
||||
json_data["audio_duration"],
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
retake_seeds=retake_seeds,
|
||||
retake_variance=retake_variance,
|
||||
task="repaint",
|
||||
repaint_start=repaint_start,
|
||||
repaint_end=repaint_end,
|
||||
src_audio_path=src_audio_path,
|
||||
)
|
||||
|
||||
repaint_bnt.click(
|
||||
fn=repaint_process_func,
|
||||
inputs=[
|
||||
input_params_json,
|
||||
repaint_input_params_json,
|
||||
retake_variance,
|
||||
retake_seeds,
|
||||
repaint_start,
|
||||
repaint_end,
|
||||
repaint_source,
|
||||
repaint_source_audio_upload,
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
],
|
||||
outputs=repaint_outputs + [repaint_input_params_json],
|
||||
)
|
||||
with gr.Tab("edit"):
|
||||
edit_prompt = gr.Textbox(lines=2, label="Edit Tags", max_lines=4)
|
||||
edit_lyrics = gr.Textbox(lines=9, label="Edit Lyrics", max_lines=13)
|
||||
retake_seeds = gr.Textbox(label="edit seeds (default None)", placeholder="", value=None)
|
||||
|
||||
edit_type = gr.Radio(["only_lyrics", "remix"], value="only_lyrics", label="Edit Type", elem_id="edit_type", info="`only_lyrics` will keep the whole song the same except lyrics difference. Make your diffrence smaller, e.g. one lyrc line change.\nremix can change the song melody and genre")
|
||||
edit_n_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.6, label="edit_n_min", interactive=True)
|
||||
edit_n_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="edit_n_max", interactive=True)
|
||||
|
||||
def edit_type_change_func(edit_type):
|
||||
if edit_type == "only_lyrics":
|
||||
n_min = 0.6
|
||||
n_max = 1.0
|
||||
elif edit_type == "remix":
|
||||
n_min = 0.2
|
||||
n_max = 0.4
|
||||
return n_min, n_max
|
||||
|
||||
edit_type.change(
|
||||
edit_type_change_func,
|
||||
inputs=[edit_type],
|
||||
outputs=[edit_n_min, edit_n_max]
|
||||
)
|
||||
|
||||
edit_source = gr.Radio(["text2music", "last_edit", "upload"], value="text2music", label="Edit Source", elem_id="edit_source")
|
||||
edit_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="edit_source_audio_upload")
|
||||
edit_source.change(
|
||||
fn=lambda x: gr.update(visible=x == "upload", elem_id="edit_source_audio_upload"),
|
||||
inputs=[edit_source],
|
||||
outputs=[edit_source_audio_upload],
|
||||
)
|
||||
|
||||
edit_bnt = gr.Button("Edit", variant="primary")
|
||||
edit_outputs, edit_input_params_json = create_output_ui("Edit")
|
||||
|
||||
def edit_process_func(
|
||||
text2music_json_data,
|
||||
edit_input_params_json,
|
||||
edit_source,
|
||||
edit_source_audio_upload,
|
||||
prompt,
|
||||
lyrics,
|
||||
edit_prompt,
|
||||
edit_lyrics,
|
||||
edit_n_min,
|
||||
edit_n_max,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
retake_seeds,
|
||||
):
|
||||
if edit_source == "upload":
|
||||
src_audio_path = edit_source_audio_upload
|
||||
audio_duration = librosa.get_duration(filename=src_audio_path)
|
||||
json_data = {
|
||||
"audio_duration": audio_duration
|
||||
}
|
||||
elif edit_source == "text2music":
|
||||
json_data = text2music_json_data
|
||||
src_audio_path = json_data["audio_path"]
|
||||
elif edit_source == "last_edit":
|
||||
json_data = edit_input_params_json
|
||||
src_audio_path = json_data["audio_path"]
|
||||
|
||||
if not edit_prompt:
|
||||
edit_prompt = prompt
|
||||
if not edit_lyrics:
|
||||
edit_lyrics = lyrics
|
||||
|
||||
return text2music_process_func(
|
||||
json_data["audio_duration"],
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
task="edit",
|
||||
src_audio_path=src_audio_path,
|
||||
edit_target_prompt=edit_prompt,
|
||||
edit_target_lyrics=edit_lyrics,
|
||||
edit_n_min=edit_n_min,
|
||||
edit_n_max=edit_n_max,
|
||||
retake_seeds=retake_seeds,
|
||||
)
|
||||
|
||||
edit_bnt.click(
|
||||
fn=edit_process_func,
|
||||
inputs=[
|
||||
input_params_json,
|
||||
edit_input_params_json,
|
||||
edit_source,
|
||||
edit_source_audio_upload,
|
||||
prompt,
|
||||
lyrics,
|
||||
edit_prompt,
|
||||
edit_lyrics,
|
||||
edit_n_min,
|
||||
edit_n_max,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
retake_seeds,
|
||||
],
|
||||
outputs=edit_outputs + [edit_input_params_json],
|
||||
)
|
||||
with gr.Tab("extend"):
|
||||
extend_seeds = gr.Textbox(label="extend seeds (default None)", placeholder="", value=None)
|
||||
left_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=0.0, label="Left Extend Length", interactive=True)
|
||||
right_extend_length = gr.Slider(minimum=0.0, maximum=240.0, step=0.01, value=30.0, label="Right Extend Length", interactive=True)
|
||||
extend_source = gr.Radio(["text2music", "last_extend", "upload"], value="text2music", label="Extend Source", elem_id="extend_source")
|
||||
|
||||
extend_source_audio_upload = gr.Audio(label="Upload Audio", type="filepath", visible=False, elem_id="extend_source_audio_upload")
|
||||
extend_source.change(
|
||||
fn=lambda x: gr.update(visible=x == "upload", elem_id="extend_source_audio_upload"),
|
||||
inputs=[extend_source],
|
||||
outputs=[extend_source_audio_upload],
|
||||
)
|
||||
|
||||
extend_bnt = gr.Button("Extend", variant="primary")
|
||||
extend_outputs, extend_input_params_json = create_output_ui("Extend")
|
||||
|
||||
def extend_process_func(
|
||||
text2music_json_data,
|
||||
extend_input_params_json,
|
||||
extend_seeds,
|
||||
left_extend_length,
|
||||
right_extend_length,
|
||||
extend_source,
|
||||
extend_source_audio_upload,
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
):
|
||||
if extend_source == "upload":
|
||||
src_audio_path = extend_source_audio_upload
|
||||
# get audio duration
|
||||
audio_duration = librosa.get_duration(filename=src_audio_path)
|
||||
json_data = {
|
||||
"audio_duration": audio_duration
|
||||
}
|
||||
elif extend_source == "text2music":
|
||||
json_data = text2music_json_data
|
||||
src_audio_path = json_data["audio_path"]
|
||||
elif extend_source == "last_extend":
|
||||
json_data = extend_input_params_json
|
||||
src_audio_path = json_data["audio_path"]
|
||||
|
||||
repaint_start = -left_extend_length
|
||||
repaint_end = json_data["audio_duration"] + right_extend_length
|
||||
return text2music_process_func(
|
||||
json_data["audio_duration"],
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
retake_seeds=extend_seeds,
|
||||
retake_variance=1.0,
|
||||
task="extend",
|
||||
repaint_start=repaint_start,
|
||||
repaint_end=repaint_end,
|
||||
src_audio_path=src_audio_path,
|
||||
)
|
||||
|
||||
extend_bnt.click(
|
||||
fn=extend_process_func,
|
||||
inputs=[
|
||||
input_params_json,
|
||||
extend_input_params_json,
|
||||
extend_seeds,
|
||||
left_extend_length,
|
||||
right_extend_length,
|
||||
extend_source,
|
||||
extend_source_audio_upload,
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
],
|
||||
outputs=extend_outputs + [extend_input_params_json],
|
||||
)
|
||||
|
||||
def sample_data():
|
||||
json_data = sample_data_func()
|
||||
return (
|
||||
json_data["audio_duration"],
|
||||
json_data["prompt"],
|
||||
json_data["lyrics"],
|
||||
json_data["infer_step"],
|
||||
json_data["guidance_scale"],
|
||||
json_data["scheduler_type"],
|
||||
json_data["cfg_type"],
|
||||
json_data["omega_scale"],
|
||||
", ".join(map(str, json_data["actual_seeds"])),
|
||||
json_data["guidance_interval"],
|
||||
json_data["guidance_interval_decay"],
|
||||
json_data["min_guidance_scale"],
|
||||
json_data["use_erg_tag"],
|
||||
json_data["use_erg_lyric"],
|
||||
json_data["use_erg_diffusion"],
|
||||
", ".join(map(str, json_data["oss_steps"])),
|
||||
json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
|
||||
json_data["guidance_scale_lyric"] if "guidance_scale_lyric" in json_data else 0.0,
|
||||
)
|
||||
|
||||
sample_bnt.click(
|
||||
sample_data,
|
||||
outputs=[
|
||||
audio_duration,
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
],
|
||||
)
|
||||
|
||||
text2music_bnt.click(
|
||||
fn=text2music_process_func,
|
||||
inputs=[
|
||||
audio_duration,
|
||||
prompt,
|
||||
lyrics,
|
||||
infer_step,
|
||||
guidance_scale,
|
||||
scheduler_type,
|
||||
cfg_type,
|
||||
omega_scale,
|
||||
manual_seeds,
|
||||
guidance_interval,
|
||||
guidance_interval_decay,
|
||||
min_guidance_scale,
|
||||
use_erg_tag,
|
||||
use_erg_lyric,
|
||||
use_erg_diffusion,
|
||||
oss_steps,
|
||||
guidance_scale_text,
|
||||
guidance_scale_lyric,
|
||||
], outputs=outputs + [input_params_json]
|
||||
)
|
||||
|
||||
|
||||
def create_main_demo_ui(
|
||||
text2music_process_func=dump_func,
|
||||
sample_data_func=dump_func,
|
||||
):
|
||||
with gr.Blocks(
|
||||
title="ACE-Step Model 1.0 DEMO",
|
||||
) as demo:
|
||||
gr.Markdown(
|
||||
"""
|
||||
<h1 style="text-align: center;">ACE-Step: A Step Towards Music Generation Foundation Model</h1>
|
||||
""")
|
||||
|
||||
with gr.Tab("text2music"):
|
||||
create_text2music_ui(
|
||||
gr=gr,
|
||||
text2music_process_func=text2music_process_func,
|
||||
sample_data_func=sample_data_func,
|
||||
)
|
||||
return demo
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo = create_main_demo_ui()
|
||||
demo.launch(
|
||||
server_name="0.0.0.0",
|
||||
server_port=7860,
|
||||
)
|
||||
Reference in New Issue
Block a user