stable-diffusion-webui-auto.../scripts/main.py

289 lines
11 KiB
Python

import gradio as gr
from pathlib import Path
from modules import script_callbacks, shared
import json
import os
from typing import List
from lru import LRU
from typing import Optional
# import asyncio
import concurrent.futures
from fastapi import FastAPI, Response, Body
from scripts.lan import lans
# from scripts.trans_google import GoogleTranslate
from scripts.trans_youdao import YoudaoTranslate
from scripts.trans_tp import TpTranslate,init_transers
transers = ['free_youdao_zh']
init_transers(transers)
transer = None
# Webui root path
ROOT_DIR = Path().absolute()
trans_file = os.path.join(ROOT_DIR, "localizations/my.json")
current_path = os.path.dirname(os.path.abspath(__file__))
config_file = os.path.join(current_path, "config.json")
trans_succ = 0
trans_config = None
textbox_info = None
end_trans_btn = None
m_temp_trans_list = None
m_temp_transed_num = 0
trans_data = None
m_trans_dict = LRU(1000)
def read_or_create_json_file(filename):
if not os.path.exists(filename):
with open(filename, 'w', encoding='utf-8') as f:
json.dump({}, f)
with open(filename, 'r', encoding='utf-8') as f:
data = json.load(f)
return data
def update_json_key(filename, data, key, value):
data[key] = value
with open(filename, 'w', encoding='utf-8') as f:
json.dump(data, f, indent=4)
# print(f'{filename} update: {key}={value}')
def update_json_file(filename,data):
# if not data:
# print('data is None')
# return
# 保存 JSON 文件
with open(filename, 'w', encoding='utf-8') as f:
json.dump(data, f, indent=4)
# print(f'{filename} update all')
def check_transer():
global transer
cur_transer = trans_config['transer']
# if cur_transer == 'free_google':
# transer = GoogleTranslate()
if cur_transer == 'free_youdao_zh':
transer = YoudaoTranslate()
else:
if cur_transer.startswith("tp_"):
tp_transer = cur_transer[3:]
transer = TpTranslate(tp_transer)
def init():
global trans_config
global trans_data
global trans_succ
trans_data = read_or_create_json_file(trans_file)
trans_succ = len(list(trans_data))
trans_config = read_or_create_json_file(config_file)
if 'show_en_enabled' not in trans_config:
# trans_config['auto_language_enabled'] = False
trans_config['show_en_enabled'] = True
trans_config['to_lan'] = 'zh-CN'
trans_config['transer'] = 'tp_bing'
trans_config['trans_succ'] = trans_succ
check_transer()
init()
#--------------------------------ui---------------------------------------
def update_info():
textbox_info.update(f'Auto Translate Num:{trans_succ}')
def auto_remove_trans():
global trans_data
global trans_succ
print('auto_remove_trans', trans_succ)
if trans_succ==0:
return
trans_succ = 0
trans_data = {}
update_json_file(trans_file,trans_data)
textbox_info.update(f'Auto Translate Num:{trans_succ}')
def on_ui_tabs():
global textbox_info
global end_trans_btn
with gr.Blocks(analytics_enabled=False) as translate_interface:
gr.HTML("<p style=\"margin-bottom:0.75em\">The untranslated characters will be translated automatically and will not affect the old translations. Use the function in the lower right corner to easily check and quickly modify the current translation.1,Save the setting;2,Click start button;3,Reload your browser.</p>")
with gr.Row():
with gr.Column(scale=80):
textbox_info = gr.Textbox(label='Translated Status', value=f'Auto Translate Num:{trans_succ}',readonly=True, elem_id="auto_textbox_info")
auto_trans_btn = gr.Button(value="Start Auto Translate", variant='primary', elem_id="auto_trans_btn")
with gr.Row():
with gr.Column(scale=45):
text_lan = gr.Textbox(label="Text", lines=3, value="", elem_id="auto_text_lan", placeholder="your select language")
with gr.Column():
with gr.Row():
auto_trans_to_en = gr.Button(value=f'-->', elem_id="auto_trans_to_en")
with gr.Row():
auto_trans_to_lan = gr.Button(value=f'<--', elem_id="auto_trans_to_lan")
with gr.Column(scale=45):
text_lan_translated = gr.Textbox(label="Translated Text", lines=3, value="", elem_id="text_lan_translated", placeholder="english")
gr.HTML("<hr />")
with gr.Row():
with gr.Column(scale=5):
with gr.Row():
tar_lang_drop = gr.Dropdown(label="To Language", choices=lans, value=trans_config['to_lan'], elem_id="auto_to_lang")
auto_langer_drop = gr.Dropdown(label="Select Translater", choices=transers, value=trans_config['transer'], elem_id="auto_langer_drop")
with gr.Row():
# auto_language_enabled = gr.Checkbox(label='enable auto trans ui language.(Need reload browser)', elem_id='auto_language_enabled', value=trans_config['auto_language_enabled'])
show_en_enabled = gr.Checkbox(label='display both english and target language', elem_id='show_en_enabled', value=trans_config['show_en_enabled'])
with gr.Row():
save_trans_setting_btn = gr.Button(value="Save Setting", elem_id='save_trans_setting_btn')
remove_trans_btn = gr.Button(value="Remove Auto Trans", elem_id='remove_trans_btn')
with gr.Column(scale=5, elem_id='auto_language_jsdiv'):
pass
end_trans_btn = gr.Button(value="Remove Auto Trans", elem_id='end_trans_btn', visible=False)
end_trans_btn.click(None,_js="auto_end_auto_trans", show_progress=False, inputs=[],outputs=[])
remove_trans_btn.click(fn=auto_remove_trans,show_progress=False, inputs=[],outputs=[])
tar_lang_drop.change(None,_js="auto_change_lan", inputs=[tar_lang_drop])
auto_langer_drop.change(None,_js="auto_change_transer", inputs=[auto_langer_drop])
# auto_language_enabled.change(None,_js="auto_change_language_enabled", inputs=[auto_language_enabled])
show_en_enabled.change(None,_js="auto_change_show_en_enabled", inputs=[show_en_enabled])
return [(translate_interface, "Auto Translate", "auto_translate")]
script_callbacks.on_ui_tabs(on_ui_tabs)
#--------------------------------api--------------------------------------
def transAPI(demo: gr.Blocks, app: FastAPI):
@app.get("/trans/config")
async def config():
res_data = {"error": 0}
res_data['config'] = trans_config
return res_data
@app.post("/trans/save_config")
async def auto_save_setting(data: dict = Body(...)):
res_data = {"error": 0}
# trans_config['auto_language_enabled'] = data['auto_language_enabled']
trans_config['show_en_enabled'] = data['show_en_enabled']
if trans_config['to_lan'] != data['to_lan']:
auto_remove_trans()
trans_config['to_lan'] = data['to_lan']
if trans_config['transer'] != data['transer']:
trans_config['transer'] = data['transer']
check_transer()
if 'appid' in data:
trans_config['appid'] = data['appid']
update_json_file(config_file, trans_config)
return res_data
@app.get("/trans/local_trans")
async def local_trans(text, save: Optional[bool] = True, en: Optional[bool] = False):
res_data = {"error": 0}
res_text = ''
global trans_succ
global m_trans_dict
global trans_data
if len(text) > 0:
print('text', text, save, en)
if text in trans_data:
res_text = trans_data[text]
elif text in m_trans_dict:
res_text = m_trans_dict[text]
else:
if en:
res_text = transer.translate(text, trans_config['to_lan'], 'en')
else:
res_text = transer.translate(text, 'en', trans_config['to_lan'])
print('local_trans', text, res_text)
if res_text != text and len(res_text)>0:
trans_succ = trans_succ + 1
m_trans_dict[text] = res_text
update_info()
if save:
update_json_key(trans_file, trans_data, text, res_text)
else:
res_data['error']=-1
res_data['txt'] = res_text
return res_data
def heavy_computation():
global trans_succ
global m_trans_dict
global m_temp_trans_list
global m_temp_transed_num
global trans_data
for text in m_temp_trans_list:
m_temp_transed_num = m_temp_transed_num + 1
if text in trans_data or text in m_trans_dict:
continue
res_text = transer.translate(text, 'en', trans_config['to_lan'])
print('local_trans', text, res_text)
if res_text != text and len(res_text)>0:
trans_succ = trans_succ + 1
m_trans_dict[text] = res_text
update_info()
update_json_key(trans_file, trans_data, text, res_text)
@app.post("/trans/local_trans_list")
async def local_trans_list(json_str: str = Body(...)):
res_data = {"error": 0}
items=json.loads(json_str)
global m_temp_trans_list
m_temp_trans_list = items
# print('local_trans_list items', items)
# asyncio.create_task(heavy_computation())
executor = concurrent.futures.ThreadPoolExecutor()
future = executor.submit(heavy_computation)
return res_data
@app.get("/trans/progress")
async def progress_get():
res_data = {"error": 0, "completed": False, "progress":0, "textinfo":"", "active":False, "trans_succ":trans_succ}
global m_temp_trans_list
global m_temp_transed_num
if m_temp_trans_list:
res_data['active'] = True
res_data['progress'] = m_temp_transed_num/len(m_temp_trans_list)
if m_temp_transed_num>0:
key = m_temp_trans_list[m_temp_transed_num-1]
if key in m_trans_dict:
res_data['textinfo'] = key+'->'+m_trans_dict[key]
if res_data['progress']==int(1):
res_data['completed'] = True
m_temp_trans_list = None
m_temp_transed_num = 0
return res_data
@app.post("/trans/progress")
async def progress():
return await progress_get()
@app.get("/trans/local_save_trans")
async def local_save_trans(text, tran):
global trans_data
res_data = {"error": 0}
if len(text) > 0:
update_json_key(trans_file,trans_data, text, tran)
else:
res_data = {"error": -1}
return res_data
script_callbacks.on_app_started(transAPI)