494 lines
24 KiB
Python
494 lines
24 KiB
Python
import gradio as gr
|
|
from datetime import datetime
|
|
import json
|
|
import re
|
|
import traceback
|
|
import os
|
|
|
|
from locales.i18n import t
|
|
from services.novel_generator import OutlineParser
|
|
from services.project_manager import ProjectManager
|
|
from services.genre_manager import GenreManager
|
|
from services.sub_genre_manager import SubGenreManager
|
|
from services.style_manager import StyleManager
|
|
from core.state import app_state
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ==================== DRAFT PERSISTENCE ====================
|
|
DRAFT_PATH = os.path.join(
|
|
os.path.dirname(os.path.dirname(__file__)), "data", "create_draft.json"
|
|
)
|
|
|
|
DRAFT_FIELDS = [
|
|
"genre", "sub_genres", "style", "title", "suggest_title_prompt",
|
|
"num_main", "num_sub", "character", "suggest_char_prompt",
|
|
"world", "suggest_world_prompt", "plot", "suggest_plot_prompt",
|
|
"custom_outline", "total_chapters", "memory_type", "memory_chapters",
|
|
"use_reflection", "outline",
|
|
]
|
|
|
|
|
|
def _save_create_draft(d: dict):
|
|
try:
|
|
os.makedirs(os.path.dirname(DRAFT_PATH), exist_ok=True)
|
|
with open(DRAFT_PATH, "w", encoding="utf-8") as f:
|
|
json.dump(d, f, ensure_ascii=False)
|
|
except Exception as e:
|
|
logger.error(f"Save create draft failed: {e}")
|
|
|
|
|
|
def _load_create_draft() -> dict:
|
|
try:
|
|
if os.path.exists(DRAFT_PATH):
|
|
with open(DRAFT_PATH, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"Load create draft failed: {e}")
|
|
return {}
|
|
|
|
|
|
def _clear_create_draft():
|
|
try:
|
|
if os.path.exists(DRAFT_PATH):
|
|
os.remove(DRAFT_PATH)
|
|
except Exception as e:
|
|
logger.error(f"Clear create draft failed: {e}")
|
|
|
|
|
|
def build_create_tab(demo=None):
|
|
"""Xây dựng giao diện Tab Sáng tác từ đầu"""
|
|
with gr.Tab(t("tabs.create")):
|
|
gr.Markdown(f"### 🪄 {t('create.header')}")
|
|
|
|
with gr.Accordion("📌 1. Thông tin cơ bản", open=True):
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
genre_choices = GenreManager.get_genre_names()
|
|
genre_dropdown = gr.Dropdown(
|
|
choices=genre_choices,
|
|
value=genre_choices[0] if genre_choices else None,
|
|
label=t("create.genre_label"),
|
|
interactive=True
|
|
)
|
|
sub_genre_choices = SubGenreManager.get_sub_genre_names()
|
|
sub_genre_dropdown = gr.Dropdown(
|
|
choices=sub_genre_choices,
|
|
label=t("create.sub_genres_label"),
|
|
multiselect=True,
|
|
interactive=True
|
|
)
|
|
|
|
style_choices = StyleManager.get_style_names()
|
|
create_style_dropdown = gr.Dropdown(
|
|
choices=style_choices,
|
|
label="Phong cách viết",
|
|
value=style_choices[0] if style_choices else None,
|
|
interactive=True
|
|
)
|
|
|
|
with gr.Column(scale=2):
|
|
title_input = gr.Textbox(
|
|
label=t("create.novel_title"),
|
|
placeholder=t("create.novel_title_default"),
|
|
interactive=True
|
|
)
|
|
with gr.Group():
|
|
suggested_titles_radio = gr.Radio(
|
|
label=t("create.suggested_titles_radio") if t("create.suggested_titles_radio") != "create.suggested_titles_radio" else "Gợi ý tên truyện",
|
|
choices=[],
|
|
visible=False,
|
|
interactive=True
|
|
)
|
|
with gr.Row():
|
|
suggest_title_prompt = gr.Textbox(
|
|
show_label=False,
|
|
placeholder=t("create.custom_prompt_placeholder"),
|
|
lines=1, interactive=True, scale=4
|
|
)
|
|
suggest_title_btn = gr.Button(t("create.suggest_title_btn"), variant="secondary", size="sm", scale=1)
|
|
suggest_title_status = gr.Textbox(show_label=False, interactive=False, visible=False)
|
|
|
|
with gr.Accordion("🎭 2. Thiết lập chi tiết", open=False):
|
|
with gr.Row():
|
|
with gr.Column():
|
|
character_input = gr.Textbox(
|
|
label=t("create.char_setting"),
|
|
placeholder=t("create.char_setting_placeholder"),
|
|
lines=5, interactive=True
|
|
)
|
|
with gr.Group():
|
|
with gr.Row():
|
|
num_main_chars = gr.Number(label=t("create.num_main_chars"), value=2, minimum=1, maximum=10, step=1, scale=1)
|
|
num_sub_chars = gr.Number(label=t("create.num_sub_chars"), value=3, minimum=0, maximum=20, step=1, scale=1)
|
|
with gr.Row():
|
|
suggest_char_prompt = gr.Textbox(show_label=False, placeholder=t("create.custom_prompt_placeholder"), lines=1, interactive=True, scale=4)
|
|
suggest_char_btn = gr.Button(t("create.suggest_btn"), variant="secondary", size="sm", scale=1, interactive=True)
|
|
suggest_char_status = gr.Textbox(show_label=False, interactive=False, visible=False)
|
|
|
|
with gr.Column():
|
|
world_input = gr.Textbox(
|
|
label=t("create.world_setting"),
|
|
placeholder=t("create.world_setting_placeholder"),
|
|
lines=5, interactive=True
|
|
)
|
|
with gr.Group():
|
|
with gr.Row():
|
|
suggest_world_prompt = gr.Textbox(show_label=False, placeholder=t("create.custom_prompt_placeholder"), lines=1, interactive=True, scale=4)
|
|
suggest_world_btn = gr.Button(t("create.suggest_btn"), variant="secondary", size="sm", scale=1, interactive=True)
|
|
suggest_world_status = gr.Textbox(show_label=False, interactive=False, visible=False)
|
|
|
|
plot_input = gr.Textbox(
|
|
label=t("create.plot_idea"),
|
|
placeholder=t("create.plot_idea_placeholder"),
|
|
lines=3, interactive=True
|
|
)
|
|
with gr.Group():
|
|
with gr.Row():
|
|
suggest_plot_prompt = gr.Textbox(show_label=False, placeholder=t("create.custom_prompt_placeholder"), lines=1, interactive=True, scale=4)
|
|
suggest_plot_btn = gr.Button(t("create.suggest_plot_btn"), variant="secondary", size="sm", scale=1, interactive=True)
|
|
suggest_plot_status = gr.Textbox(show_label=False, interactive=False, visible=False)
|
|
|
|
with gr.Accordion("📝 3. Dàn ý truyện", open=False):
|
|
custom_outline_prompt = gr.Textbox(
|
|
label="Yêu cầu riêng/Chỉ dẫn bổ sung cho Dàn ý (Tùy chọn)",
|
|
placeholder="VD: Không có nữ chính, kết cục mở, nvc là người tu ma nhưng tính cách hài hước...",
|
|
lines=2, interactive=True
|
|
)
|
|
with gr.Row():
|
|
total_chapters = gr.Number(
|
|
label=t("create.chapter_count"), value=20, minimum=1, maximum=200, step=1, scale=1, interactive=True
|
|
)
|
|
generate_outline_btn = gr.Button(t("create.gen_outline_btn"), variant="primary", size="lg", scale=3, interactive=True)
|
|
|
|
outline_output = gr.Textbox(label=t("create.outline_display"), lines=15, interactive=True)
|
|
outline_status = gr.Textbox(label=t("create.gen_status"), interactive=False)
|
|
|
|
with gr.Accordion("🚀 4. Sáng tác tự động", open=False):
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
memory_type = gr.Radio(label="Cách nhớ ngữ cảnh", choices=["Toàn văn", "Tóm tắt"], value="Toàn văn")
|
|
with gr.Column(scale=1):
|
|
memory_chapters = gr.Number(label="Số chương ghi nhớ", value=3, minimum=1, maximum=20, step=1)
|
|
with gr.Column(scale=1):
|
|
use_reflection_checkbox = gr.Checkbox(label="Bật chế độ Tự kiểm duyệt (Self-Reflection)", value=False, info="AI tự đọc lại và sửa lỗi nháp. Tốn gấp đôi thời gian & Token.")
|
|
with gr.Row():
|
|
auto_generate_btn = gr.Button(t("create.start_gen_btn"), variant="primary", size="lg", interactive=True)
|
|
stop_btn = gr.Button(t("create.pause_gen_btn"), variant="stop", size="lg", interactive=True)
|
|
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
generation_progress = gr.Textbox(label=t("create.gen_status"), lines=15, interactive=False)
|
|
with gr.Column(scale=2):
|
|
chapter_selector = gr.Dropdown(label="Danh sách chương đã tạo", choices=[], interactive=True, allow_custom_value=True)
|
|
chapter_content_display = gr.Textbox(label="Nội dung chương", lines=15, interactive=False)
|
|
|
|
def on_suggest_title(genre, sub_genres, custom_prompt):
|
|
yield gr.Radio(visible=False), gr.update(value="⏳ Đang gọi AI xử lý... Vui lòng chờ.", visible=True), gr.update(interactive=False)
|
|
try:
|
|
gen = app_state.get_generator()
|
|
content, msg = gen.suggest_title(genre, sub_genres, custom_prompt)
|
|
if content:
|
|
try:
|
|
match = re.search(r'\[\s*\{[\s\S]*\}\s*\]', content)
|
|
if match:
|
|
data = {"suggestions": json.loads(match.group(0))}
|
|
else:
|
|
clean_content = re.sub(r'```(?:json)?\s*', '', content)
|
|
clean_content = re.sub(r'```\s*', '', clean_content)
|
|
data = json.loads(clean_content)
|
|
|
|
if "suggestions" in data:
|
|
choices = [f"{item.get('title', '')} - {item.get('description', '')}" for item in data["suggestions"]]
|
|
yield gr.Radio(choices=choices, visible=True), gr.update(visible=False), gr.update(interactive=True)
|
|
return
|
|
except Exception as e:
|
|
logger.error(f"Failed to parse title JSON: {e}")
|
|
yield gr.Radio(choices=[content, f"Parse error: {str(e)}"], visible=True), gr.update(visible=False), gr.update(interactive=True)
|
|
return
|
|
yield gr.Radio(choices=[msg], visible=True), gr.update(visible=False), gr.update(interactive=True)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
yield gr.Radio(choices=[f"Lỗi: {str(e)}"], visible=True), gr.update(value=f"❌ Lỗi: {str(e)}", visible=True), gr.update(interactive=True)
|
|
|
|
def on_title_select(selected):
|
|
if selected and " - " in selected:
|
|
return selected.split(" - ", 1)[0]
|
|
return selected
|
|
|
|
def on_suggest_char(title, genre, sub_genres, custom_prompt, num_main, num_sub):
|
|
yield gr.update(), gr.update(value="⏳ Đang gọi AI xử lý... Vui lòng chờ.", visible=True), gr.update(interactive=False)
|
|
try:
|
|
gen = app_state.get_generator()
|
|
nm = int(num_main) if num_main else 2
|
|
ns = int(num_sub) if num_sub else 3
|
|
content, msg = gen.suggest_content("char", title, genre, sub_genres, "", "", custom_prompt, nm, ns)
|
|
if content:
|
|
yield gr.update(value=content), gr.update(visible=False), gr.update(interactive=True)
|
|
else:
|
|
yield gr.update(), gr.update(value=f"❌ {msg}", visible=True), gr.update(interactive=True)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
yield gr.update(), gr.update(value=f"❌ Lỗi: {str(e)}", visible=True), gr.update(interactive=True)
|
|
|
|
def on_suggest_world(title, genre, sub_genres, custom_prompt):
|
|
yield gr.update(), gr.update(value="⏳ Đang gọi AI xử lý... Vui lòng chờ.", visible=True), gr.update(interactive=False)
|
|
try:
|
|
gen = app_state.get_generator()
|
|
content, msg = gen.suggest_content("world", title, genre, sub_genres, "", "", custom_prompt)
|
|
if content:
|
|
yield gr.update(value=content), gr.update(visible=False), gr.update(interactive=True)
|
|
else:
|
|
yield gr.update(), gr.update(value=f"❌ {msg}", visible=True), gr.update(interactive=True)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
yield gr.update(), gr.update(value=f"❌ Lỗi: {str(e)}", visible=True), gr.update(interactive=True)
|
|
|
|
def on_suggest_plot(title, genre, sub_genres, char_setting, world_setting, custom_prompt):
|
|
yield gr.update(), gr.update(value="⏳ Đang gọi AI xử lý... Vui lòng chờ.", visible=True), gr.update(interactive=False)
|
|
try:
|
|
gen = app_state.get_generator()
|
|
content, msg = gen.suggest_content("plot", title, genre, sub_genres, char_setting, world_setting, custom_prompt)
|
|
if content:
|
|
yield gr.update(value=content), gr.update(visible=False), gr.update(interactive=True)
|
|
else:
|
|
yield gr.update(), gr.update(value=f"❌ {msg}", visible=True), gr.update(interactive=True)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
yield gr.update(), gr.update(value=f"❌ Lỗi: {str(e)}", visible=True), gr.update(interactive=True)
|
|
|
|
def on_generate_outline(title, genre, sub_genres, num_chapters, char_setting, world_setting, plot_idea, custom_outline, progress=gr.Progress()):
|
|
progress(0.1, desc="Đang gọi AI...")
|
|
gen = app_state.get_generator()
|
|
content, msg = gen.generate_outline(
|
|
title, genre, sub_genres or [],
|
|
int(num_chapters), char_setting, world_setting, plot_idea, custom_outline
|
|
)
|
|
return content, msg
|
|
|
|
def on_auto_generate(title, genre, sub_genres, char_setting, world_setting, plot_idea, outline_text, mem_type, mem_chaps, use_reflection, selected_style, progress=gr.Progress()):
|
|
"""Tự động tạo toàn bộ tiểu thuyết"""
|
|
_clear_create_draft()
|
|
gen = app_state.get_generator()
|
|
|
|
# Cập nhật style vào config hiện tại
|
|
if selected_style:
|
|
gen.config.generation.writing_style = selected_style
|
|
|
|
chapters, parse_msg = OutlineParser.parse(outline_text)
|
|
if not chapters:
|
|
yield f"❌ {parse_msg}", gr.update()
|
|
return
|
|
|
|
project, create_msg = ProjectManager.create_project(
|
|
title, genre, sub_genres or [],
|
|
char_setting, world_setting, plot_idea
|
|
)
|
|
if not project:
|
|
yield f"❌ {create_msg}", gr.update()
|
|
return
|
|
|
|
project.chapters = chapters
|
|
app_state.current_project = project
|
|
app_state.is_generating = True
|
|
app_state.stop_requested = False
|
|
|
|
# Save immediately to ensure chapter titles are stored even if generation is stopped early
|
|
ProjectManager.save_project(project)
|
|
|
|
results = [f"📋 Đã phân tích {len(chapters)} chương", f"💾 {create_msg}"]
|
|
yield "\n".join(results), gr.update(choices=[], value=None)
|
|
|
|
generated_chapters = []
|
|
|
|
for i, chapter in enumerate(chapters):
|
|
if app_state.stop_requested:
|
|
results.append("\n⚠️ Đã dừng sinh!")
|
|
yield "\n".join(results), gr.update(choices=generated_chapters)
|
|
break
|
|
|
|
results.append(f"\n✍️ Đang sinh Chương {chapter.num}: {chapter.title}...")
|
|
progress((i + 1) / len(chapters))
|
|
yield "\n".join(results), gr.update(choices=generated_chapters)
|
|
|
|
mem_ch = int(mem_chaps) if mem_chaps else 3
|
|
start_idx = max(0, i - mem_ch)
|
|
past_chapters = chapters[start_idx:i]
|
|
|
|
prev_content = ""
|
|
context_summary = ""
|
|
|
|
if mem_type == "Toàn văn":
|
|
prev_texts = [c.content for c in past_chapters if hasattr(c, 'content') and c.content]
|
|
prev_content = "\n\n".join(prev_texts)
|
|
prev_content = prev_content[-4000:]
|
|
else:
|
|
summaries = []
|
|
for c in past_chapters:
|
|
if not hasattr(c, 'content') or not c.content:
|
|
continue
|
|
if not hasattr(c, 'summary') or not c.summary:
|
|
summ, _ = gen.generate_chapter_summary(c.content, c.title)
|
|
c.summary = summ
|
|
if c.summary:
|
|
summaries.append(f"Chương {c.num} - {c.title}: {c.summary}")
|
|
|
|
if summaries:
|
|
context_summary = "\n".join(summaries)
|
|
|
|
if i > 0 and chapters[i-1].content:
|
|
prev_content = chapters[i-1].content[-1500:]
|
|
|
|
content, msg = gen.generate_chapter(
|
|
chapter_num=chapter.num, chapter_title=chapter.title,
|
|
chapter_desc=chapter.desc, novel_title=title,
|
|
character_setting=char_setting, world_setting=world_setting,
|
|
plot_idea=plot_idea, genre=genre, sub_genres=sub_genres,
|
|
previous_content=prev_content, context_summary=context_summary,
|
|
use_reflection=use_reflection
|
|
)
|
|
|
|
if content:
|
|
chapter.content = content
|
|
chapter.word_count = len(content)
|
|
chapter.generated_at = datetime.now().isoformat()
|
|
results.append(f"✅ Chương {chapter.num}: {len(content)} từ")
|
|
ProjectManager.save_project(project)
|
|
chapter_name = f"Chương {chapter.num}: {chapter.title}"
|
|
generated_chapters.append(chapter_name)
|
|
yield "\n".join(results), gr.update(choices=generated_chapters, value=chapter_name)
|
|
else:
|
|
results.append(f"❌ Chương {chapter.num}: {msg}")
|
|
yield "\n".join(results), gr.update(choices=generated_chapters)
|
|
|
|
app_state.is_generating = False
|
|
total_words = sum(ch.word_count for ch in chapters if ch.content)
|
|
results.append(f"\n🎉 Hoàn thành! Tổng: {total_words} từ")
|
|
yield "\n".join(results), gr.update(choices=generated_chapters)
|
|
|
|
def on_chapter_select(chapter_title):
|
|
if not app_state.current_project or not chapter_title:
|
|
return ""
|
|
for ch in app_state.current_project.chapters:
|
|
if f"Chương {ch.num}: {ch.title}" == chapter_title:
|
|
return ch.content or ""
|
|
return ""
|
|
|
|
def on_stop():
|
|
app_state.stop_requested = True
|
|
return "⏸️ Đang dừng..."
|
|
|
|
# ==================== DRAFT PERSISTENCE ====================
|
|
draft_inputs = [
|
|
genre_dropdown, sub_genre_dropdown, create_style_dropdown,
|
|
title_input, suggest_title_prompt, num_main_chars, num_sub_chars,
|
|
character_input, suggest_char_prompt, world_input,
|
|
suggest_world_prompt, plot_input, suggest_plot_prompt,
|
|
custom_outline_prompt, total_chapters, memory_type,
|
|
memory_chapters, use_reflection_checkbox, outline_output,
|
|
]
|
|
|
|
def on_draft_change(*values):
|
|
d = dict(zip(DRAFT_FIELDS, values))
|
|
_save_create_draft(d)
|
|
return
|
|
|
|
for _comp in draft_inputs:
|
|
_comp.change(fn=on_draft_change, inputs=draft_inputs, outputs=[])
|
|
|
|
def on_load_draft():
|
|
d = _load_create_draft()
|
|
title = d.get("title", "")
|
|
character = d.get("character", "")
|
|
world = d.get("world", "")
|
|
plot = d.get("plot", "")
|
|
outline = d.get("outline", "")
|
|
return (
|
|
d.get("genre"),
|
|
d.get("sub_genres"),
|
|
d.get("style"),
|
|
title,
|
|
d.get("suggest_title_prompt"),
|
|
d.get("num_main", 2),
|
|
d.get("num_sub", 3),
|
|
d.get("character", ""),
|
|
gr.update(interactive=True),
|
|
gr.update(interactive=True),
|
|
d.get("world", ""),
|
|
gr.update(interactive=True),
|
|
gr.update(interactive=True),
|
|
d.get("plot", ""),
|
|
gr.update(interactive=True),
|
|
gr.update(interactive=True),
|
|
d.get("custom_outline", ""),
|
|
d.get("total_chapters", 20),
|
|
gr.update(interactive=True),
|
|
d.get("outline", ""),
|
|
gr.update(interactive=True),
|
|
gr.update(interactive=True),
|
|
d.get("memory_type", "Toàn văn"),
|
|
d.get("memory_chapters", 3),
|
|
bool(d.get("use_reflection", False)),
|
|
)
|
|
|
|
if demo is not None:
|
|
demo.load(
|
|
fn=on_load_draft,
|
|
inputs=None,
|
|
outputs=[
|
|
genre_dropdown, sub_genre_dropdown, create_style_dropdown,
|
|
title_input, suggest_title_prompt, num_main_chars, num_sub_chars,
|
|
character_input, suggest_char_prompt, suggest_char_btn,
|
|
world_input, suggest_world_prompt, suggest_world_btn,
|
|
plot_input, suggest_plot_prompt, suggest_plot_btn,
|
|
custom_outline_prompt, total_chapters, generate_outline_btn,
|
|
outline_output, auto_generate_btn, stop_btn,
|
|
memory_type, memory_chapters, use_reflection_checkbox,
|
|
],
|
|
)
|
|
|
|
# Bind events
|
|
suggest_title_btn.click(
|
|
fn=on_suggest_title,
|
|
inputs=[genre_dropdown, sub_genre_dropdown, suggest_title_prompt],
|
|
outputs=[suggested_titles_radio, suggest_title_status, suggest_title_btn]
|
|
)
|
|
suggested_titles_radio.change(
|
|
fn=on_title_select,
|
|
inputs=[suggested_titles_radio],
|
|
outputs=[title_input]
|
|
)
|
|
suggest_char_btn.click(
|
|
fn=on_suggest_char,
|
|
inputs=[title_input, genre_dropdown, sub_genre_dropdown, suggest_char_prompt, num_main_chars, num_sub_chars],
|
|
outputs=[character_input, suggest_char_status, suggest_char_btn]
|
|
)
|
|
suggest_world_btn.click(
|
|
fn=on_suggest_world,
|
|
inputs=[title_input, genre_dropdown, sub_genre_dropdown, suggest_world_prompt],
|
|
outputs=[world_input, suggest_world_status, suggest_world_btn]
|
|
)
|
|
suggest_plot_btn.click(
|
|
fn=on_suggest_plot,
|
|
inputs=[title_input, genre_dropdown, sub_genre_dropdown, character_input, world_input, suggest_plot_prompt],
|
|
outputs=[plot_input, suggest_plot_status, suggest_plot_btn]
|
|
)
|
|
generate_outline_btn.click(
|
|
fn=on_generate_outline,
|
|
inputs=[title_input, genre_dropdown, sub_genre_dropdown, total_chapters, character_input, world_input, plot_input, custom_outline_prompt],
|
|
outputs=[outline_output, outline_status],
|
|
show_progress="full"
|
|
)
|
|
auto_generate_btn.click(
|
|
fn=on_auto_generate,
|
|
inputs=[title_input, genre_dropdown, sub_genre_dropdown, character_input, world_input, plot_input, outline_output, memory_type, memory_chapters, use_reflection_checkbox, create_style_dropdown],
|
|
outputs=[generation_progress, chapter_selector]
|
|
)
|
|
chapter_selector.change(
|
|
fn=on_chapter_select,
|
|
inputs=[chapter_selector],
|
|
outputs=[chapter_content_display]
|
|
)
|
|
stop_btn.click(fn=on_stop, outputs=[generation_progress])
|