233 lines
8.7 KiB
Python
233 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
|
from textual.app import ComposeResult
|
|
from textual.widgets import (
|
|
Header,
|
|
Footer,
|
|
Label,
|
|
Input,
|
|
Select,
|
|
Button,
|
|
Markdown,
|
|
Static,
|
|
ProgressBar,
|
|
)
|
|
from textual.containers import Container, Horizontal, Center
|
|
from textual.containers import Container
|
|
from textual.screen import Screen
|
|
import pathlib
|
|
|
|
import heurams.kernel.particles as pt
|
|
import heurams.services.hasher as hasher
|
|
from heurams.context import *
|
|
from textual.worker import Worker, get_current_worker
|
|
|
|
class PrecachingScreen(Screen):
|
|
"""预缓存音频文件屏幕
|
|
|
|
缓存记忆单元音频文件, 全部(默认) 或部分记忆单元(可选参数传入)
|
|
|
|
Args:
|
|
nucleons (list): 可选列表, 仅包含 Nucleon 对象
|
|
desc (list): 可选字符串, 包含对此次调用的文字描述
|
|
"""
|
|
BINDINGS = [("q", "go_back", "返回")]
|
|
|
|
def __init__(self, nucleons: list = [], desc: str = ""):
|
|
super().__init__(name=None, id=None, classes=None)
|
|
self.nucleons = nucleons
|
|
self.is_precaching = False
|
|
self.current_file = ""
|
|
self.current_item = ""
|
|
self.progress = 0
|
|
self.total = len(nucleons)
|
|
self.processed = 0
|
|
self.precache_worker = None
|
|
self.cancel_flag = 0
|
|
self.desc = desc
|
|
for i in nucleons:
|
|
i: pt.Nucleon
|
|
i.do_eval()
|
|
#print("完成 EVAL")
|
|
|
|
def compose(self) -> ComposeResult:
|
|
yield Header(show_clock=True)
|
|
with Container(id="precache_container"):
|
|
yield Label("[b]音频预缓存[/b]", classes="title-label")
|
|
|
|
if self.nucleons:
|
|
yield Static(f"目标单元归属: [b]{self.desc}[/b]", classes="target-info")
|
|
yield Static(f"单元数量: {len(self.nucleons)}", classes="target-info")
|
|
else:
|
|
yield Static("目标: 所有单元", classes="target-info")
|
|
|
|
yield Static(id="status", classes="status-info")
|
|
yield Static(id="current_item", classes="current-item")
|
|
yield ProgressBar(total=100, show_eta=False, id="progress_bar")
|
|
|
|
with Horizontal(classes="button-group"):
|
|
if not self.is_precaching:
|
|
yield Button("开始预缓存", id="start_precache", variant="primary")
|
|
else:
|
|
yield Button("取消预缓存", id="cancel_precache", variant="error")
|
|
yield Button("清空缓存", id="clear_cache", variant="warning")
|
|
yield Button("返回", id="go_back", variant="default")
|
|
|
|
yield Static("若您离开此界面, 未完成的缓存进程会自动停止.")
|
|
yield Static("缓存程序支持 \"断点续传\".")
|
|
|
|
yield Footer()
|
|
|
|
def on_mount(self):
|
|
"""挂载时初始化状态"""
|
|
self.update_status("就绪", "等待开始...")
|
|
|
|
def update_status(self, status, current_item="", progress=None):
|
|
"""更新状态显示"""
|
|
status_widget = self.query_one("#status", Static)
|
|
item_widget = self.query_one("#current_item", Static)
|
|
progress_bar = self.query_one("#progress_bar", ProgressBar)
|
|
|
|
status_widget.update(f"状态: {status}")
|
|
item_widget.update(f"当前项目: {current_item}" if current_item else "")
|
|
|
|
if progress is not None:
|
|
progress_bar.progress = progress
|
|
progress_bar.advance(0) # 刷新显示
|
|
|
|
def precache_by_text(self, text: str):
|
|
"""预缓存单段文本的音频"""
|
|
from heurams.context import rootdir, workdir, config_var
|
|
cache_dir = pathlib.Path(config_var.get()["paths"]["cache_dir"])
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
cache_file = cache_dir / f"{hasher.get_md5(text)}.wav"
|
|
if not cache_file.exists():
|
|
try: # TODO: 调用模块消除tts耦合
|
|
import edge_tts as tts
|
|
communicate = tts.Communicate(text, "zh-CN-XiaoxiaoNeural")
|
|
communicate.save_sync(str(cache_file))
|
|
return 1
|
|
except Exception as e:
|
|
print(f"预缓存失败 '{text}': {e}")
|
|
return 0
|
|
return 1
|
|
|
|
def precache_by_nucleon(self, nucleon: pt.Nucleon):
|
|
"""依据 Nucleon 缓存"""
|
|
#print(nucleon.metadata['formation']['tts_text'])
|
|
ret = self.precache_by_text(nucleon.metadata['formation']['tts_text'])
|
|
return ret
|
|
#print(f"TTS 缓存: {nucleon.metadata['formation']['tts_text']}")
|
|
|
|
def precache_by_list(self, nucleons: list):
|
|
"""依据 Nucleons 列表缓存"""
|
|
for idx, nucleon in enumerate(nucleons):
|
|
#print(f"PROC: {nucleon}")
|
|
worker = get_current_worker()
|
|
if worker and worker.is_cancelled: # 函数在worker中执行且已被取消
|
|
return False
|
|
text = nucleon.metadata['formation']['tts_text']
|
|
#self.current_item = text[:30] + "..." if len(text) > 50 else text
|
|
#print(text)
|
|
self.processed += 1
|
|
#print(self.processed)
|
|
#print(self.total)
|
|
progress = int((self.processed / self.total) * 100) if self.total > 0 else 0
|
|
#print(progress)
|
|
self.update_status(
|
|
f"正处理 ({idx + 1}/{len(nucleons)})",
|
|
text,
|
|
progress
|
|
)
|
|
ret = self.precache_by_nucleon(nucleon)
|
|
if not ret:
|
|
self.update_status(
|
|
"出错",
|
|
f"处理失败, 跳过: {self.current_item}",
|
|
)
|
|
import time
|
|
time.sleep(1)
|
|
if self.cancel_flag:
|
|
worker.cancel()
|
|
self.cancel_flag = 0
|
|
return False
|
|
return True
|
|
|
|
def precache_by_nucleons(self):
|
|
#print("开始缓存")
|
|
ret = self.precache_by_list(self.nucleons)
|
|
#print(f"返回 {ret}")
|
|
return ret
|
|
|
|
def precache_by_filepath(self, path: pathlib.Path):
|
|
"""预缓存单个文件的所有内容"""
|
|
lst = list()
|
|
for i in pt.load_nucleon(path):
|
|
lst.append(i[0])
|
|
return self.precache_by_list(lst)
|
|
|
|
|
|
def precache_all_files(self):
|
|
"""预缓存所有文件"""
|
|
from heurams.context import rootdir, workdir, config_var
|
|
nucleon_path = pathlib.Path(config_var.get()["paths"]["nucleon_dir"])
|
|
nucleon_files = [f for f in nucleon_path.iterdir() if f.suffix == ".toml"] # TODO: 解耦合
|
|
|
|
# 计算总项目数
|
|
self.total = 0
|
|
nu = list()
|
|
for file in nucleon_files:
|
|
try:
|
|
for i in pt.load_nucleon(file):
|
|
nu.append(i[0])
|
|
except:
|
|
continue
|
|
self.total = len(nu)
|
|
for i in nu:
|
|
i: pt.Nucleon
|
|
i.do_eval()
|
|
return self.precache_by_list(nu)
|
|
|
|
def on_button_pressed(self, event: Button.Pressed) -> None:
|
|
event.stop()
|
|
if event.button.id == "start_precache" and not self.is_precaching:
|
|
# 开始预缓存
|
|
if self.nucleons:
|
|
self.precache_worker = self.run_worker(self.precache_by_nucleons, thread=True, exclusive=True, exit_on_error=True)
|
|
else:
|
|
self.precache_worker = self.run_worker(self.precache_all_files, thread=True, exclusive=True, exit_on_error=True)
|
|
|
|
elif event.button.id == "cancel_precache" and self.is_precaching:
|
|
# 取消预缓存
|
|
if self.precache_worker:
|
|
self.precache_worker.cancel()
|
|
self.is_precaching = False
|
|
self.processed = 0
|
|
self.progress = 0
|
|
self.update_status("已取消", "预缓存操作被用户取消", 0)
|
|
|
|
elif event.button.id == "clear_cache":
|
|
# 清空缓存
|
|
try:
|
|
import shutil
|
|
from heurams.context import rootdir, workdir, config_var
|
|
shutil.rmtree(f"{config_var.get()["paths"]["cache_dir"]}", ignore_errors=True)
|
|
self.update_status("已清空", "音频缓存已清空", 0)
|
|
except Exception as e:
|
|
self.update_status("错误", f"清空缓存失败: {e}")
|
|
self.cancel_flag = 1
|
|
self.processed = 0
|
|
self.progress = 0
|
|
|
|
elif event.button.id == "go_back":
|
|
self.action_go_back()
|
|
|
|
def action_go_back(self):
|
|
if self.is_precaching and self.precache_worker:
|
|
self.precache_worker.cancel()
|
|
self.app.pop_screen()
|
|
|
|
def action_quit_app(self):
|
|
if self.is_precaching and self.precache_worker:
|
|
self.precache_worker.cancel()
|
|
self.app.exit()
|