diff --git a/.gitignore b/.gitignore index 8cb6184..b39c704 100644 --- a/.gitignore +++ b/.gitignore @@ -40,3 +40,4 @@ resources/scripts/db/* resources/scripts/localWhisper/.venv/* resources/scripts/localWhisper/_internal/* resources/scripts/localWhisper/local_whisper.exe +resources/scripts/joyCaption/.venv/* diff --git a/electron.vite.config.mjs b/electron.vite.config.mjs index aba2ad8..85e350d 100644 --- a/electron.vite.config.mjs +++ b/electron.vite.config.mjs @@ -2,16 +2,17 @@ import { resolve } from 'path' import { defineConfig, externalizeDepsPlugin, bytecodePlugin } from 'electron-vite' import vue from '@vitejs/plugin-vue' import Jsx from '@vitejs/plugin-vue-jsx' +import tsconfigPaths from 'vite-tsconfig-paths'; export default defineConfig({ main: { - plugins: [externalizeDepsPlugin(), bytecodePlugin()] + plugins: [externalizeDepsPlugin(), bytecodePlugin(),tsconfigPaths()] }, discord: { - plugins: [externalizeDepsPlugin(), bytecodePlugin()] + plugins: [externalizeDepsPlugin(), bytecodePlugin(),tsconfigPaths()] }, preload: { - plugins: [externalizeDepsPlugin(), bytecodePlugin()] + plugins: [externalizeDepsPlugin(), bytecodePlugin(),tsconfigPaths()] }, renderer: { resolve: { diff --git a/package-lock.json b/package-lock.json index d0f5135..ecbd5fc 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "laitool", - "version": "3.1.8", + "version": "3.1.9", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "laitool", - "version": "3.1.8", + "version": "3.1.9", "hasInstallScript": true, "dependencies": { "@alicloud/alimt20181012": "^1.2.0", @@ -39,6 +39,7 @@ "systeminformation": "^5.22.10", "tencentcloud-sdk-nodejs": "^4.0.821", "uuid": "^9.0.1", + "vite-tsconfig-paths": "^5.0.1", "vue-router": "^4.2.5", "wav-file-info": "^0.0.10", "winston": "^3.13.0", @@ -5084,6 +5085,12 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/globrex": { + "version": "0.1.2", + "resolved": "https://registry.npmmirror.com/globrex/-/globrex-0.1.2.tgz", + "integrity": "sha512-uHJgbwAMwNFf5mLst7IWLNg14x1CkeqglJb/K3doi4dw6q2IvAAmM/Y81kevy83wP+Sst+nutFTYOGg3d1lsxg==", + "license": "MIT" + }, "node_modules/gopd": { "version": "1.0.1", "license": "MIT", @@ -10341,6 +10348,26 @@ "typescript": ">=4.2.0" } }, + "node_modules/tsconfck": { + "version": "3.1.4", + "resolved": "https://registry.npmmirror.com/tsconfck/-/tsconfck-3.1.4.tgz", + "integrity": "sha512-kdqWFGVJqe+KGYvlSO9NIaWn9jT1Ny4oKVzAJsKii5eoE9snzTJzL4+MMVOMn+fikWGFmKEylcXL710V/kIPJQ==", + "license": "MIT", + "bin": { + "tsconfck": "bin/tsconfck.js" + }, + "engines": { + "node": "^18 || >=20" + }, + "peerDependencies": { + "typescript": "^5.0.0" + }, + "peerDependenciesMeta": { + "typescript": { + "optional": true + } + } + }, "node_modules/tslib": { "version": "2.6.2", "license": "0BSD", @@ -10532,6 +10559,25 @@ } } }, + "node_modules/vite-tsconfig-paths": { + "version": "5.0.1", + "resolved": "https://registry.npmmirror.com/vite-tsconfig-paths/-/vite-tsconfig-paths-5.0.1.tgz", + "integrity": "sha512-yqwv+LstU7NwPeNqajZzLEBVpUFU6Dugtb2P84FXuvaoYA+/70l9MHE+GYfYAycVyPSDYZ7mjOFuYBRqlEpTig==", + "license": "MIT", + "dependencies": { + "debug": "^4.1.1", + "globrex": "^0.1.2", + "tsconfck": "^3.0.3" + }, + "peerDependencies": { + "vite": "*" + }, + "peerDependenciesMeta": { + "vite": { + "optional": true + } + } + }, "node_modules/vooks": { "version": "0.2.12", "dev": true, diff --git a/package.json b/package.json index 1bfa724..341d2de 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "laitool", - "version": "3.1.8", + "version": "3.1.9", "description": "An AI tool for image processing, video processing, and other functions.", "main": "./out/main/index.js", "author": "laitool.cn", @@ -47,6 +47,7 @@ "systeminformation": "^5.22.10", "tencentcloud-sdk-nodejs": "^4.0.821", "uuid": "^9.0.1", + "vite-tsconfig-paths": "^5.0.1", "vue-router": "^4.2.5", "wav-file-info": "^0.0.10", "winston": "^3.13.0", @@ -93,4 +94,4 @@ "icon": "./resources/icon.ico" } } -} \ No newline at end of file +} diff --git a/resources/scripts/Lai.py b/resources/scripts/Lai.py index bdadd4e..b0c85ef 100644 --- a/resources/scripts/Lai.py +++ b/resources/scripts/Lai.py @@ -16,9 +16,33 @@ if len(sys.argv) < 2: sys.argv = [ "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe", "-c", - "D:/来推项目集/7.4/娱乐：江湖大哥退休，去拍电影/scripts/output_crop_00001.json", - "NVIDIA", - ] + "C:/Users/27698/Desktop/LAITool/project/73bca563-fd8f-42e3-b08b-e4b2ecbd0d71/scripts/output_00008_video_config.json", + "OTHER", + ] + # sys.argv = [ + # "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe", + # "-p", + # "C:/Users/27698/Desktop/LAITool/resources/config/sd_config.json", + # "input", + # "C:\\Users\\27698\\Desktop\\测试\\123", + # ] + # sys.argv = [ + # "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe", + # "-pf", + # "C:\\Users\\27698\\Desktop\\测试\\123\\tmp\\input_crop", + # ] + # sys.argv = [ + # "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe", + # "-ps", + # "C:\\Users\\27698\\Desktop\\测试\\123\\tmp\\input_crop\\00008.png", + # ] + # sys.argv = [ + # "C:\\Users\\27698\\Desktop\\LAITool\\resources\\scripts\\Lai.exe", + # "-pt", + # "C:\\Users\\27698\\Desktop\\测试\\123\\tmp\\input_crop\\1.txt", + # ] + + print(sys.argv) @@ -73,11 +97,26 @@ elif sys.argv[1] == "-f": with open(sys.argv[2], "w", encoding="utf-8") as file: json.dump(data, file, ensure_ascii=False, indent=4) -# 反推提示词 +# 反推提示词（指定项目下面的input_crop文件夹） elif sys.argv[1] == "-p": Push_back_Prompt.init(sys.argv[2], sys.argv[3], sys.argv[4]) pass +# 反推指定的文件夹 +elif sys.argv[1] == "-pf": + Push_back_Prompt.getAssignDir(sys.argv[2]) + pass + +# 反推指定的图片文件 +elif sys.argv[1] == "-ps": + Push_back_Prompt.getAssignImage(sys.argv[2]) + pass + +# 反推指定的图片文件 +elif sys.argv[1] == "-pt": + Push_back_Prompt.getAssignTxt(sys.argv[2]) + pass + elif sys.argv[1] == "-ka": shotSplit.get_fram(sys.argv[2], sys.argv[3], sys.argv[4]) pass diff --git a/resources/scripts/Push_back_Prompt.py b/resources/scripts/Push_back_Prompt.py index b7ad5ac..a8197cf 100644 --- a/resources/scripts/Push_back_Prompt.py +++ b/resources/scripts/Push_back_Prompt.py @@ -148,10 +148,11 @@ class WaifuDiffusionInterrogator(Interrogator): Dict[str, float], Dict[str, float] # rating confidents # tag confidents ]: # init model - if not hasattr(self, 'model') or self.model is None: + if not hasattr(self, "model") or self.model is None: model_path = os.path.join(cript_directory, "model/tag/model.onnx") tags_path = os.path.join(cript_directory, "model/tag/selected_tags.csv") from onnxruntime import InferenceSession + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] self.model = InferenceSession(str(model_path), providers=providers) print(f"从{model_path} 读取 {self.name}模型") @@ -223,7 +224,89 @@ def filter_action(tag_actions: [], tags: []): return action_tags, other_tags -def init(sd_setting,m,project_path): +def getAssignTxt(txtPath): + if not os.path.exists(txtPath): + os.makedirs(txtPath) + + # load model + model = WaifuDiffusionInterrogator( + "wd14-convnextv2-v2", + repo_id="SmilingWolf/wd-v1-4-convnextv2-tagger-v2", + revision="v2.0", + ) + + frame_files = [] + with open(txtPath, 'r', encoding='utf-8') as file: + for line in file: + frame_files.append(line.strip()) # 使用 strip() 去除每行的换行符和多余的空白 + + # 轮询开始输出 + frame_files.sort() + + for frame_file in frame_files: + txt = getTags(model, frame_file) + # tags = txt.split(",") + + # save tag + txt_file = os.path.join(os.path.dirname(frame_file), f"{Path(frame_file).stem}.txt") + with open(txt_file, "w", encoding="utf-8") as tags: + tags.write(txt) + print(f"{frame_file} 提示词反推完成") + sys.stdout.flush() + + + +def getAssignImage(imagePath): + if not os.path.exists(imagePath): + os.makedirs(imagePath) + + # load model + model = WaifuDiffusionInterrogator( + "wd14-convnextv2-v2", + repo_id="SmilingWolf/wd-v1-4-convnextv2-tagger-v2", + revision="v2.0", + ) + + txt = getTags(model, imagePath) + # tags = txt.split(",") + + # save tag + txt_file = os.path.join(os.path.dirname(imagePath), f"{Path(imagePath).stem}.txt") + with open(txt_file, "w", encoding="utf-8") as tags: + tags.write(txt) + print(f"{imagePath} 提示词反推完成") + sys.stdout.flush() + + +def getAssignDir(imagePath): + if not os.path.exists(imagePath): + os.makedirs(imagePath) + + # load model + model = WaifuDiffusionInterrogator( + "wd14-convnextv2-v2", + repo_id="SmilingWolf/wd-v1-4-convnextv2-tagger-v2", + revision="v2.0", + ) + + # 轮询开始输出 + frame_files = [f for f in os.listdir(imagePath) if f.endswith(".png")] + frame_files.sort() + + for frame in frame_files: + frame_file = os.path.join(imagePath, frame) + txt = getTags(model, frame_file) + # tags = txt.split(",") + + # save tag + txt_file = os.path.join(imagePath, f"{Path(frame_file).stem}.txt") + with open(txt_file, "w", encoding="utf-8") as tags: + tags.write(txt) + print(f"{frame} 提示词反推完成") + sys.stdout.flush() + + +def init(sd_setting, m, project_path): try: setting_json = public_tools.read_config(sd_setting, webui=False) except Exception as e: diff --git a/resources/scripts/__pycache__/Push_back_Prompt.cpython-310.pyc b/resources/scripts/__pycache__/Push_back_Prompt.cpython-310.pyc index 9a691ae..e9dfb4d 100644 Binary files a/resources/scripts/__pycache__/Push_back_Prompt.cpython-310.pyc and b/resources/scripts/__pycache__/Push_back_Prompt.cpython-310.pyc differ diff --git a/resources/scripts/db/book.realm.lock b/resources/scripts/db/book.realm.lock index 8f0002c..8c40e99 100644 Binary files a/resources/scripts/db/book.realm.lock and b/resources/scripts/db/book.realm.lock differ diff --git a/resources/scripts/db/software.realm.lock b/resources/scripts/db/software.realm.lock index 734ea88..4d29dfd 100644 Binary files a/resources/scripts/db/software.realm.lock and b/resources/scripts/db/software.realm.lock differ diff --git a/resources/scripts/joyCaption/0.26.0 b/resources/scripts/joyCaption/0.26.0 new file mode 100644 index 0000000..60ac347 --- /dev/null +++ b/resources/scripts/joyCaption/0.26.0 @@ -0,0 +1,29 @@ +Collecting accelerate + Downloading accelerate-1.0.1-py3-none-any.whl.metadata (19 kB) +Requirement already satisfied: numpy<3.0.0,>=1.17 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (2.1.2) +Requirement already satisfied: packaging>=20.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (24.1) +Collecting psutil (from accelerate) + Downloading psutil-6.1.0-cp37-abi3-win_amd64.whl.metadata (23 kB) +Requirement already satisfied: pyyaml in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (6.0.2) +Requirement already satisfied: torch>=1.10.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (2.5.0) +Requirement already satisfied: huggingface-hub>=0.21.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (0.26.1) +Requirement already satisfied: safetensors>=0.4.3 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from accelerate) (0.4.5) +Requirement already satisfied: filelock in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (3.16.1) +Requirement already satisfied: fsspec>=2023.5.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (2024.10.0) +Requirement already satisfied: requests in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (2.32.3) +Requirement already satisfied: tqdm>=4.42.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (4.66.5) +Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from huggingface-hub>=0.21.0->accelerate) (4.12.2) +Requirement already satisfied: networkx in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (3.4.2) +Requirement already satisfied: jinja2 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (3.1.4) +Requirement already satisfied: sympy==1.13.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from torch>=1.10.0->accelerate) (1.13.1) +Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from sympy==1.13.1->torch>=1.10.0->accelerate) (1.3.0) +Requirement already satisfied: colorama in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from tqdm>=4.42.1->huggingface-hub>=0.21.0->accelerate) (0.4.6) +Requirement already satisfied: MarkupSafe>=2.0 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from jinja2->torch>=1.10.0->accelerate) (3.0.2) +Requirement already satisfied: charset-normalizer<4,>=2 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (3.4.0) +Requirement already satisfied: idna<4,>=2.5 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (3.10) +Requirement already satisfied: urllib3<3,>=1.21.1 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (2.2.3) +Requirement already satisfied: certifi>=2017.4.17 in c:\users\27698\desktop\laitool\resources\scripts\joycaption\.venv\lib\site-packages (from requests->huggingface-hub>=0.21.0->accelerate) (2024.8.30) +Downloading accelerate-1.0.1-py3-none-any.whl (330 kB) +Downloading psutil-6.1.0-cp37-abi3-win_amd64.whl (254 kB) +Installing collected packages: psutil, accelerate +Successfully installed accelerate-1.0.1 psutil-6.1.0 diff --git a/resources/scripts/joyCaption/0.26.0' b/resources/scripts/joyCaption/0.26.0' new file mode 100644 index 0000000..e69de29 diff --git a/resources/scripts/joyCaption/app.py b/resources/scripts/joyCaption/app.py new file mode 100644 index 0000000..ae30b62 --- /dev/null +++ b/resources/scripts/joyCaption/app.py @@ -0,0 +1,336 @@ +import spaces +import gradio as gr +from huggingface_hub import InferenceClient +from torch import nn +from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM +from pathlib import Path +import torch +import torch.amp.autocast_mode +from PIL import Image +import os +import torchvision.transforms.functional as TVF + + +CLIP_PATH = "google/siglip-so400m-patch14-384" +CHECKPOINT_PATH = Path("cgrkzexw-599808") +TITLE = "

JoyCaption Alpha Two (2024-09-26a)

" +CAPTION_TYPE_MAP = { + "Descriptive": [ + "Write a descriptive caption for this image in a formal tone.", + "Write a descriptive caption for this image in a formal tone within {word_count} words.", + "Write a {length} descriptive caption for this image in a formal tone.", + ], + "Descriptive (Informal)": [ + "Write a descriptive caption for this image in a casual tone.", + "Write a descriptive caption for this image in a casual tone within {word_count} words.", + "Write a {length} descriptive caption for this image in a casual tone.", + ], + "Training Prompt": [ + "Write a stable diffusion prompt for this image.", + "Write a stable diffusion prompt for this image within {word_count} words.", + "Write a {length} stable diffusion prompt for this image.", + ], + "MidJourney": [ + "Write a MidJourney prompt for this image.", + "Write a MidJourney prompt for this image within {word_count} words.", + "Write a {length} MidJourney prompt for this image.", + ], + "Booru tag list": [ + "Write a list of Booru tags for this image.", + "Write a list of Booru tags for this image within {word_count} words.", + "Write a {length} list of Booru tags for this image.", + ], + "Booru-like tag list": [ + "Write a list of Booru-like tags for this image.", + "Write a list of Booru-like tags for this image within {word_count} words.", + "Write a {length} list of Booru-like tags for this image.", + ], + "Art Critic": [ + "Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc.", + "Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it within {word_count} words.", + "Analyze this image like an art critic would with information about its composition, style, symbolism, the use of color, light, any artistic movement it might belong to, etc. Keep it {length}.", + ], + "Product Listing": [ + "Write a caption for this image as though it were a product listing.", + "Write a caption for this image as though it were a product listing. Keep it under {word_count} words.", + "Write a {length} caption for this image as though it were a product listing.", + ], + "Social Media Post": [ + "Write a caption for this image as if it were being used for a social media post.", + "Write a caption for this image as if it were being used for a social media post. Limit the caption to {word_count} words.", + "Write a {length} caption for this image as if it were being used for a social media post.", + ], +} + +HF_TOKEN = os.environ.get("HF_TOKEN", None) + + +class ImageAdapter(nn.Module): + def __init__(self, input_features: int, output_features: int, ln1: bool, pos_emb: bool, num_image_tokens: int, deep_extract: bool): + super().__init__() + self.deep_extract = deep_extract + + if self.deep_extract: + input_features = input_features * 5 + + self.linear1 = nn.Linear(input_features, output_features) + self.activation = nn.GELU() + self.linear2 = nn.Linear(output_features, output_features) + self.ln1 = nn.Identity() if not ln1 else nn.LayerNorm(input_features) + self.pos_emb = None if not pos_emb else nn.Parameter(torch.zeros(num_image_tokens, input_features)) + + # Other tokens (<|image_start|>, <|image_end|>, <|eot_id|>) + self.other_tokens = nn.Embedding(3, output_features) + self.other_tokens.weight.data.normal_(mean=0.0, std=0.02) # Matches HF's implementation of llama3 + + def forward(self, vision_outputs: torch.Tensor): + if self.deep_extract: + x = torch.concat(( + vision_outputs[-2], + vision_outputs[3], + vision_outputs[7], + vision_outputs[13], + vision_outputs[20], + ), dim=-1) + assert len(x.shape) == 3, f"Expected 3, got {len(x.shape)}" # batch, tokens, features + assert x.shape[-1] == vision_outputs[-2].shape[-1] * 5, f"Expected {vision_outputs[-2].shape[-1] * 5}, got {x.shape[-1]}" + else: + x = vision_outputs[-2] + + x = self.ln1(x) + + if self.pos_emb is not None: + assert x.shape[-2:] == self.pos_emb.shape, f"Expected {self.pos_emb.shape}, got {x.shape[-2:]}" + x = x + self.pos_emb + + x = self.linear1(x) + x = self.activation(x) + x = self.linear2(x) + + # <|image_start|>, IMAGE, <|image_end|> + other_tokens = self.other_tokens(torch.tensor([0, 1], device=self.other_tokens.weight.device).expand(x.shape[0], -1)) + assert other_tokens.shape == (x.shape[0], 2, x.shape[2]), f"Expected {(x.shape[0], 2, x.shape[2])}, got {other_tokens.shape}" + x = torch.cat((other_tokens[:, 0:1], x, other_tokens[:, 1:2]), dim=1) + + return x + + def get_eot_embedding(self): + return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0) + + + +# Load CLIP +print("Loading CLIP") +clip_processor = AutoProcessor.from_pretrained(CLIP_PATH) +clip_model = AutoModel.from_pretrained(CLIP_PATH) +clip_model = clip_model.vision_model + +assert (CHECKPOINT_PATH / "clip_model.pt").exists() +print("Loading VLM's custom vision model") +checkpoint = torch.load(CHECKPOINT_PATH / "clip_model.pt", map_location='cpu') +checkpoint = {k.replace("_orig_mod.module.", ""): v for k, v in checkpoint.items()} +clip_model.load_state_dict(checkpoint) +del checkpoint + +clip_model.eval() +clip_model.requires_grad_(False) +clip_model.to("cuda") + + +# Tokenizer +print("Loading tokenizer") +tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_PATH / "text_model", use_fast=True) +assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance(tokenizer, PreTrainedTokenizerFast), f"Tokenizer is of type {type(tokenizer)}" + +# LLM +print("Loading LLM") +print("Loading VLM's custom text model") +text_model = AutoModelForCausalLM.from_pretrained(CHECKPOINT_PATH / "text_model", device_map=0, torch_dtype=torch.bfloat16) +text_model.eval() + +# Image Adapter +print("Loading image adapter") +image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False) +image_adapter.load_state_dict(torch.load(CHECKPOINT_PATH / "image_adapter.pt", map_location="cpu")) +image_adapter.eval() +image_adapter.to("cuda") + + +@spaces.GPU() +@torch.no_grad() +def stream_chat(input_image: Image.Image, caption_type: str, caption_length: str | int, extra_options: list[str], name_input: str, custom_prompt: str) -> tuple[str, str]: + torch.cuda.empty_cache() + + # 'any' means no length specified + length = None if caption_length == "any" else caption_length + + if isinstance(length, str): + try: + length = int(length) + except ValueError: + pass + + # Build prompt + if length is None: + map_idx = 0 + elif isinstance(length, int): + map_idx = 1 + elif isinstance(length, str): + map_idx = 2 + else: + raise ValueError(f"Invalid caption length: {length}") + + prompt_str = CAPTION_TYPE_MAP[caption_type][map_idx] + + # Add extra options + if len(extra_options) > 0: + prompt_str += " " + " ".join(extra_options) + + # Add name, length, word_count + prompt_str = prompt_str.format(name=name_input, length=caption_length, word_count=caption_length) + + if custom_prompt.strip() != "": + prompt_str = custom_prompt.strip() + + # For debugging + print(f"Prompt: {prompt_str}") + + # Preprocess image + # NOTE: I found the default processor for so400M to have worse results than just using PIL directly + #image = clip_processor(images=input_image, return_tensors='pt').pixel_values + image = input_image.resize((384, 384), Image.LANCZOS) + pixel_values = TVF.pil_to_tensor(image).unsqueeze(0) / 255.0 + pixel_values = TVF.normalize(pixel_values, [0.5], [0.5]) + pixel_values = pixel_values.to('cuda') + + # Embed image + # This results in Batch x Image Tokens x Features + with torch.amp.autocast_mode.autocast('cuda', enabled=True): + vision_outputs = clip_model(pixel_values=pixel_values, output_hidden_states=True) + embedded_images = image_adapter(vision_outputs.hidden_states) + embedded_images = embedded_images.to('cuda') + + # Build the conversation + convo = [ + { + "role": "system", + "content": "You are a helpful image captioner.", + }, + { + "role": "user", + "content": prompt_str, + }, + ] + + # Format the conversation + convo_string = tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = True) + assert isinstance(convo_string, str) + + # Tokenize the conversation + # prompt_str is tokenized separately so we can do the calculations below + convo_tokens = tokenizer.encode(convo_string, return_tensors="pt", add_special_tokens=False, truncation=False) + prompt_tokens = tokenizer.encode(prompt_str, return_tensors="pt", add_special_tokens=False, truncation=False) + assert isinstance(convo_tokens, torch.Tensor) and isinstance(prompt_tokens, torch.Tensor) + convo_tokens = convo_tokens.squeeze(0) # Squeeze just to make the following easier + prompt_tokens = prompt_tokens.squeeze(0) + + # Calculate where to inject the image + eot_id_indices = (convo_tokens == tokenizer.convert_tokens_to_ids("<|eot_id|>")).nonzero(as_tuple=True)[0].tolist() + assert len(eot_id_indices) == 2, f"Expected 2 <|eot_id|> tokens, got {len(eot_id_indices)}" + + preamble_len = eot_id_indices[1] - prompt_tokens.shape[0] # Number of tokens before the prompt + + # Embed the tokens + convo_embeds = text_model.model.embed_tokens(convo_tokens.unsqueeze(0).to('cuda')) + + # Construct the input + input_embeds = torch.cat([ + convo_embeds[:, :preamble_len], # Part before the prompt + embedded_images.to(dtype=convo_embeds.dtype), # Image + convo_embeds[:, preamble_len:], # The prompt and anything after it + ], dim=1).to('cuda') + + input_ids = torch.cat([ + convo_tokens[:preamble_len].unsqueeze(0), + torch.zeros((1, embedded_images.shape[1]), dtype=torch.long), # Dummy tokens for the image (TODO: Should probably use a special token here so as not to confuse any generation algorithms that might be inspecting the input) + convo_tokens[preamble_len:].unsqueeze(0), + ], dim=1).to('cuda') + attention_mask = torch.ones_like(input_ids) + + # Debugging + print(f"Input to model: {repr(tokenizer.decode(input_ids[0]))}") + + #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None) + #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, top_k=10, temperature=0.5, suppress_tokens=None) + generate_ids = text_model.generate(input_ids, inputs_embeds=input_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=True, suppress_tokens=None) # Uses the default which is temp=0.6, top_p=0.9 + + # Trim off the prompt + generate_ids = generate_ids[:, input_ids.shape[1]:] + if generate_ids[0][-1] == tokenizer.eos_token_id or generate_ids[0][-1] == tokenizer.convert_tokens_to_ids("<|eot_id|>"): + generate_ids = generate_ids[:, :-1] + + caption = tokenizer.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0] + + return prompt_str, caption.strip() + + +with gr.Blocks() as demo: + gr.HTML(TITLE) + + with gr.Row(): + with gr.Column(): + input_image = gr.Image(type="pil", label="Input Image") + + caption_type = gr.Dropdown( + choices=["Descriptive", "Descriptive (Informal)", "Training Prompt", "MidJourney", "Booru tag list", "Booru-like tag list", "Art Critic", "Product Listing", "Social Media Post"], + label="Caption Type", + value="Descriptive", + ) + + caption_length = gr.Dropdown( + choices=["any", "very short", "short", "medium-length", "long", "very long"] + + [str(i) for i in range(20, 261, 10)], + label="Caption Length", + value="long", + ) + + extra_options = gr.CheckboxGroup( + choices=[ + "If there is a person/character in the image you must refer to them as {name}.", + "Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).", + "Include information about lighting.", + "Include information about camera angle.", + "Include information about whether there is a watermark or not.", + "Include information about whether there are JPEG artifacts or not.", + "If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.", + "Do NOT include anything sexual; keep it PG.", + "Do NOT mention the image's resolution.", + "You MUST include information about the subjective aesthetic quality of the image from low to very high.", + "Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.", + "Do NOT mention any text that is in the image.", + "Specify the depth of field and whether the background is in focus or blurred.", + "If applicable, mention the likely use of artificial or natural lighting sources.", + "Do NOT use any ambiguous language.", + "Include whether the image is sfw, suggestive, or nsfw.", + "ONLY describe the most important elements of the image." + ], + label="Extra Options" + ) + + name_input = gr.Textbox(label="Person/Character Name (if applicable)") + gr.Markdown("**Note:** Name input is only used if an Extra Option is selected that requires it.") + + custom_prompt = gr.Textbox(label="Custom Prompt (optional, will override all other settings)") + gr.Markdown("**Note:** Alpha Two is not a general instruction follower and will not follow prompts outside its training data well. Use this feature with caution.") + + run_button = gr.Button("Caption") + + with gr.Column(): + output_prompt = gr.Textbox(label="Prompt that was used") + output_caption = gr.Textbox(label="Caption") + + run_button.click(fn=stream_chat, inputs=[input_image, caption_type, caption_length, extra_options, name_input, custom_prompt], outputs=[output_prompt, output_caption]) + + +if __name__ == "__main__": + demo.launch() \ No newline at end of file diff --git a/resources/scripts/joyCaption/batch-caption.py b/resources/scripts/joyCaption/batch-caption.py new file mode 100644 index 0000000..30306d0 --- /dev/null +++ b/resources/scripts/joyCaption/batch-caption.py @@ -0,0 +1,423 @@ +#!/usr/bin/env python3 +""" +Use JoyCaption to caption images. +""" +import argparse +import dataclasses +import json +import logging +import os +import random +from pathlib import Path + +import PIL.Image +import torch +import torch.amp +import torchvision.transforms.functional as TVF +from PIL import Image +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm +from transformers import ( + AutoTokenizer, + LlavaForConditionalGeneration, + PreTrainedTokenizer, + PreTrainedTokenizerFast, +) + + +def none_or_type(value, desired_type): + if value == "None": + return None + return desired_type(value) + + +parser = argparse.ArgumentParser() +parser.add_argument("--glob", type=str, help="Glob pattern to find images") +parser.add_argument("--filelist", type=str, help="File containing list of images") +parser.add_argument("--prompt", type=str, help="Prompt to use") +parser.add_argument( + "--prompt-file", type=str, help="JSON file containing prompts to use" +) +parser.add_argument("--batch-size", type=int, default=1, help="Batch size") +parser.add_argument( + "--greedy", action="store_true", help="Use greedy decoding instead of sampling" +) +parser.add_argument( + "--temperature", type=float, default=0.6, help="Sampling temperature" +) +parser.add_argument( + "--top-p", type=lambda x: none_or_type(x, float), default=0.9, help="Top-p sampling" +) +parser.add_argument( + "--top-k", type=lambda x: none_or_type(x, int), default=None, help="Top-k sampling" +) +parser.add_argument( + "--max-new-tokens", + type=int, + default=256, + help="Maximum length of the generated caption (in tokens)", +) +parser.add_argument( + "--num-workers", + type=int, + default=4, + help="Number of workers loading images in parallel", +) +parser.add_argument( + "--model", + type=str, + default="fancyfeast/llama-joycaption-alpha-two-hf-llava", + help="Model to use", +) + + +PIL.Image.MAX_IMAGE_PIXELS = 933120000 # Quiets Pillow from giving warnings on really large images (WARNING: Exposes a risk of DoS from malicious images) + + +@dataclasses.dataclass +class Prompt: + prompt: str + weight: float + + +@torch.no_grad() +def main(): + # Logging + logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s" + ) + + # Parse arguments + args = parser.parse_args() + logging.info(f"Arguments: {args}") + + args.prompt = "Please describe the image." + # Make sure we have a prompt or a prompt file + prompts = parse_prompts(args.prompt, args.prompt_file) + + args.filelist = ( + "C:\\Users\\27698\\Desktop\\node\\12\\12.txt" + ) + # Find the images + image_paths = find_images(args.glob, args.filelist) + if len(image_paths) == 0: + logging.warning("No images found") + return + logging.info(f"Found {len(image_paths)} images") + + # Ignore all images that already have captions + image_paths = [ + path for path in image_paths if not Path(path).with_suffix(".txt").exists() + ] + + # Load JoyCaption + tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=True) + assert isinstance(tokenizer, PreTrainedTokenizer) or isinstance( + tokenizer, PreTrainedTokenizerFast + ), f"Tokenizer is of type {type(tokenizer)}" + llava_model = LlavaForConditionalGeneration.from_pretrained( + args.model, torch_dtype="bfloat16" + ) + assert isinstance(llava_model, LlavaForConditionalGeneration) + + dataset = ImageDataset( + prompts, + image_paths, + tokenizer, + llava_model.config.image_token_index, + llava_model.config.image_seq_length, + ) + dataloader = DataLoader( + dataset, + collate_fn=dataset.collate_fn, + num_workers=args.num_workers, + shuffle=False, + drop_last=False, + batch_size=args.batch_size, + ) + end_of_header_id = tokenizer.convert_tokens_to_ids("<|end_header_id|>") + end_of_turn_id = tokenizer.convert_tokens_to_ids("<|eot_id|>") + assert isinstance(end_of_header_id, int) and isinstance(end_of_turn_id, int) + + pbar = tqdm(total=len(image_paths), desc="Captioning images...", dynamic_ncols=True) + for batch in dataloader: + vision_dtype = ( + llava_model.vision_tower.vision_model.embeddings.patch_embedding.weight.dtype + ) + vision_device = ( + llava_model.vision_tower.vision_model.embeddings.patch_embedding.weight.device + ) + language_device = ( + llava_model.language_model.get_input_embeddings().weight.device + ) + + # Move to GPU + pixel_values = batch["pixel_values"].to(vision_device, non_blocking=True) + input_ids = batch["input_ids"].to(language_device, non_blocking=True) + attention_mask = batch["attention_mask"].to(language_device, non_blocking=True) + + # Normalize the image + pixel_values = pixel_values / 255.0 + pixel_values = TVF.normalize(pixel_values, [0.5], [0.5]) + pixel_values = pixel_values.to(vision_dtype) + + # Generate the captions + generate_ids = llava_model.generate( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + max_new_tokens=args.max_new_tokens, + do_sample=not args.greedy, + suppress_tokens=None, + use_cache=True, + temperature=args.temperature, + top_k=args.top_k, + top_p=args.top_p, + ) + + # Trim off the prompts + assert isinstance(generate_ids, torch.Tensor) + generate_ids = generate_ids.tolist() + generate_ids = [ + trim_off_prompt(ids, end_of_header_id, end_of_turn_id) + for ids in generate_ids + ] + + # Decode the captions + captions = tokenizer.batch_decode( + generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False + ) + captions = [c.strip() for c in captions] + + for path, caption in zip(batch["paths"], captions): + write_caption(Path(path), caption) + + pbar.update(len(captions)) + + +def trim_off_prompt(input_ids: list[int], eoh_id: int, eot_id: int) -> list[int]: + # Trim off the prompt + while True: + try: + i = input_ids.index(eoh_id) + except ValueError: + break + + input_ids = input_ids[i + 1 :] + + # Trim off the end + try: + i = input_ids.index(eot_id) + except ValueError: + return input_ids + + return input_ids[:i] + + +def write_caption(image_path: Path, caption: str): + caption_path = image_path.with_suffix(".txt") + + try: + f = os.open( + caption_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL + ) # Write-only, create if not exist, fail if exists + except FileExistsError: + logging.warning(f"Caption file '{caption_path}' already exists") + return + except Exception as e: + logging.error(f"Failed to open caption file '{caption_path}': {e}") + return + + try: + os.write(f, caption.encode("utf-8")) + os.close(f) + except Exception as e: + logging.error(f"Failed to write caption to '{caption_path}': {e}") + return + + +class ImageDataset(Dataset): + def __init__( + self, + prompts: list[Prompt], + paths: list[Path], + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast, + image_token_id: int, + image_seq_length: int, + ): + self.prompts = prompts + self.paths = paths + self.tokenizer = tokenizer + self.image_token_id = image_token_id + self.image_seq_length = image_seq_length + self.pad_token_id = tokenizer.pad_token_id + + def __len__(self): + return len(self.paths) + + def __getitem__(self, idx: int) -> dict: + path = self.paths[idx] + + # Pick a prompt + prompt_str = random.choices( + self.prompts, weights=[p.weight for p in self.prompts] + )[0].prompt + + # Preprocess image + # NOTE: I don't use the Processor here and instead do it manually. + # This is because in my testing a simple resize in Pillow yields higher quality results than the Processor, + # and the Processor had some buggy behavior on some images. + # And yes, with the so400m model, the model expects the image to be squished into a square, not padded. + try: + image = Image.open(path) + if image.size != (384, 384): + image = image.resize((384, 384), Image.LANCZOS) + image = image.convert("RGB") + pixel_values = TVF.pil_to_tensor(image) + except Exception as e: + logging.error(f"Failed to load image '{path}': {e}") + pixel_values = None # Will be filtered out later + + # Build the conversation + convo = [ + { + "role": "system", + "content": "You are a helpful image captioner.", + }, + { + "role": "user", + "content": prompt_str, + }, + ] + + # Format the conversation + convo_string = self.tokenizer.apply_chat_template( + convo, tokenize=False, add_generation_prompt=True + ) + assert isinstance(convo_string, str) + + # Tokenize the conversation + convo_tokens = self.tokenizer.encode( + convo_string, add_special_tokens=False, truncation=False + ) + + # Repeat the image tokens + input_tokens = [] + for token in convo_tokens: + if token == self.image_token_id: + input_tokens.extend([self.image_token_id] * self.image_seq_length) + else: + input_tokens.append(token) + + input_ids = torch.tensor(input_tokens, dtype=torch.long) + attention_mask = torch.ones_like(input_ids) + + return { + "path": path, + "pixel_values": pixel_values, + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + def collate_fn(self, batch: list[dict]) -> dict: + # Filter out images that failed to load + batch = [item for item in batch if item["pixel_values"] is not None] + + # Pad input_ids and attention_mask + # Have to use left padding because HF's generate can't handle right padding it seems + max_length = max(item["input_ids"].shape[0] for item in batch) + n_pad = [max_length - item["input_ids"].shape[0] for item in batch] + input_ids = torch.stack( + [ + torch.nn.functional.pad( + item["input_ids"], (n, 0), value=self.pad_token_id + ) + for item, n in zip(batch, n_pad) + ] + ) + attention_mask = torch.stack( + [ + torch.nn.functional.pad(item["attention_mask"], (n, 0), value=0) + for item, n in zip(batch, n_pad) + ] + ) + + # Stack pixel values + pixel_values = torch.stack([item["pixel_values"] for item in batch]) + + # Paths + paths = [item["path"] for item in batch] + + return { + "paths": paths, + "pixel_values": pixel_values, + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + +def parse_prompts(prompt_str: str | None, prompt_file: str | None) -> list[Prompt]: + if prompt_str is not None and prompt_file is not None: + raise ValueError("Cannot specify both --prompt and --prompt-file") + + if prompt_str is not None: + return [Prompt(prompt=prompt_str, weight=1.0)] + + if prompt_file is None: + raise ValueError("Must specify either --prompt or --prompt-file") + + data = json.loads(Path(prompt_file).read_text()) + + if not isinstance(data, list): + raise ValueError("Expected JSON file to contain a list of prompts") + + prompts = [] + + for item in data: + if isinstance(item, str): + prompts.append(Prompt(prompt=item, weight=1.0)) + elif ( + isinstance(item, dict) + and "prompt" in item + and "weight" in item + and isinstance(item["prompt"], str) + and isinstance(item["weight"], (int, float)) + ): + prompts.append(Prompt(prompt=item["prompt"], weight=item["weight"])) + else: + raise ValueError( + f"Invalid prompt in JSON file. Should be either a string or an object with 'prompt' and 'weight' fields: {item}" + ) + + if len(prompts) == 0: + raise ValueError("No prompts found in JSON file") + + if sum(p.weight for p in prompts) <= 0.0: + raise ValueError("Prompt weights must sum to a positive number") + + return prompts + + +def find_images(glob: str | None, filelist: str | Path | None) -> list[Path]: + if glob is None and filelist is None: + raise ValueError("Must specify either --glob or --filelist") + + paths = [] + + if glob is not None: + paths.extend(Path(".").glob(glob)) + + if filelist is not None: + paths.extend( + ( + Path(line.strip()) + for line in Path(filelist).read_text().strip().splitlines() + if line.strip() != "" + ) + ) + + return paths + + +if __name__ == "__main__": + main() diff --git a/resources/scripts/joyCaption/joy.py b/resources/scripts/joyCaption/joy.py new file mode 100644 index 0000000..51dcf75 --- /dev/null +++ b/resources/scripts/joyCaption/joy.py @@ -0,0 +1,75 @@ +import torch +import torch.amp +import torchvision.transforms.functional as TVF +from PIL import Image +from transformers import AutoTokenizer, LlavaForConditionalGeneration + + +IMAGE_PATH = "C:/Users/27698/Desktop/node/12/00001.png" +PROMPT = "Write a long descriptive caption for this image in a formal tone." +MODEL_NAME = "fancyfeast/llama-joycaption-alpha-two-hf-llava" + + +# Load JoyCaption +# bfloat16 is the native dtype of the LLM used in JoyCaption (Llama 3.1) +# device_map=0 loads the model into the first GPU +tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True) +llava_model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME, torch_dtype="bfloat16", device_map="cuda:0") +llava_model.eval() + +with torch.no_grad(): + # Load and preprocess image + # Normally you would use the Processor here, but the image module's processor + # has some buggy behavior and a simple resize in Pillow yields higher quality results + image = Image.open(IMAGE_PATH) + + if image.size != (384, 384): + image = image.resize((384, 384), Image.LANCZOS) + + image = image.convert("RGB") + pixel_values = TVF.pil_to_tensor(image) + + # Normalize the image + pixel_values = pixel_values / 255.0 + pixel_values = TVF.normalize(pixel_values, [0.5], [0.5]) + pixel_values = pixel_values.to(torch.bfloat16).unsqueeze(0) + + # Build the conversation + convo = [ + { + "role": "system", + "content": "You are a helpful image captioner.", + }, + { + "role": "user", + "content": PROMPT, + }, + ] + + # Format the conversation + convo_string = tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=True) + + # Tokenize the conversation + convo_tokens = tokenizer.encode(convo_string, add_special_tokens=False, truncation=False) + + # Repeat the image tokens + input_tokens = [] + for token in convo_tokens: + if token == llava_model.config.image_token_index: + input_tokens.extend([llava_model.config.image_token_index] * llava_model.config.image_seq_length) + else: + input_tokens.append(token) + + input_ids = torch.tensor(input_tokens, dtype=torch.long).unsqueeze(0) + attention_mask = torch.ones_like(input_ids) + + # Generate the caption + generate_ids = llava_model.generate(input_ids=input_ids.to('cuda'), pixel_values=pixel_values.to('cuda'), attention_mask=attention_mask.to('cuda'), max_new_tokens=300, do_sample=True, suppress_tokens=None, use_cache=True)[0] + + # Trim off the prompt + generate_ids = generate_ids[input_ids.shape[1]:] + + # Decode the caption + caption = tokenizer.decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) + caption = caption.strip() + print(caption) \ No newline at end of file diff --git a/src/define/Tools/common.ts b/src/define/Tools/common.ts index de20486..b909f42 100644 --- a/src/define/Tools/common.ts +++ b/src/define/Tools/common.ts @@ -1,3 +1,4 @@ +import { escapeRegExp } from "lodash"; /** * 检查字符串中是不是包含中文或者标点符号 @@ -66,4 +67,25 @@ export async function ExecuteConcurrently(tasks: Array<() => Promise>, conc } } return Promise.all(results); +} + +/** + * 替换主字符串中的子字符串 + * @param mainString 主字符串 + * @param substringArray 子字符串数组 + * @returns 替换后的字符串 + */ +export function ReplaceSubstrings(mainString: string, substringArray: string[], replacement: string = " "): string { + // 按长度降序排序子字符串数组，以确保较长的子字符串先被替换 + substringArray.sort((a, b) => b.length - a.length); + + // 对每个子字符串进行替换 + for (const substring of substringArray) { + // 创建一个正则表达式，用于全局匹配子字符串 + const regex = new RegExp(escapeRegExp(substring), 'g'); + // 将匹配到的子字符串替换为等长的空格 + mainString = mainString.replace(regex, replacement); + } + + return mainString; } \ No newline at end of file diff --git a/src/define/Tools/validate.ts b/src/define/Tools/validate.ts index 9260f55..f462742 100644 --- a/src/define/Tools/validate.ts +++ b/src/define/Tools/validate.ts @@ -5,7 +5,6 @@ * @returns 可以解析返回true，否则返回false */ export function ValidateJson(str: string): boolean { - try { JSON.parse(str); return true diff --git a/src/define/db/service/Book/bookTaskDetailService.ts b/src/define/db/service/Book/bookTaskDetailService.ts index 5621bfb..32d4b1a 100644 --- a/src/define/db/service/Book/bookTaskDetailService.ts +++ b/src/define/db/service/Book/bookTaskDetailService.ts @@ -244,7 +244,7 @@ export class BookTaskDetailService extends BaseRealmService { * 删除满足条件的对象吗，必传小说ID和小说任务ID * @param condition bookId，bookTaskId，name，id */ - DeleteBookTaskDetail(condition) { + DeleteBookTaskDetail(condition: Book.DeleteBookTaskDetailCondition) { try { if (isEmpty(condition.id) && isEmpty(condition.bookTaskId) && isEmpty(condition.bookId)) { throw new Error('删除小说分镜信息失败，没有必要参数') diff --git a/src/define/define_string/bookDefineString.ts b/src/define/define_string/bookDefineString.ts new file mode 100644 index 0000000..06c9eba --- /dev/null +++ b/src/define/define_string/bookDefineString.ts @@ -0,0 +1,164 @@ +const BOOK = { + MAIN_DATA_RETURN: 'MAIN_DATA_RETURN', // 监听任务返回 + REPLACE_VIDEO_CURRENT_FRAME: 'REPLACE_VIDEO_CURRENT_FRAME', + GET_BOOK_TYPE: 'GET_BOOK_TYPE', + ADD_OR_MODIFY_BOOK: 'ADD_OR_MODIFY_BOOK', + GET_BOOK_DATA: 'GET_BOOK_DATA', + GET_FRAME_DATA: 'GET_FRAME_DATA', + GET_BOOK_TASK_DATA: 'GET_BOOK_TASK_DATA', + AUTO_ACTION: 'AUTO_ACTION', + SAVE_BOOK_SUBTITLE_POSITION: 'SAVE_BOOK_SUBTITLE_POSITION', + OPEN_BOOK_SUBTITLE_POSITION_SCREENSHOT: 'OPEN_BOOK_SUBTITLE_POSITION_SCREENSHOT', + GET_CURRENT_FRAME_TEXT: 'GET_CURRENT_FRAME_TEXT', + GET_VIDEO_FRAME_TEXT: 'GET_VIDEO_FRAME_TEXT', + GET_BOOK_TASK_DETAIL: 'GET_BOOK_TASK_DETAIL', + REVERSE_PROMPT_TO_GPT_PROMPT: 'REVERSE_PROMPT_TO_GPT_PROMPT', + SINGLE_REVERSE_TO_GPT_PROMPT: 'SINGLE_REVERSE_TO_GPT_PROMPT', + SAVE_IMAGE_STYLE: 'SAVE_IMAGE_STYLE', + IMAGE_LOCK_OPERATION: "IMAGE_LOCK_OPERATION", + DOWNLOAD_IMAGE_AND_SPLIT: "DOWNLOAD_IMAGE_AND_SPLIT", + ONE_TO_FOUR_BOOK_TASK: "ONE_TO_FOUR_BOOK_TASK", + RESET_BOOK_TASK: "RESET_BOOK_TASK", + DELETE_BOOK_TASK: "DELETE_BOOK_TASK", + GENERATE_IMAGE_ALL: "GENERATE_IMAGE_ALL", + CHECK_IMAGE_FILE_SIZE: "CHECK_IMAGE_FILE_SIZE", + HD_IMAGE: "HD_IMAGE", + USE_BOOK_VIDEO_DATA_TO_BOOK_TASK: "USE_BOOK_VIDEO_DATA_TO_BOOK_TASK", + ADD_JIANYING_DRAFT: "ADD_JIANYING_DRAFT", + EXPORT_COPYWRITING: "EXPORT_COPYWRITING", + IMPORT_COPYWRITING: 'IMPORT_COPYWRITING', + MERGE_PROMPT: "MERGE_PROMPT", + RESET_BOOK_DATA: "RESET_BOOK_DATA", + DELETE_BOOK_DATA: "DELETE_BOOK_DATA", + CLEAR_IMPORT_WORD: "CLEAR_IMPORT_WORD", + RESET_GPT_REVERSE_DATA: "RESET_GPT_REVERSE_DATA", + REMOVE_MERGE_PROMPT_DATA: "REMOVE_MERGE_PROMPT_DATA", + REMOVE_GENERATE_IMAGE: 'REMOVE_GENERATE_IMAGE', + ADD_NEW_BOOK_TASK: "ADD_NEW_BOOK_TASK", + REPLACE_BOOK_DATA: "REPLACE_BOOK_DATA", + SAVE_COPYWRITING: 'SAVE_COPYWRITING', + + //#region 通用 + + /** 重置小说批次详细信息的数据 */ + RESET_BOOK_TASK_DETAIL_DATA: "RESET_BOOK_TASK_DETAIL_DATA", + + //#endregion + + //#region 分镜 + + /** 剪映抽帧，并判断是不是切割视频 */ + JIANYING_FRAME: "JIANYING_FRAME", + + //#endregion + + //#region 提示词 + /** + * 原创推理提示词 + */ + ORIGINAL_GPT_PROMPT: "ORIGINAL_GPT_PROMPT", + + /** * 原创推理提示词返回 */ + ORIGINAL_GPT_PROMPT_RETURN: "ORIGINAL_GPT_PROMPT_RETURN", + + /** * 导入提示词，通用 */ + IMPORT_GPT_PROMPT: "IMPORT_GPT_PROMPT", + + /** 移除不想要的提示词 */ + REMOVE_BAD_PROMPT: "REMOVE_BAD_PROMPT", + //#endregion + + //#region 生图返回相关 + + /** + * MJ生图返回信息 + */ + MJ_IMAGE_GENERATE_RETURN: 'MJ_IMAGE_GENERATE_RETURN', + + /** + * SD生图返回信息 + */ + SD_IMAGE_GENERATE_RETURN: 'SD_IMAGE_GENERATE_RETURN', + + /** + * D3 出图返回信息 + */ + D3_IMAGE_GENERATE_RETURN: 'D3_IMAGE_GENERATE_RETURN', + + /** + * flux forge 生图返回信息 + */ + FLUX_FORGE_IMAGE_GENERATE_RETURN: "FLUX_FORGE_IMAGE_GENERATE_RETURN", + + /** + * flux api 生图返回信息 + */ + FLUX_API_IMAGE_GENERATE_RETURN: "FLUX_API_IMAGE_GENERATE_RETURN", + + //#endregion + + //#region 图片相关 + + /** + * 获取MJ的消息ID，下载图片并分割 + */ + GET_IMAGE_URL_AND_DOWNLOAD: "GET_IMAGE_URL_AND_DOWNLOAD", + + //#endregion + + COMPUTE_STORYBOARD: 'COMPUTE_STORYBOARD', + + GET_FRAME: 'GET_FRAME', + + FRAMING: 'FRAMING', + + //#region 字幕相关 + + /** + * 获取copywriting + */ + GET_COPYWRITING: 'GET_COPYWRITING', + /** + * 获取copywriting的返回信息 + */ + GET_COPYWRITING_RETURN: 'GET_COPYWRITING_RETURN', + + //#endregion + + REMOVE_WATERMARK: 'REMOVE_WATERMARK', + REMOVE_WATERMARK_RETURN: 'REMOVE_WATERMARK_RETURN', + + SPLI_TAUDIO: 'SPLI_TAUDIO', + SPLI_TAUDIO_RETURN: 'SPLI_TAUDIO_RETURN', + + //#region 反推相关 + /** + * 添加反推任务 + */ + ADD_REVERSE_PROMPT: 'ADD_REVERSE_PROMPT', + /** + * 反推任务返回信息 + */ + REVERSE_PROMPT_RETURN: 'REVERSE_PROMPT_RETURN', + //#endregion + + //#region 翻译相关 + + /** + * 反推翻译返回信息 + */ + REVERSE_PROMPT_TRANSLATE_RETURN: 'REVERSE_PROMPT_TRANSLATE_RETURN', + + //#endregion + + //#region + + /** 合成视频相关 */ + ADD_GENERATE_VIDEO_TASK: "ADD_GENERATE_VIDEO_TASK", + /** 生成视频返回 */ + GENERATE_VIDEO_RETURN: "GENERATE_VIDEO_RETURN", + //#endregion + + +} +export default BOOK; \ No newline at end of file diff --git a/src/define/define_string/index.ts b/src/define/define_string/index.ts index 95ba783..d10a215 100644 --- a/src/define/define_string/index.ts +++ b/src/define/define_string/index.ts @@ -2,11 +2,13 @@ import { SYSTEM } from "./systemDefineString" import TASK from "./taskDefineString" import TTS from "./ttsDefineString" import SETTING from "./settingDefineString" +import BOOK from "./bookDefineString" export const DEFINE_STRING = { SYSTEM: SYSTEM, TASK: TASK, TTS: TTS, + BOOK: BOOK, SETTING: SETTING, SHOW_GLOBAL_MESSAGE: "SHOW_GLOBAL_MESSAGE", SHOW_GLOBAL_MAIN_NOTIFICATION: 'SHOW_GLOBAL_MAIN_NOTIFICATION', @@ -225,146 +227,6 @@ export const DEFINE_STRING = { BATCH_PROCESS_IMAGE_RESULT: 'BATCH_PROCESS_IMAGE_RESULT', PROCESS_IMAGE_WATERMASK_CHECK: "PROCESS_IMAGE_WATERMASK_CHECK" }, - BOOK: { - MAIN_DATA_RETURN: 'MAIN_DATA_RETURN', // 监听任务返回 - REPLACE_VIDEO_CURRENT_FRAME: 'REPLACE_VIDEO_CURRENT_FRAME', - GET_BOOK_TYPE: 'GET_BOOK_TYPE', - ADD_OR_MODIFY_BOOK: 'ADD_OR_MODIFY_BOOK', - GET_BOOK_DATA: 'GET_BOOK_DATA', - GET_FRAME_DATA: 'GET_FRAME_DATA', - GET_BOOK_TASK_DATA: 'GET_BOOK_TASK_DATA', - AUTO_ACTION: 'AUTO_ACTION', - SAVE_BOOK_SUBTITLE_POSITION: 'SAVE_BOOK_SUBTITLE_POSITION', - OPEN_BOOK_SUBTITLE_POSITION_SCREENSHOT: 'OPEN_BOOK_SUBTITLE_POSITION_SCREENSHOT', - GET_CURRENT_FRAME_TEXT: 'GET_CURRENT_FRAME_TEXT', - GET_VIDEO_FRAME_TEXT: 'GET_VIDEO_FRAME_TEXT', - GET_BOOK_TASK_DETAIL: 'GET_BOOK_TASK_DETAIL', - REVERSE_PROMPT_TO_GPT_PROMPT: 'REVERSE_PROMPT_TO_GPT_PROMPT', - SINGLE_REVERSE_TO_GPT_PROMPT: 'SINGLE_REVERSE_TO_GPT_PROMPT', - SAVE_IMAGE_STYLE: 'SAVE_IMAGE_STYLE', - IMAGE_LOCK_OPERATION: "IMAGE_LOCK_OPERATION", - DOWNLOAD_IMAGE_AND_SPLIT: "DOWNLOAD_IMAGE_AND_SPLIT", - ONE_TO_FOUR_BOOK_TASK: "ONE_TO_FOUR_BOOK_TASK", - RESET_BOOK_TASK: "RESET_BOOK_TASK", - DELETE_BOOK_TASK: "DELETE_BOOK_TASK", - GENERATE_IMAGE_ALL: "GENERATE_IMAGE_ALL", - CHECK_IMAGE_FILE_SIZE: "CHECK_IMAGE_FILE_SIZE", - HD_IMAGE: "HD_IMAGE", - USE_BOOK_VIDEO_DATA_TO_BOOK_TASK: "USE_BOOK_VIDEO_DATA_TO_BOOK_TASK", - ADD_JIANYING_DRAFT: "ADD_JIANYING_DRAFT", - EXPORT_COPYWRITING: "EXPORT_COPYWRITING", - IMPORT_COPYWRITING: 'IMPORT_COPYWRITING', - MERGE_PROMPT: "MERGE_PROMPT", - RESET_BOOK_DATA: "RESET_BOOK_DATA", - DELETE_BOOK_DATA: "DELETE_BOOK_DATA", - CLEAR_IMPORT_WORD: "CLEAR_IMPORT_WORD", - RESET_GPT_REVERSE_DATA: "RESET_GPT_REVERSE_DATA", - REMOVE_MERGE_PROMPT_DATA: "REMOVE_MERGE_PROMPT_DATA", - REMOVE_GENERATE_IMAGE: 'REMOVE_GENERATE_IMAGE', - ADD_NEW_BOOK_TASK: "ADD_NEW_BOOK_TASK", - REPLACE_BOOK_DATA: "REPLACE_BOOK_DATA", - SAVE_COPYWRITING: 'SAVE_COPYWRITING', - - //#region 提示词 - /** - * 原创推理提示词 - */ - ORIGINAL_GPT_PROMPT: "ORIGINAL_GPT_PROMPT", - /** - * 原创推理提示词返回 - */ - ORIGINAL_GPT_PROMPT_RETURN: "ORIGINAL_GPT_PROMPT_RETURN", - - /** - * 导入提示词，通用 - */ - IMPORT_GPT_PROMPT: "IMPORT_GPT_PROMPT", - //#endregion - - //#region 生图返回相关 - - /** - * MJ生图返回信息 - */ - MJ_IMAGE_GENERATE_RETURN: 'MJ_IMAGE_GENERATE_RETURN', - - /** - * SD生图返回信息 - */ - SD_IMAGE_GENERATE_RETURN: 'SD_IMAGE_GENERATE_RETURN', - - /** - * D3 出图返回信息 - */ - D3_IMAGE_GENERATE_RETURN: 'D3_IMAGE_GENERATE_RETURN', - - /** - * flux forge 生图返回信息 - */ - FLUX_FORGE_IMAGE_GENERATE_RETURN: "FLUX_FORGE_IMAGE_GENERATE_RETURN", - - /** - * flux api 生图返回信息 - */ - FLUX_API_IMAGE_GENERATE_RETURN: "FLUX_API_IMAGE_GENERATE_RETURN", - - //#endregion - - //#region 图片相关 - - /** - * 获取MJ的消息ID，下载图片并分割 - */ - GET_IMAGE_URL_AND_DOWNLOAD: "GET_IMAGE_URL_AND_DOWNLOAD", - - //#endregion - - COMPUTE_STORYBOARD: 'COMPUTE_STORYBOARD', - - GET_FRAME: 'GET_FRAME', - - FRAMING: 'FRAMING', - - //#region 字幕相关 - - /** - * 获取copywriting - */ - GET_COPYWRITING: 'GET_COPYWRITING', - /** - * 获取copywriting的返回信息 - */ - GET_COPYWRITING_RETURN: 'GET_COPYWRITING_RETURN', - - //#endregion - - REMOVE_WATERMARK: 'REMOVE_WATERMARK', - REMOVE_WATERMARK_RETURN: 'REMOVE_WATERMARK_RETURN', - - SPLI_TAUDIO: 'SPLI_TAUDIO', - SPLI_TAUDIO_RETURN: 'SPLI_TAUDIO_RETURN', - - //#region 反推相关 - /** - * 添加反推任务 - */ - ADD_REVERSE_PROMPT: 'ADD_REVERSE_PROMPT', - /** - * 反推任务返回信息 - */ - REVERSE_PROMPT_RETURN: 'REVERSE_PROMPT_RETURN', - //#endregion - - //#region 翻译相关 - - /** - * 反推翻译返回信息 - */ - REVERSE_PROMPT_TRANSLATE_RETURN: 'REVERSE_PROMPT_TRANSLATE_RETURN', - - //#endregion - - }, PROMPT: { GET_SORT_OPTIONS: 'GET_SORT_OPTIONS', SAVE_PROMPT_SORT_DATA: 'SAVE_PROMPT_SORT_DATA', diff --git a/src/define/enum/bookEnum.ts b/src/define/enum/bookEnum.ts index 51401a5..49ae8af 100644 --- a/src/define/enum/bookEnum.ts +++ b/src/define/enum/bookEnum.ts @@ -208,7 +208,12 @@ export enum BookTaskStatus { COMPOSING_DONE = 'composing_done', // 合成视频失败 - COMPOSING_FAIL = 'composing_fail' + COMPOSING_FAIL = 'composing_fail', + /** 添加草稿完成 */ + DRAFT_DONE = 'draft_done', + /** 添加草稿失败 */ + DRAFT_FAIL = 'draft_fail', + } export enum TagDefineType { diff --git a/src/define/enum/softwareEnum.ts b/src/define/enum/softwareEnum.ts index ed0d2bd..a0d58a4 100644 --- a/src/define/enum/softwareEnum.ts +++ b/src/define/enum/softwareEnum.ts @@ -60,6 +60,7 @@ export enum ResponseMessageType { GET_TEXT = 'getText', // 获取文案 REMOVE_WATERMARK = "REMOVE_WATERMARK",// 删除水印 MJ_REVERSE = 'MJ_REVERSE',// MJ反推，返回反推结果 + SD_REVERSE = 'SD_REVERSE',// MJ反推，返回反推结果 REVERSE_PROMPT_TRANSLATE = 'REVERSE_PROMPT_TRANSLATE',// 反推提示词翻译 GPT_PROMPT_TRANSLATE = 'GPT_PROMPT_TRANSLATE', // GPT提示词翻译 MJ_IMAGE = 'MJ_IMAGE',// MJ 生成图片 diff --git a/src/main/IPCEvent/bookIpc.ts b/src/main/IPCEvent/bookIpc.ts index 090f3b6..cb3f9de 100644 --- a/src/main/IPCEvent/bookIpc.ts +++ b/src/main/IPCEvent/bookIpc.ts @@ -92,10 +92,11 @@ export function BookIpc() { //#endregion - //#region 小说通用操作 + //#region 小说详情通用操作 ipcMain.handle(DEFINE_STRING.BOOK.REPLACE_BOOK_DATA, async (event, bookTaskId, replaceData) => await bookGeneral.ReplaceBookData(bookTaskId, replaceData)) + //#endregion //#region 分镜相关 @@ -116,6 +117,9 @@ export function BookIpc() { await bookFrame.ReplaceVideoCurrentFrame(bookTaskDetailId, currentTime) ) + // 剪映抽帧，并判断是不是切割视频 + ipcMain.handle(DEFINE_STRING.BOOK.JIANYING_FRAME, async (event, bookTaskId, draftName) => await bookFrame.JianyingFrame(bookTaskId, draftName)) + //#endregion //#region 提示词相关 @@ -154,6 +158,12 @@ export function BookIpc() { async (event, bookTaskId: string, txtPath: string) => await bookPrompt.ImportGPTPrompt(bookTaskId, txtPath) ) + /** 移除不想要的提示词 */ + ipcMain.handle( + DEFINE_STRING.BOOK.REMOVE_BAD_PROMPT, + async (event, bookTaskId: string) => await bookPrompt.RemoveBadPrompt(bookTaskId) + ) + //#endregion //#region 文案相关 @@ -290,13 +300,13 @@ export function BookIpc() { // 重置小说批次数据 ipcMain.handle( DEFINE_STRING.BOOK.RESET_BOOK_TASK, - async (event, bookTaskId) => await bookTask.ReSetBookTask(bookTaskId) + async (event, bookTaskId, operateBookType) => await bookTask.ReSetBookTask(bookTaskId, operateBookType) ) // 删除小说批次数据 ipcMain.handle( DEFINE_STRING.BOOK.DELETE_BOOK_TASK, - async (event, bookTaskId) => await bookTask.DeleteBookTask(bookTaskId) + async (event, bookTaskId, operateBookType) => await bookTask.DeleteBookTask(bookTaskId, operateBookType) ) @@ -312,5 +322,10 @@ export function BookIpc() { DEFINE_STRING.BOOK.ADD_JIANYING_DRAFT, async (event, id: string | string[], operateBookType: OperateBookType) => await bookVideo.AddJianyingDraft(id, operateBookType) ) + + ipcMain.handle( + DEFINE_STRING.BOOK.ADD_GENERATE_VIDEO_TASK, + async (event, id: string | string[], operateBookType: OperateBookType) => await bookVideo.AddGenerateVideoTask(id, operateBookType) + ) //#endregion } diff --git a/src/main/Service/Book/ReverseBook.ts b/src/main/Service/Book/ReverseBook.ts index 55bb375..2571aec 100644 --- a/src/main/Service/Book/ReverseBook.ts +++ b/src/main/Service/Book/ReverseBook.ts @@ -236,7 +236,7 @@ export class ReverseBook { if (operateBookType == OperateBookType.BOOKTASK) { bookTaskDetails = await this.bookServiceBasic.GetBookTaskDetailData({ bookTaskId: id - }) + }); } else if (operateBookType == OperateBookType.BOOKTASKDETAIL) { let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(id); bookTaskDetails = [bookTaskDetail] @@ -263,7 +263,11 @@ export class ReverseBook { if (bookTaskDetailIds.length <= 0) { throw new Error("没有需要反推的数据，请检查") } - await this.AddReversePromptTask(bookTaskDetailIds, type); + // 判断是不是SD反推并且是反推的全部的数据 + if (type == BookType.SD_REVERSE && operateBookType == OperateBookType.BOOKTASK) { + bookTaskDetailIds = [null]; + } + await this.AddReversePromptTask(bookTaskDetailIds, type, id); return successMessage(null, "添加反推任务成功", 'ReverseBook_AddReversePrompt') } catch (error) { return errorMessage("添加反推任务失败，错误信息如下：" + error.message, "ReverseBook_AddReversePrompt") @@ -276,8 +280,14 @@ export class ReverseBook { * @param type 反推的类型 * @returns */ - async AddReversePromptTask(bookTaskDetailIds: string[], type: BookType): Promise { + async AddReversePromptTask(bookTaskDetailIds: string[], type: BookType, bookTaskId: string): Promise { try { + if (bookTaskDetailIds.length == 1 && bookTaskDetailIds[0] == null && type == BookType.SD_REVERSE) { + let bookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskId); + await this.bookServiceBasic.AddBookBackTask(bookTask.bookId, BookBackTaskType.SD_REVERSE, TaskExecuteType.AUTO, bookTaskId, bookTaskDetailIds[0], DEFINE_STRING.BOOK.REVERSE_PROMPT_RETURN + ); + return; + } for (let index = 0; index < bookTaskDetailIds.length; index++) { const bookTaskDetailId = bookTaskDetailIds[index]; let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(bookTaskDetailId) @@ -327,7 +337,7 @@ export class ReverseBook { res = await this.mjOpt.MJImage2Text(task); break case BookBackTaskType.SD_REVERSE: - res = await this.SDReversePrompt(task); + res = await this.sdOpt.SDReversePrompt(task); break default: throw new Error("未知的任务类型") diff --git a/src/main/Service/Book/bookFrame.ts b/src/main/Service/Book/bookFrame.ts index 502e0d9..e23eb9c 100644 --- a/src/main/Service/Book/bookFrame.ts +++ b/src/main/Service/Book/bookFrame.ts @@ -11,6 +11,10 @@ import { LogScheduler } from '../task/logScheduler'; import { BookBasic } from "./BooKBasic"; import { LoggerStatus, OtherData } from "../../../define/enum/softwareEnum"; import { BasicReverse } from "./basicReverse"; +import JianyingService from "../jianying/jianyingService"; +import { ValidateJson } from "../../../define/Tools/validate"; +import { BookTaskStatus } from "../../../define/enum/bookEnum"; +import { define } from "../../../define/define"; export class BookFrame { bookServiceBasic: BookServiceBasic @@ -18,14 +22,86 @@ export class BookFrame { logScheduler: LogScheduler basicReverse: BasicReverse bookBasic: BookBasic + jianyingService: JianyingService constructor() { this.bookServiceBasic = new BookServiceBasic(); this.ffmpegOptions = new FfmpegOptions(); this.logScheduler = new LogScheduler() this.bookBasic = new BookBasic() this.basicReverse = new BasicReverse() + this.jianyingService = new JianyingService() } + /** + * 剪映抽帧，并判断是不是切割视频 + * @param draftName 草稿名字 + * @param clipVideo 是不是切割视频 + */ + async JianyingFrame(bookTaskId: string, draftName: string) { + try { + console.log(draftName); + let draftPath = path.join(global.config.draft_path, draftName); + if (!await CheckFileOrDirExist(draftPath)) { + throw new Error("未找到对应的草稿文件，请检查"); + } + let bookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskId); + let book = await this.bookServiceBasic.GetBookDataById(bookTask.bookId); + + // 在执行操作之前，先删除之前的文件 + await this.bookServiceBasic.DeleteBookTaskDetailData({ + bookTaskId: bookTask.id + }) + // 在删除对应的图片和视频文件 + let input = path.resolve(book.bookFolderPath, "tmp/input"); + if (await CheckFileOrDirExist(input)) { + await DeleteFolderAllFile(input, true); + } + + await this.jianyingService.GetDraftFrameAndText(draftPath, book.bookFolderPath, 'input'); + let draftTimeLinePath = path.resolve(book.bookFolderPath, "draftFrameData.json"); + if (!await CheckFileOrDirExist(draftTimeLinePath)) { + throw new Error("没有找到对应的草稿时间轴数据，请检查"); + } + let draftTimeLineString = await fs.promises.readFile(draftTimeLinePath, "utf-8"); + if (!ValidateJson(draftTimeLineString)) { + throw new Error("草稿时间轴数据格式错误，请检查"); + } + let draftTimeLine = JSON.parse(draftTimeLineString); + + // 开始修改数据 + // 获取SD设置 + let sdConifg = JSON.parse(await fs.promises.readFile(define.sd_setting, 'utf-8')); + let adetailer = false; + if (sdConifg && sdConifg?.webui?.adetailer) { + adetailer = true; + } + // 写入数据 + for (let i = 0; i < draftTimeLine.length; i++) { + const element = draftTimeLine[i]; + await this.bookServiceBasic.AddBookTaskDetail({ + bookTaskId: bookTaskId, + bookId: bookTask.bookId, + startTime: Math.ceil(element.startTime / 1000), + endTime: Math.ceil(element.endTime / 1000), + status: BookTaskStatus.WAIT, + word: element.word, + videoPath: element.videoPath, + oldImage: path.relative(define.project_path, element.framePath), + afterGpt: element.text, + subValue: undefined, + timeLimit: `${element.startTime} -- ${element.endTime}`, + // 新增修脸跟随 + adetailer: adetailer + }) + } + return successMessage(null, '剪映抽帧成功', 'BookFrame_JianyingFrame') + } catch (error) { + return errorMessage('剪映抽帧失败，失败信息如下：' + error.toString(), 'BookFrame_JianyingFrame') + } + } + + + //#region 替换指定分镜的视频当前帧 /** * 替换指定分镜的视频当前帧 * @param bookTaskDetailId 指定的小说分镜ID @@ -65,6 +141,8 @@ export class BookFrame { return errorMessage('替换指定分镜的视频当前帧失败，失败信息如下：' + error.toString(), 'BookFrame_ReplaceVideoCurrentFrame'); } } + //#endregion + //#region 进行分镜截取的一些操作 diff --git a/src/main/Service/Book/bookGeneral.ts b/src/main/Service/Book/bookGeneral.ts index 0a12f27..ce3e0ed 100644 --- a/src/main/Service/Book/bookGeneral.ts +++ b/src/main/Service/Book/bookGeneral.ts @@ -13,7 +13,7 @@ export class BookGeneral { constructor() { this.bookServiceBasic = new BookServiceBasic() } - + //#region 批量替换数据的操作 /** diff --git a/src/main/Service/Book/bookImage.ts b/src/main/Service/Book/bookImage.ts index 4b785d0..9a1d5ac 100644 --- a/src/main/Service/Book/bookImage.ts +++ b/src/main/Service/Book/bookImage.ts @@ -73,14 +73,6 @@ export class BookImage { for (let i = 0; i < bookTaskDetails.length; i++) { const element = bookTaskDetails[i]; await this.bookServiceBasic.DeleteBoookTaskDetailGenerateImage(element.id); - // if (bookTask.imageCategory == BookImageCategory.MJ) { - // } else if (bookTask.imageCategory == BookImageCategory.D3) { - // throw new Error('暂时不支持D3生成的图片删除') - // } else if (bookTask.imageCategory == BookImageCategory.SD) { - // await this.bookServiceBasic.DeleteBoookTaskDetailGenerateImage(element.id); - // } else { - // throw new Error('未知的小说类型，请检查') - // } // 上面的信息重置完毕之后，开始删除文件信息 let outImage = element.outImagePath; if (await CheckFileOrDirExist(outImage)) { diff --git a/src/main/Service/Book/bookPrompt.ts b/src/main/Service/Book/bookPrompt.ts index 0bf55cc..2f58fec 100644 --- a/src/main/Service/Book/bookPrompt.ts +++ b/src/main/Service/Book/bookPrompt.ts @@ -5,12 +5,14 @@ import { GeneralResponse } from "../../../model/generalResponse"; import { errorMessage, SendReturnMessage, successMessage } from "../../Public/generalTools"; import { BookServiceBasic } from "../ServiceBasic/bookServiceBasic"; import { GptService } from "../GPT/gpt"; -import { ExecuteConcurrently } from "../../../define/Tools/common"; +import { ExecuteConcurrently, ReplaceSubstrings } from "../../../define/Tools/common"; import { DEFINE_STRING } from "../../../define/define_string"; import { CheckFileOrDirExist } from "../../../define/Tools/file"; import fs from 'fs'; import path from 'path' import readline from 'readline'; +import { define } from "../../../define/define"; +import { ValidateJson } from "../../../define/Tools/validate"; export class BookPrompt { @@ -21,6 +23,51 @@ export class BookPrompt { this.gptService = new GptService() } + + //#region 移除不想要的提示词 + /** + * 移除不想要的提示词 + * @param bookTaskId 小说任务的ID + */ + async RemoveBadPrompt(bookTaskId: string) { + try { + let bookTaskDetails = await this.bookServiceBasic.GetBookTaskDetailData({ + bookTaskId: bookTaskId + }) + + let sdConfigString = await fs.promises.readFile(define.sd_setting, 'utf-8') + if (isEmpty(sdConfigString)) { + throw new Error('读取SD配置文件失败') + } + if (!ValidateJson(sdConfigString)) { + throw new Error('SD配置文件格式错误') + } + let sdConfig = JSON.parse(sdConfigString) + let badPrompt = sdConfig.tag.badPrompt; + let badPromptList = badPrompt.split(',') + let result = [] + for (let i = 0; i < bookTaskDetails.length; i++) { + const element = bookTaskDetails[i]; + if (isEmpty(element.gptPrompt)) { + continue; + } + let newText = ReplaceSubstrings(element.gptPrompt, badPromptList, ""); + newText.replaceAll(",,", ","); + await this.bookServiceBasic.UpdateBookTaskDetail(element.id, { + gptPrompt: newText + }) + result.push({ + id: element.id, + gptPrompt: newText + }) + } + return successMessage(result, "移除不想要的提示词成功", "Book_RemoveBadPrompt") + } catch (error) { + return errorMessage("移除不想要的提示词失败，错误信息如下：" + error.message, "Book_RemoveBadPrompt") + } + } + //#endregion + //#region 提示词通用 /** @@ -191,7 +238,13 @@ export class BookPrompt { } } }) - } else if (type == BookType.SD_REVERSE) { } + } else if (type == BookType.SD_REVERSE) { + // SD反推 + for (let i = 0; i < bookTaskDetails.length; i++) { + const element = bookTaskDetails[i]; + await this.bookServiceBasic.UpdateBookTaskDetail(element.id, { gptPrompt: undefined }); + } + } else { throw new Error("SD反推删除还不支持") } diff --git a/src/main/Service/Book/bookTask.ts b/src/main/Service/Book/bookTask.ts index 3fae924..c993f0a 100644 --- a/src/main/Service/Book/bookTask.ts +++ b/src/main/Service/Book/bookTask.ts @@ -233,17 +233,36 @@ export class BookTask { * 重置小说任务数据 * @param bookTaskId 小说任务ID */ - async ReSetBookTask(bookTaskId: string): Promise { + async ReSetBookTask(bookTaskId: string | string[], operateBookType = OperateBookType.BOOKTASK as OperateBookType): Promise { try { - let bookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskId); - await this.bookServiceBasic.ResetBookTask(bookTaskId); - // 数据库删除完毕，看是删除对应的文件（主要是图片） - let imageFolder = bookTask.imageFolder; - // 整个删掉在重建 - if (await CheckFileOrDirExist(imageFolder)) { - await DeleteFolderAllFile(imageFolder) + let bookTasks = [] as Book.SelectBookTask[] + if (operateBookType == OperateBookType.BOOKTASK) { + let tempBookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskId as string); + bookTasks.push(tempBookTask) + } else if (operateBookType == OperateBookType.ASSIGNBOOKTASK) { + for (let i = 0; i < bookTaskId.length; i++) { + const element = bookTaskId[i]; + let tempBookTask = await this.bookServiceBasic.GetBookTaskDataById(element); + bookTasks.push(tempBookTask); + } + } + else { + throw new Error("无效的操作类型") + } + if (bookTasks.length <= 0) { + throw new Error("没有找到对应的小说任务") + } + for (let i = 0; i < bookTasks.length; i++) { + const element = bookTasks[i]; + await this.bookServiceBasic.ResetBookTask(element.id); + // 数据库删除完毕，看是删除对应的文件（主要是图片） + let imageFolder = element.imageFolder; + // 整个删掉在重建 + if (await CheckFileOrDirExist(imageFolder)) { + await DeleteFolderAllFile(imageFolder) + } + await CheckFolderExistsOrCreate(imageFolder) } - await CheckFolderExistsOrCreate(imageFolder) return successMessage(null, "重置小说数据成功", "BookTask_ReSetBookTask") } catch (error) { return errorMessage('重置小说批次数据失败，错误信息如下' + error.toString(), "BookTask_ReSetBookTask"); @@ -254,15 +273,32 @@ export class BookTask { * 删除对应的小说批次数据 * @param bookTaskId 小说批次ID */ - async DeleteBookTask(bookTaskId: string): Promise { + async DeleteBookTask(bookTaskId: string | string[], operateBookType = OperateBookType.BOOKTASK as OperateBookType): Promise { try { - let bookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskId); - let imageFolder = bookTask.imageFolder; - // 先删除每个批次对应的数据，然后删除批次 - await this.bookServiceBasic.DeleteBookTaskData(bookTaskId); - // 删除成功，直接把对应的出图文件夹删掉 - if (await CheckFileOrDirExist(imageFolder)) { - await DeleteFolderAllFile(imageFolder, true) + let bookTasks = [] as Book.SelectBookTask[] + if (operateBookType == OperateBookType.BOOKTASK) { + let tempBookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskId as string); + bookTasks.push(tempBookTask) + } else if (operateBookType == OperateBookType.ASSIGNBOOKTASK) { + for (let i = 0; i < bookTaskId.length; i++) { + const element = bookTaskId[i]; + let tempBookTask = await this.bookServiceBasic.GetBookTaskDataById(element); + bookTasks.push(tempBookTask); + } + } else { + throw new Error("无效的操作类型") + } + + for (let i = 0; i < bookTasks.length; i++) { + const element = bookTasks[i]; + let bookTask = await this.bookServiceBasic.GetBookTaskDataById(element.id); + let imageFolder = bookTask.imageFolder; + // 先删除每个批次对应的数据，然后删除批次 + await this.bookServiceBasic.DeleteBookTaskData(element.id); + // 删除成功，直接把对应的出图文件夹删掉 + if (await CheckFileOrDirExist(imageFolder)) { + await DeleteFolderAllFile(imageFolder, true) + } } return successMessage(null, "删除小说批次数据成功", "BookTask_DeleteBookTask") } catch (error) { diff --git a/src/main/Service/Book/bookVideo.ts b/src/main/Service/Book/bookVideo.ts index 89a71c3..394049c 100644 --- a/src/main/Service/Book/bookVideo.ts +++ b/src/main/Service/Book/bookVideo.ts @@ -1,26 +1,36 @@ -import { OperateBookType } from "../../../define/enum/bookEnum"; +import { BookBackTaskStatus, BookBackTaskType, BookTaskStatus, OperateBookType, TaskExecuteType } from "../../../define/enum/bookEnum"; import { Book } from "../../../model/book"; import { errorMessage, successMessage } from "../../Public/generalTools"; import { GeneralResponse } from "../../../model/generalResponse"; import { Setting } from '../../../main/setting/setting' import path from 'path' -import { CheckFolderExistsOrCreate } from "../../../define/Tools/file"; +import { CheckFileOrDirExist, CheckFolderExistsOrCreate, CopyFileOrFolder, GetFilesWithExtensions } from "../../../define/Tools/file"; import fs from 'fs' import { ClipDraft } from '../../Public/clipDraft' import { BookServiceBasic } from "../ServiceBasic/bookServiceBasic"; import { isEmpty } from "lodash"; import JianyingService from "../jianying/jianyingService"; +import { DEFINE_STRING } from "../../../define/define_string"; +import { define } from "../../../define/define"; +import { ValidateJson } from "../../../define/Tools/validate"; +import BookSetting from "@/main/setting/bookSetting"; +import util from 'util'; +import { spawn, exec } from 'child_process'; +import { SendMessageToRenderer } from "../globalService"; +const execAsync = util.promisify(exec); export class BookVideo { setting: Setting bookServiceBasic: BookServiceBasic jianyingService: JianyingService + bookSetting: BookSetting constructor() { this.setting = new Setting(global) this.bookServiceBasic = new BookServiceBasic() this.jianyingService = new JianyingService() } + //#region 引用主小说相关数据 /** * 将小说中的生成视频的数据，写道小说批次任务中 @@ -29,7 +39,6 @@ export class BookVideo { */ async UseBookVideoDataToBookTask(id: string, operateBookType: OperateBookType) { try { - console.log(id, operateBookType) let book = undefined as Book.SelectBook; let bookTasks = undefined as Book.SelectBookTask[]; if (operateBookType == OperateBookType.BOOK) { @@ -61,7 +70,7 @@ export class BookVideo { // 将修改数据放在一个事务中 for (let i = 0; i < bookTasks.length; i++) { const element = bookTasks[i]; - this.bookServiceBasic.UpdetedBookTaskData(element.id, { + await this.bookServiceBasic.UpdetedBookTaskData(element.id, { backgroundMusic: book.backgroundMusic, friendlyReminder: book.friendlyReminder, draftSrtStyle: book.draftSrtStyle, @@ -81,6 +90,10 @@ export class BookVideo { } } + //#endregion + + //#region 生成配置文件，用于生成草稿或者是合成视频 + /** * 生成配置文件，用于生成草稿或者是合成视频 * @param book 小说数据 @@ -98,7 +111,7 @@ export class BookVideo { } // 开始生成配置文件 - let configPath = path.join(book.bookFolderPath, "scripts/config.json"); + let configPath = path.join(book.bookFolderPath, `scripts/${bookTask.name}_config.json`); await CheckFolderExistsOrCreate(path.dirname(configPath)); // 开始配置 @@ -145,11 +158,16 @@ export class BookVideo { } // 完毕，将数据写出 await fs.promises.writeFile(configPath, JSON.stringify(configData)); + // 复制一个到config.json中 + await CopyFileOrFolder(configPath, path.join(book.bookFolderPath, 'scripts/config.json')); } catch (error) { throw error } } + //#endregion + + //#region 添加剪映草稿 /** * 添加剪映草稿 @@ -159,7 +177,6 @@ export class BookVideo { */ async AddJianyingDraft(id: string | string[], operateBookType: OperateBookType): Promise { try { - console.log(id, operateBookType) let book = undefined as Book.SelectBook let bookTasks = [] as Book.SelectBookTask[] if (operateBookType == OperateBookType.ASSIGNBOOKTASK) { @@ -231,4 +248,211 @@ export class BookVideo { return errorMessage('添加剪映草稿失败，错误信息如下：' + error.toString(), "BookTask_AddJianyingDraft"); } } + + //#endregion + + //#region 添加合成视频任务 + + async AddGenerateVideoTask(id: string | string[], operateBookType: OperateBookType): Promise { + try { + let book = undefined as Book.SelectBook + let bookTasks = [] as Book.SelectBookTask[] + if (operateBookType == OperateBookType.ASSIGNBOOKTASK) { + if (id.length <= 0) { + throw new Error("没有找到小说批次任务，请检查"); + } + for (let i = 0; i < id.length; i++) { + const element = id[i]; + let tempBookTask = await this.bookServiceBasic.GetBookTaskDataById(element as string); + bookTasks.push(tempBookTask) + if (book == undefined) { + book = await this.bookServiceBasic.GetBookDataById(tempBookTask.bookId); + } + // 这边要将主小说的数据写入到选中批次中 + let useRes = await this.UseBookVideoDataToBookTask(element, OperateBookType.BOOKTASK); + if (useRes.code == 0) { + throw new Error(useRes.message); + } + } + + + } else if (operateBookType == OperateBookType.BOOKTASK) { + // 直接获取对应的数据 + let tempBookTask = await this.bookServiceBasic.GetBookTaskDataById(id as string); + if (tempBookTask == null) { + throw new Error("没有找到小说批次任务，请检查"); + } + bookTasks = [tempBookTask] + book = await this.bookServiceBasic.GetBookDataById(tempBookTask.bookId); + if (book == null) { + throw new Error + } + } else { + throw new Error("未知的操作类型"); + } + + if (bookTasks.length <= 0) { + throw new Error("没有找到小说批次任务，请检查") + } + + let result = [] + for (let i = 0; i < bookTasks.length; i++) { + const element = bookTasks[i]; + await this.GenerateConfigFile(book, element); + // 数据处理完毕，开始添加任务 + await this.bookServiceBasic.AddBookBackTask(book.id, BookBackTaskType.COMPOSING, TaskExecuteType.AUTO, element.id, null, DEFINE_STRING.BOOK.GENERATE_VIDEO_RETURN); + // 修改状态 + await this.bookServiceBasic.UpdetedBookTaskData(element.id, { + status: BookTaskStatus.COMPOSING, + }); + result.push(element.name); + } + + // 添加完毕 + return successMessage(result, `${result.join('\n')} ${'\n'} 合成视频任务添加成功`, "BookTask_AddGenerateVideoTask") + } catch (error) { + return errorMessage('添加剪映草稿失败，错误信息如下：' + error.toString(), "BookTask_AddJianyingDraft"); + } + } + //#endregion + + //#region 合成视频任务执行 + /** + * 合成视频任务执行 + * @param task + */ + async GenerateVideo(task: TaskModal.Task): Promise { + try { + console.log(task); + this.bookSetting = new BookSetting(); + let bookTask = await this.bookServiceBasic.GetBookTaskDataById(task.bookTaskId); + let book = await this.bookServiceBasic.GetBookDataById(bookTask.bookId); + if (isEmpty(bookTask.srtPath) || !await CheckFileOrDirExist(bookTask.srtPath)) { + throw new Error("没有找到对应的SRT文件，请检查"); + } + if (isEmpty(bookTask.audioPath) || !await CheckFileOrDirExist(bookTask.audioPath)) { + throw new Error("没有找到对应的音频文件，请检查"); + } + if (isEmpty(bookTask.backgroundMusic)) { + throw new Error("没有找到对应的背景音乐文件夹，请检查"); + } + let videoConfigString = await fs.promises.readFile(define.video_config, 'utf-8'); + if (!ValidateJson(videoConfigString)) { + throw new Error("视频配置文件格式错误，请检查"); + } + + // 检查图片是不是出完了 + let bookTaskDetails = await this.bookServiceBasic.GetBookTaskDetailData({ + bookTaskId: bookTask.id + }); + for (let i = 0; i < bookTaskDetails.length; i++) { + const element = bookTaskDetails[i]; + if (isEmpty(element.outImagePath) || !await CheckFileOrDirExist(element.outImagePath)) { + throw new Error("存在没有出完的图片或者图片不存在，请检查"); + } + } + // 检查分镜信息和出入的图片数量是不是一致 + let outImages = await GetFilesWithExtensions(bookTask.imageFolder, ['.png']); + if (outImages.length != bookTaskDetails.length) { + throw new Error("分镜信息和出图文件夹中图片数量不一致，请检查"); + } + + let videoConfig = JSON.parse(videoConfigString); + let assRandomIndex = Math.floor(Math.random() * videoConfig.assConfig.length) + + let outMp4 = path.join( + book.bookFolderPath, + bookTask.name + '.mp4' + ) + // 将所有的数据写入到配置文件中 + let obj = { + srt_path: bookTask.srtPath, + // 字幕样式（需要随机） + srt_style: isEmpty(bookTask.draftSrtStyle) ? videoConfig.assConfig[assRandomIndex] : videoConfig.assConfig[assRandomIndex], + audio_path: bookTask.audioPath, + background_music_folder: await this.bookSetting.GetBackgroundMusicFolder(bookTask.backgroundMusic), + // 水印设置（需要随机） + friendly_reminder: await this.bookSetting.GetWatermarkSetting(), + video_resolution_x: videoConfig.video_resolution_x, + video_resolution_y: videoConfig.video_resolution_y, + outpue_file: outMp4, + image_folder: bookTask.imageFolder, + srt_config: path.join(book.bookFolderPath, `scripts/${bookTask.name}_config.json`), + mp4_file_txt: path.join(book.bookFolderPath, `scripts/${bookTask.name}.txt`), + status: 'no', + audio_sound_size: videoConfig.audioSoundSize, + background_music_sound_size: videoConfig.backgroundMusicSoundSize, + keyFrame: videoConfig.keyframe, + frameRate: videoConfig.frameRate, + bitRate: videoConfig.bitRate, + stdout: '', + stderr: "" + } + + // 写出配置文件 + let configPath = path.join(book.bookFolderPath, `scripts/${bookTask.name}_video_config.json`); + await fs.promises.writeFile(configPath, JSON.stringify(obj), 'utf-8'); + + let scriptPath = path.join(define.scripts_path, 'Lai.exe') + let gpu = global.gpu.type + if (videoConfig.libx264) { + gpu = 'OTHER' + } + // 执行生成图片的脚本 + let script = `cd "${define.scripts_path}" && "${scriptPath}" -c "${configPath.replaceAll("\\", '/')}" "${gpu}"` + const output = await execAsync(script, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' }) + + if (output.stderr != '') { + obj.status = 'video_error' + obj.stdout = output.stdout + obj.stderr = output.stderr + throw new Error(output.stderr) + } else { + obj.status = 'video_ok' + obj.stdout = output.stdout + obj.stderr = output.stderr + } + + // 检查输出文件是不是存在 + if (!CheckFileOrDirExist(outMp4)) { + throw new Error("生成的视频文件不存在，请检查") + } + + // 修改数据 + await this.bookServiceBasic.UpdetedBookTaskData(task.bookTaskId, { + status: BookTaskStatus.COMPOSING_DONE, + generateVideoPath: path.relative(define.project_path, outMp4) + }) + await this.bookServiceBasic.UpdateTaskStatus({ + id: task.id, + status: BookBackTaskStatus.DONE + }) + SendMessageToRenderer({ + code: 1, + id: task.bookTaskDetailId, + data: BookTaskStatus.COMPOSING_DONE + }, task.messageName); + return successMessage(null, "合成视频任务执行成功", "BookTask_GenerateVideo"); + } catch (error) { + let message = '合成视频任务执行失败，错误信息如下：' + error.toString() + // 失败 + await this.bookServiceBasic.UpdetedBookTaskData(task.bookTaskId, { + status: BookTaskStatus.COMPOSING_FAIL, + errorMsg: error.message + }) + await this.bookServiceBasic.UpdateTaskStatus({ + id: task.id, + status: BookBackTaskStatus.FAIL, + errorMessage: message + }) + SendMessageToRenderer({ + code: 1, + id: task.bookTaskDetailId, + data: BookTaskStatus.COMPOSING_FAIL, + message: message + }, task.messageName); + return errorMessage(message, "BookTask_GenerateVideo"); + } + } + //#endregion } \ No newline at end of file diff --git a/src/main/Service/MJ/mj.ts b/src/main/Service/MJ/mj.ts index 5d931e5..7c95699 100644 --- a/src/main/Service/MJ/mj.ts +++ b/src/main/Service/MJ/mj.ts @@ -401,7 +401,7 @@ export class MJOpt { let characterString = '' // 人物 let characterUrl = '' let sceneString = "" // 场景 - + // 获取当前的自定义风格的垫图字符串 if (book.type == BookType.ORIGINAL) { let sceneIds = element.sceneTags ? element.sceneTags : [] @@ -588,7 +588,6 @@ export class MJOpt { } else { if (task_res.progress == 100) { task_res.type == MJRespoonseType.FINISHED; - console.log(task.id, "22222") await this.bookServiceBasic.UpdateTaskStatus({ id: task.id, status: BookBackTaskStatus.DONE @@ -597,7 +596,6 @@ export class MJOpt { let imagePath = path.join(book.bookFolderPath, `data\\MJOriginalImage\\${task_res.messageId}.png`); await CheckFolderExistsOrCreate(path.dirname(imagePath)) await this.tools.downloadFileUrl(task_res.imageClick, imagePath) - debugger // 进行图片裁剪 let imageRes = await ImageSplit(imagePath, bookTaskDetail.name, path.join(book.bookFolderPath, 'data\\MJOriginalImage')); if (imageRes && imageRes.length < 4) { @@ -606,7 +604,7 @@ export class MJOpt { // 修改数据库数据，将图片保存到对应的文件夹中 let firstImage = imageRes[0]; - if (book.type == BookType.ORIGINAL) { + if (book.type == BookType.ORIGINAL && bookTask.name == "output_00001") { await CopyFileOrFolder(firstImage, path.join(book.bookFolderPath, `tmp\\input\\${bookTaskDetail.name}.png`)); } let out_file = path.join(bookTask.imageFolder, `${bookTaskDetail.name}.png`) diff --git a/src/main/Service/SD/sd.ts b/src/main/Service/SD/sd.ts index a019aa5..029c98f 100644 --- a/src/main/Service/SD/sd.ts +++ b/src/main/Service/SD/sd.ts @@ -4,15 +4,20 @@ import { checkStringValueAddSuffix, errorMessage, successMessage } from "../../P import { define } from '../../../define/define' import fs from "fs"; import { ImageStyle } from "../Book/imageStyle"; -import { BookBackTaskStatus, BookType, OperateBookType } from "../../../define/enum/bookEnum"; +import { BookBackTaskStatus, BookTaskStatus, BookType, OperateBookType } from "../../../define/enum/bookEnum"; import { isEmpty } from "lodash"; import { BookServiceBasic } from "../ServiceBasic/bookServiceBasic"; import { PresetService } from "../presetService"; import path from "path"; import axios from "axios"; -import { CheckFolderExistsOrCreate, CopyFileOrFolder, DeleteFileExifData } from "../../../define/Tools/file"; +import { CheckFileOrDirExist, CheckFolderExistsOrCreate, CopyFileOrFolder, DeleteFileExifData } from "../../../define/Tools/file"; import { Base64ToFile } from "../../../define/Tools/image"; import { MJAction, MJImageType } from "../../../define/enum/mjEnum"; +import { DEFINE_STRING } from "../../../define/define_string"; +import { OtherData, ResponseMessageType } from "../../../define/enum/softwareEnum"; +import util from 'util'; +import { exec } from 'child_process'; +const execAsync = util.promisify(exec) const fspromise = fs.promises export class SDOpt { @@ -35,6 +40,14 @@ export class SDOpt { return sdSetting } + private sendChangeMessage(data: GeneralResponse.MessageResponse, message_name: string = DEFINE_STRING.BOOK.MAIN_DATA_RETURN) { + if (!message_name) { + message_name = DEFINE_STRING.BOOK.MAIN_DATA_RETURN + } + global.newWindow[0].win.webContents.send(message_name, data) + } + + //#region 合并提示词 /** @@ -300,7 +313,7 @@ export class SDOpt { // 这边去图片信息 await DeleteFileExifData(path.join(define.package_path, 'exittool/exiftool.exe'), infoImgPath, imgPath); // 写出去 - if (bookTask.name == 'output_00001') { + if (bookTask.name == 'output_00001' && book.type == BookType.ORIGINAL) { // 复制一个到input let inputImgPath = path.join(inputFolder, `${bookTaskDetail.name}.png`) await CopyFileOrFolder(imgPath, inputImgPath) @@ -359,6 +372,16 @@ export class SDOpt { id: task.bookTaskDetailId }) + await this.bookServiceBasic.UpdateBookTaskDetail(task.bookTaskDetailId, { + status: BookTaskStatus.IMAGE_FAIL, + }) + + await this.bookServiceBasic.UpdateTaskStatus({ + id: task.id, + status: BookBackTaskStatus.FAIL, + errorMessage: errorMsg + }) + global.newWindow[0].win.webContents.send(task.messageName, { code: 0, message: errorMsg, @@ -371,4 +394,126 @@ export class SDOpt { } } //#endregion + + //#region 反推 + + /** + * SD 执行单个反推任务 + * @param task 对应的任务 + * @returns + */ + async SDReversePrompt(task: TaskModal.Task): Promise { + try { + console.log(task) + if (task.bookTaskDetailId == OtherData.DEFAULT) { // 反推整个任务的提示词 + let bookTask = await this.bookServiceBasic.GetBookTaskDataById(task.bookTaskId); + let book = await this.bookServiceBasic.GetBookDataById(bookTask.bookId); + let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailData({ + bookTaskId: task.bookTaskId + }); + let oldImages = [] + for (let i = 0; i < bookTaskDetail.length; i++) { + const element = bookTaskDetail[i]; + oldImages.push({ + id: element.id, + name: element.name, + imagePath: element.oldImage + }) + } + // 判断是不是有为空的 + let imagePaths = [] as string[] + for (let i = 0; i < oldImages.length; i++) { + const element = oldImages[i]; + if (!isEmpty(element.imagePath)) { + imagePaths.push(element.imagePath) + } + } + + if (imagePaths.length == 0) { + throw new Error("没有找到需要反推的图片，请先生成图片") + } + + let inputTxt = path.join(book.bookFolderPath, 'tmp/input/input.txt') + await fspromise.writeFile(inputTxt, imagePaths.join('\n'), 'utf-8') + let exe = path.join(define.scripts_path, 'Lai.exe') + let scripts = `cd "${define.scripts_path}" && ${exe} -pt "${inputTxt}"`; + + // 完成了反推 + // 开始写出数据 + const output = await execAsync(scripts, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' }) + + for (let i = 0; i < oldImages.length; i++) { + const element = oldImages[i]; + let tagTxt = path.join(path.dirname(element.imagePath), `${element.name}.txt`) + if (!await CheckFileOrDirExist(tagTxt)) { + throw new Error("反推失败，没有找到反推出来的提示词文件，请检查") + } + let content = await fspromise.readFile(tagTxt, 'utf-8'); + content = content.replace(/\r?\n/g, ''); + // 更新数据库 + await this.bookServiceBasic.UpdateBookTaskDetail(element.id, { + gptPrompt: content + }); + // 主动返回任务状态 + this.sendChangeMessage({ + code: 1, + type: ResponseMessageType.SD_REVERSE, + id: element.id, + data: content + }, task.messageName); + } + } else { // 反推单个分镜的提示词 + let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(task.bookTaskDetailId); + let imagePath = bookTaskDetail.oldImage; + if (isEmpty(imagePath)) { + throw new Error("当前任务没有生成图片，请先生成图片") + } + if (!await CheckFileOrDirExist(imagePath)) { + throw new Error("当前任务生成图片不存在，请先生成图片") + } + + // 开始执行exe进行反推 + let exe = path.join(define.scripts_path, 'Lai.exe') + let scripts = `cd "${define.scripts_path}" && ${exe} -ps "${imagePath}"`; + + const output = await execAsync(scripts, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' }) + + // 反推成功 + let imageName = path.basename(imagePath, path.extname(imagePath)); + let txt = path.join(path.dirname(imagePath), `${imageName}.txt`) + if (!await CheckFileOrDirExist(txt)) { + throw new Error("反推失败，没有找到反推出来的提示词文件，请检查") + } + let content = await fspromise.readFile(txt, 'utf-8'); + content = content.replace(/\r?\n/g, ''); + + // 更新数据库 + await this.bookServiceBasic.UpdateBookTaskDetail(task.bookTaskDetailId, { + gptPrompt: content + }); + + await this.bookServiceBasic.UpdateTaskStatus({ + id: task.id, + status: BookBackTaskStatus.DONE + }) + // 主动返回任务状态 + this.sendChangeMessage({ + code: 1, + type: ResponseMessageType.SD_REVERSE, + id: task.bookTaskDetailId, + data: content + }, task.messageName); + } + } catch (error) { + let errorMsg = "SD反推失败，错误信息如下：" + error.toString() + await this.bookServiceBasic.UpdateTaskStatus({ + id: task.id, + status: BookBackTaskStatus.FAIL, + errorMessage: errorMsg + }) + return errorMessage(errorMsg, "SDOpt_SDReversePrompt"); + } + } + + //#endregion } \ No newline at end of file diff --git a/src/main/Service/ServiceBasic/bookServiceBasic.ts b/src/main/Service/ServiceBasic/bookServiceBasic.ts index 1479af5..2122454 100644 --- a/src/main/Service/ServiceBasic/bookServiceBasic.ts +++ b/src/main/Service/ServiceBasic/bookServiceBasic.ts @@ -32,8 +32,8 @@ class BookServiceBasic { } //#region 事务操作 - transaction = async (realm : any) => await this.bookBasic.transaction(realm); - + transaction = async (realm: any) => await this.bookBasic.transaction(realm); + // (callback: (realm: any) => void) { // this.bookService.transaction(() => { // callback(this.bookService.realm) @@ -64,16 +64,16 @@ class BookServiceBasic { //#region 分镜任务 - AddBookTaskDetail = async (bookTaskDetail: Book.SelectBookTaskDetail) => await this.bookTaskDetailServiceBasic.AddBookTaskDetail(bookTaskDetail); GetBookTaskDetailDataById = async (bookTaskDetailId: string) => await this.bookTaskDetailServiceBasic.GetBookTaskDetailDataById(bookTaskDetailId); GetBookTaskDetailData = async (condition: Book.QueryBookTaskDetailCondition, returnEmpty: boolean = false) => await this.bookTaskDetailServiceBasic.GetBookTaskDetailData(condition, returnEmpty); UpdateBookTaskDetail = async (bookTaskDetailId: string, data: Book.SelectBookTaskDetail) => await this.bookTaskDetailServiceBasic.UpdateBookTaskDetail(bookTaskDetailId, data); - UpdateBookTaskStatus = async (bookTaskDetailId: string, status: BookTaskStatus,errorMsg? : string) => await this.bookTaskDetailServiceBasic.UpdateBookTaskStatus(bookTaskDetailId, status,errorMsg); - DeleteBookTaskDetailReversePromptById = async (bookTaskDetailId: string, reversePromptId: string) => await this.bookTaskDetailServiceBasic.DeleteBookTaskDetailReversePromptById(bookTaskDetailId); + UpdateBookTaskStatus = async (bookTaskDetailId: string, status: BookTaskStatus, errorMsg?: string) => await this.bookTaskDetailServiceBasic.UpdateBookTaskStatus(bookTaskDetailId, status, errorMsg); + DeleteBookTaskDetailReversePromptById = async (bookTaskDetailId: string) => await this.bookTaskDetailServiceBasic.DeleteBookTaskDetailReversePromptById(bookTaskDetailId); DeleteBoookTaskDetailGenerateImage = async (bookTaskDetailId: string) => await this.bookTaskDetailServiceBasic.DeleteBoookTaskDetailGenerateImage(bookTaskDetailId); UpdateBookTaskDetailReversePrompt = async (bookTaskDetailId: string, reversePromptId: string, data: Book.ReversePrompt) => await this.bookTaskDetailServiceBasic.UpdateBookTaskDetailReversePrompt(bookTaskDetailId, reversePromptId, data); UpdateBookTaskDetailMjMessage = async (bookTaskDetailId: string, mjMessage: Book.MJMessage) => await this.bookTaskDetailServiceBasic.UpdateBookTaskDetailMjMessage(bookTaskDetailId, mjMessage); + DeleteBookTaskDetailData = async (condition: Book.DeleteBookTaskDetailCondition) => await this.bookTaskDetailServiceBasic.DeleteBookTaskDetail(condition); //#endregion diff --git a/src/main/Service/ServiceBasic/bookTaskDetailServiceBasic.ts b/src/main/Service/ServiceBasic/bookTaskDetailServiceBasic.ts index 677374d..b47c8b7 100644 --- a/src/main/Service/ServiceBasic/bookTaskDetailServiceBasic.ts +++ b/src/main/Service/ServiceBasic/bookTaskDetailServiceBasic.ts @@ -89,6 +89,16 @@ export default class BookTaskDetailServiceBasic { this.bookTaskDetailService.DeleteBookTaskDetailReversePromptById(bookTaskDetailId) } + /** + * 删除指定的小说分镜数据 + * @param condition + */ + async DeleteBookTaskDetail(condition: Book.QueryBookTaskDetailCondition): Promise { + await this.InitService(); + this.bookTaskDetailService.DeleteBookTaskDetail(condition) + } + + /** * 删除指定的小说分镜生成的图片 * @param bookTaskDetailId diff --git a/src/main/Service/Subtitle/subtitleService.ts b/src/main/Service/Subtitle/subtitleService.ts index 3144c2d..e400eaa 100644 --- a/src/main/Service/Subtitle/subtitleService.ts +++ b/src/main/Service/Subtitle/subtitleService.ts @@ -354,7 +354,6 @@ export class SubtitleService { bookTaskId: bookTaskId }, true) - debugger // 获取SD设置 let sdConifg = JSON.parse(await fs.promises.readFile(define.sd_setting, 'utf-8')); let adetailer = false; diff --git a/src/main/Service/globalService.ts b/src/main/Service/globalService.ts new file mode 100644 index 0000000..d64f40d --- /dev/null +++ b/src/main/Service/globalService.ts @@ -0,0 +1,14 @@ +import { DEFINE_STRING } from "@/define/define_string" +import { GeneralResponse } from "@/model/generalResponse" + +/** + * 主动返回前端的消息 + * @param data 返回的数据 + * @param message_name 消息名称 + */ +export function SendMessageToRenderer(data: GeneralResponse.MessageResponse, message_name: string = DEFINE_STRING.BOOK.MAIN_DATA_RETURN) { + if (!message_name) { + message_name = DEFINE_STRING.BOOK.MAIN_DATA_RETURN + } + global.newWindow[0].win.webContents.send(message_name, data) +} \ No newline at end of file diff --git a/src/main/Service/jianying/jianyingService.ts b/src/main/Service/jianying/jianyingService.ts index 355d19e..1c6f8fc 100644 --- a/src/main/Service/jianying/jianyingService.ts +++ b/src/main/Service/jianying/jianyingService.ts @@ -36,9 +36,8 @@ class JianyingService { * 获取剪映草稿的关键帧和文本 * @param draftDir 草稿目录 * @param projectDir 项目目录 - * @param packagePath 包路径 */ - public async GetDraftFrameAndText(draftDir: string, projectDir: string, packagePath: string) { + public async GetDraftFrameAndText(draftDir: string, projectDir: string, inputName: string = "input_crop"): Promise { try { // 获取草稿文件路径 let draftJsonPath = path.resolve(draftDir, "draft_content.json"); @@ -59,7 +58,7 @@ class JianyingService { await DeleteFolderAllFile(projectTmp); } // 创建输出文件夹 - let projectInput = path.resolve(projectTmp, "input_crop"); + let projectInput = path.resolve(projectTmp, inputName); console.log("projectInput", projectInput); // 获取剪映的轨道，并且判断是否包含一个video轨道和一个text轨道 @@ -98,6 +97,8 @@ class JianyingService { console.log("GetDraftFrameAndText", jsonPath); } catch (error) { throw error; + } finally { + this.draftTimeLine = []; } } diff --git a/src/main/Service/task/taskManage.ts b/src/main/Service/task/taskManage.ts index 75e2619..1bfaca4 100644 --- a/src/main/Service/task/taskManage.ts +++ b/src/main/Service/task/taskManage.ts @@ -13,6 +13,7 @@ import { FluxOpt } from '../Flux/flux' import { AsyncQueue } from '../../quene' import { SoftWareServiceBasic } from '../ServiceBasic/softwareServiceBasic' import { MJSetting } from '../../../model/Setting/mjSetting' +import { BookVideo } from '../Book/bookVideo' export class TaskManager { isExecuting: boolean = false; @@ -24,6 +25,7 @@ export class TaskManager { bookBackTaskListService!: BookBackTaskListService; eventListeners: Record = {}; softWareServiceBasic: SoftWareServiceBasic + bookVideo: BookVideo mjSetting: MJSetting.MjSetting spaceTime: number = 5000; @@ -252,6 +254,15 @@ export class TaskManager { }, `${batch}_${task.id}`, batch, `${batch}_${task.id}_${new Date().getTime()}`) } + /** 添加生成视频的后台任务 */ + async AddGenerateVideo(task: TaskModal.Task) { + let batch = task.messageName; + global.requestQuene.enqueue(async () => { + this.bookVideo = new BookVideo(); + await this.bookVideo.GenerateVideo(task); + }, `${batch}_${task.id}`, batch, `${batch}_${task.id}_${new Date().getTime()}`) + } + /** * 添加 flux forge 任务到内存队列中 @@ -299,7 +310,8 @@ export class TaskManager { case BookBackTaskType.RECOGNIZE: this.AddExtractSubtitlesData(task); break; - case BookBackTaskType.MJ_REVERSE || BookBackTaskType.SD_REVERSE: + case BookBackTaskType.MJ_REVERSE: + case BookBackTaskType.SD_REVERSE: this.AddSingleReversePrompt(task); break; case BookBackTaskType.FLUX_FORGE_IMAGE: @@ -317,6 +329,10 @@ export class TaskManager { case BookBackTaskType.D3_IMAGE: this.AddD3Image(task); break; + + case BookBackTaskType.COMPOSING: // 合成视频 + this.AddGenerateVideo(task); + break; default: throw new Error('未知的任务类型'); } diff --git a/src/main/func.js b/src/main/func.js index 2f3dcbf..52a5a5e 100644 --- a/src/main/func.js +++ b/src/main/func.js @@ -506,7 +506,7 @@ async function getFrame(value) { let package_path = define.package_path.replaceAll('\\', '/') let jianying = new JianyingService() - await jianying.GetDraftFrameAndText(draft_path, out_dir, package_path) + await jianying.GetDraftFrameAndText(draft_path, out_dir) global.newWindow[0].win.webContents.send(DEFINE_STRING.SHOW_MESSAGE_DIALOG, { code: 1, @@ -539,7 +539,7 @@ async function PushBackPrompt() { let script = `cd "${define.scripts_path}" && "${py_path}" -p "${sd_config_path.replaceAll( '\\', '/' - )}" "input" "${global.config.project_path}"` + )}" "input_crop" "${global.config.project_path}"` const output = await execAsync(script, { maxBuffer: 1024 * 1024 * 10, encoding: 'utf-8' }) global.newWindow[0].win.webContents.send(DEFINE_STRING.SHOW_MESSAGE_DIALOG, { code: 1, diff --git a/src/main/setting/bookSetting.js b/src/main/setting/bookSetting.js deleted file mode 100644 index 89b52a3..0000000 --- a/src/main/setting/bookSetting.js +++ /dev/null @@ -1,5 +0,0 @@ - -export class BookSetting { - constructor() { - } -} \ No newline at end of file diff --git a/src/main/setting/bookSetting.ts b/src/main/setting/bookSetting.ts new file mode 100644 index 0000000..acf7052 --- /dev/null +++ b/src/main/setting/bookSetting.ts @@ -0,0 +1,57 @@ +import { define } from "@/define/define"; +import { CheckFileOrDirExist } from "@/define/Tools/file"; +import { ValidateJson } from "@/define/Tools/validate"; +import fs from 'fs'; + +class BookSetting { + constructor() { } + + /** + * 通过ID获取设置的背景音乐文件夹 + * @param id 背景音乐文件夹ID + */ + async GetBackgroundMusicFolder(id: string): Promise { + let clipSettingString = await fs.promises.readFile(define.clip_setting, 'utf-8'); + if (!ValidateJson(clipSettingString)) { + throw new Error('Clip设置文件格式错误'); + } + let clipSetting = JSON.parse(clipSettingString); + let backgroundMusicSetting = clipSetting.background_music_setting; + if (!backgroundMusicSetting) { + throw new Error('没有找到对应的背景音乐文件夹地址设置，请检查'); + } + let clip = backgroundMusicSetting.find((clip: any) => clip.id === id); + if (clip == undefined) { + throw new Error('没有找到对应的背景音乐文件夹地址，请检查'); + } + if (!await CheckFileOrDirExist(clip.folder_path)) { + throw new Error('背景音乐文件夹不存在，请检查'); + } + return clip.folder_path; + } + + /** + * 获取指定的水印设置，没有的话，随机返回一个 + * @param id 水印设置ID + */ + async GetWatermarkSetting(id?: string) { + let videoSettingString = await fs.promises.readFile(define.video_config, 'utf-8'); + if (!ValidateJson(videoSettingString)) { + throw new Error('Video设置文件格式错误'); + } + let videoSetting = JSON.parse(videoSettingString); + let watermarkConfig = videoSetting.watermarkConfig; + if (!watermarkConfig) { + throw new Error('没有找到对应的水印设置，请检查'); + } + if (id == undefined) { + return watermarkConfig[Math.floor(Math.random() * watermarkConfig.length)]; + } + let watermark = watermarkConfig.find((watermark: any) => watermark.id === id); + if (watermark == undefined) { + throw new Error('没有找到对应的水印设置，请检查'); + } + return watermark; + } +} +export default BookSetting; \ No newline at end of file diff --git a/src/model/book.d.ts b/src/model/book.d.ts index 4fdf4c7..8873687 100644 --- a/src/model/book.d.ts +++ b/src/model/book.d.ts @@ -272,6 +272,13 @@ declare namespace Book { duration: number } + type DeleteBookTaskDetailCondition = { + id?: string; + bookId?: string; + bookTaskId?: string; + name?: string; + } + //#region 替换小说数据 type BookReplaceData = { diff --git a/src/preload/book.ts b/src/preload/book.ts index d1319bc..81a4e35 100644 --- a/src/preload/book.ts +++ b/src/preload/book.ts @@ -54,7 +54,7 @@ const book = { //#endregion - //#region 小说通用操作 + //#region bookTaskDetail通用操作 // 替换文案数据 ReplaceBookData: async (bookTaskId: string, replaceData: Book.BookReplaceData) => @@ -82,6 +82,9 @@ const book = { bookTaskDetailId, currentTime ), + + /** 剪映抽帧，并判断是不是切割视频 */ + JianyingFrame: async (bookTaskId: string, draftName: string) => await ipcRenderer.invoke(DEFINE_STRING.BOOK.JIANYING_FRAME, bookTaskId, draftName), //#endregion //#region 文案相关信息 @@ -160,6 +163,10 @@ const book = { // 导入提示词 ImportGPTPrompt: async (bookTaskId: string, txtPath: string) => await ipcRenderer.invoke(DEFINE_STRING.BOOK.IMPORT_GPT_PROMPT, bookTaskId, txtPath), + + /** 移除不想要的提示词 */ + RemoveBadPrompt: async (bookTaskId: string) => await ipcRenderer.invoke(DEFINE_STRING.BOOK.REMOVE_BAD_PROMPT, bookTaskId), + //#endregion //#region 图片相关 @@ -227,12 +234,12 @@ const book = { await ipcRenderer.invoke(DEFINE_STRING.BOOK.ADD_NEW_BOOK_TASK, addBookTaskData), // 重置小说批次数据 - ReSetBookTask: async (bookTaskId) => - await ipcRenderer.invoke(DEFINE_STRING.BOOK.RESET_BOOK_TASK, bookTaskId), + ReSetBookTask: async (bookTaskId: string | string[], operateBookType: OperateBookType) => + await ipcRenderer.invoke(DEFINE_STRING.BOOK.RESET_BOOK_TASK, bookTaskId, operateBookType), // 删除小说批次数据 - DeleteBookTask: async (bookTaskId) => - await ipcRenderer.invoke(DEFINE_STRING.BOOK.DELETE_BOOK_TASK, bookTaskId), + DeleteBookTask: async (bookTaskId: string | string[], operateBookType: OperateBookType) => + await ipcRenderer.invoke(DEFINE_STRING.BOOK.DELETE_BOOK_TASK, bookTaskId, operateBookType), // 一键生成所有图片 GenerateImageAll: async (bookTaskId, imageCategory) => @@ -261,8 +268,11 @@ const book = { // 添加数据到剪映草稿 AddJianyingDraft: async (id: string | string[], operateBookType: OperateBookType) => - await ipcRenderer.invoke(DEFINE_STRING.BOOK.ADD_JIANYING_DRAFT, id, operateBookType) + await ipcRenderer.invoke(DEFINE_STRING.BOOK.ADD_JIANYING_DRAFT, id, operateBookType), + /** 合成视频相关 */ + AddGenerateVideoTask: async (id: string | string[], operateBookType: OperateBookType) => + await ipcRenderer.invoke(DEFINE_STRING.BOOK.ADD_GENERATE_VIDEO_TASK, id, operateBookType) //#endregion } export { book } diff --git a/src/renderer/src/components/Book/Components/DatatableGptPromptButton.vue b/src/renderer/src/components/Book/Components/DatatableGptPromptButton.vue deleted file mode 100644 index 92acf01..0000000 --- a/src/renderer/src/components/Book/Components/DatatableGptPromptButton.vue +++ /dev/null @@ -1,300 +0,0 @@ - - - diff --git a/src/renderer/src/components/Book/Components/DatatableHeaderGptPrompt.vue b/src/renderer/src/components/Book/Components/DatatableHeaderGptPrompt.vue deleted file mode 100644 index ee2e331..0000000 --- a/src/renderer/src/components/Book/Components/DatatableHeaderGptPrompt.vue +++ /dev/null @@ -1,165 +0,0 @@ - - - diff --git a/src/renderer/src/components/Book/Components/DatatablePrompt.vue b/src/renderer/src/components/Book/Components/DatatablePrompt.vue deleted file mode 100644 index 8eefb3f..0000000 --- a/src/renderer/src/components/Book/Components/DatatablePrompt.vue +++ /dev/null @@ -1,215 +0,0 @@ - - - diff --git a/src/renderer/src/components/Book/Components/ManageBook/BookTaskListAction.vue b/src/renderer/src/components/Book/Components/ManageBook/BookTaskListAction.vue index a0d6352..a5c34b6 100644 --- a/src/renderer/src/components/Book/Components/ManageBook/BookTaskListAction.vue +++ b/src/renderer/src/components/Book/Components/ManageBook/BookTaskListAction.vue @@ -33,7 +33,7 @@ - diff --git a/src/renderer/src/components/Book/Components/ManageBook/ManageBookTaskGenerateInformation.vue b/src/renderer/src/components/Book/Components/ManageBook/ManageBookTaskGenerateInformation.vue index 7692175..eb273e0 100644 --- a/src/renderer/src/components/Book/Components/ManageBook/ManageBookTaskGenerateInformation.vue +++ b/src/renderer/src/components/Book/Components/ManageBook/ManageBookTaskGenerateInformation.vue @@ -46,7 +46,7 @@ clearable /> - +
注意：选择的草稿主轨道的图片数量要和当前的相同，会生成新的草稿并且只会替换图片，保留其余的数据 @@ -67,12 +67,23 @@ > 保存数据 -
注意：在生成草稿前要先保存数据
+
+ 注意：在 生成草稿/合成视频 前要先保存数据 +
- 注意：在生成草稿前要先保存数据，当前会生成选择的草稿，全部会使用上面的参数 + 注意：在 生成草稿/合成视频 前要先保存数据，当前会生成选择的 草稿/视频，全部会使用上面的参数
- 生成草稿 + 添加合成视频任务 + 生成草稿 @@ -81,18 +92,21 @@ import { ref, onMounted } from 'vue' import { useMessage, NForm, NFormItem, NButton, NIcon, NInput, NSelect } from 'naive-ui' import { FolderOpen } from '@vicons/ionicons5' -import { OperateBookType } from '../../../../../../define/enum/bookEnum' +import { BookTaskStatus, OperateBookType } from '../../../../../../define/enum/bookEnum' import { useReverseManageStore } from '../../../../../../stores/reverseManage' let props = defineProps({ bookTask: undefined, type: undefined, - selectBookTask: [] + selectBookTask: [], + optionType: undefined }) let bookTask = ref(props.bookTask) let type = ref(props.type) +let optionType = ref(props.optionType) + let backgroundMusicOptions = ref([]) let message = useMessage() let draftSelect = ref([]) @@ -217,23 +231,78 @@ async function SaveVideoData() { } } -/** - * 添加草稿 - */ -async function AddJianyingDraft() { +async function AddGenerateVideoTask() { + let res = undefined if (props.type == 'book') { - let res = await window.book.AddJianyingDraft( + // 单个 + res = await window.book.AddGenerateVideoTask( props.selectBookTask, OperateBookType.ASSIGNBOOKTASK ) - window.api.showGlobalMessageDialog(res) + if (res.code == 1) { + reverseManageStore.UpdatedBookTaskData(props.selectBookTask, { + status: BookTaskStatus.COMPOSING + }) + } else { + reverseManageStore.UpdatedBookTaskData(props.selectBookTask, { + status: BookTaskStatus.COMPOSING_FAIL, + errorMsg: res.message + }) + } } else if (props.type == 'bookTask') { - let res = await window.book.AddJianyingDraft(bookTask.value.id, OperateBookType.BOOKTASK) - window.api.showGlobalMessageDialog(res) + // 批次 + res = await window.book.AddGenerateVideoTask(bookTask.value.id, OperateBookType.BOOKTASK) + if (res.code == 1) { + reverseManageStore.UpdatedBookTaskData(bookTask.value.id, { + status: BookTaskStatus.COMPOSING + }) + } else { + reverseManageStore.UpdatedBookTaskData(bookTask.value.id, { + status: BookTaskStatus.COMPOSING_FAIL, + errorMsg: res.message + }) + } } else { message.error('未知的操作类型，请检查') return } + window.api.showGlobalMessageDialog(res) +} + +/** + * 添加草稿 + */ +async function AddJianyingDraft() { + let res = undefined + if (props.type == 'book') { + res = await window.book.AddJianyingDraft(props.selectBookTask, OperateBookType.ASSIGNBOOKTASK) + if (res.code == 1) { + reverseManageStore.UpdatedBookTaskData(props.selectBookTask, { + status: BookTaskStatus.DRAFT_DONE + }) + } else { + reverseManageStore.UpdatedBookTaskData(props.selectBookTask, { + status: BookTaskStatus.DRAFT_FAIL, + errorMsg: res.message + }) + } + } else if (props.type == 'bookTask') { + res = await window.book.AddJianyingDraft(bookTask.value.id, OperateBookType.BOOKTASK) + if (res.code == 1) { + reverseManageStore.UpdatedBookTaskData(bookTask.value.id, { + status: BookTaskStatus.DRAFT_DONE + }) + } else { + reverseManageStore.UpdatedBookTaskData(bookTask.value.id, { + status: BookTaskStatus.DRAFT_FAIL, + errorMsg: res.message + }) + } + } else { + message.error('未知的操作类型，请检查') + return + } + window.api.showGlobalMessageDialog(res) } let rules = ref({ srtPath: [{ required: true, message: '请选择背景音乐', trigger: ['input', 'blur', 'change'] }], diff --git a/src/renderer/src/components/Book/Components/ManageBookDetailButton.vue b/src/renderer/src/components/Book/Components/ManageBookDetailButton.vue deleted file mode 100644 index 1bfc1d7..0000000 --- a/src/renderer/src/components/Book/Components/ManageBookDetailButton.vue +++ /dev/null @@ -1,960 +0,0 @@ - - - - - diff --git a/src/renderer/src/components/Book/Components/DatatableAfterGpt.vue b/src/renderer/src/components/Book/MJReverse/DatatableAfterGpt.vue similarity index 96% rename from src/renderer/src/components/Book/Components/DatatableAfterGpt.vue rename to src/renderer/src/components/Book/MJReverse/DatatableAfterGpt.vue index d50f47a..34bf468 100644 --- a/src/renderer/src/components/Book/Components/DatatableAfterGpt.vue +++ b/src/renderer/src/components/Book/MJReverse/DatatableAfterGpt.vue @@ -35,7 +35,6 @@ export default defineComponent({ props: ['initData', 'index'], setup(props) { let data = ref(props.initData) - console.log('data', props.initData, props.index) onMounted(async () => {}) let InputDebounced = debounce(handleInput, 1000) diff --git a/src/renderer/src/components/Book/Components/DatatableGenerateImage.vue b/src/renderer/src/components/Book/MJReverse/DatatableGenerateImage.vue similarity index 100% rename from src/renderer/src/components/Book/Components/DatatableGenerateImage.vue rename to src/renderer/src/components/Book/MJReverse/DatatableGenerateImage.vue diff --git a/src/renderer/src/components/Book/MJReverse/DatatableGptPromptButton.vue b/src/renderer/src/components/Book/MJReverse/DatatableGptPromptButton.vue new file mode 100644 index 0000000..224db51 --- /dev/null +++ b/src/renderer/src/components/Book/MJReverse/DatatableGptPromptButton.vue @@ -0,0 +1,284 @@ + + + diff --git a/src/renderer/src/components/Book/Components/DatatableHeaderAfterGpt.vue b/src/renderer/src/components/Book/MJReverse/DatatableHeaderAfterGpt.vue similarity index 100% rename from src/renderer/src/components/Book/Components/DatatableHeaderAfterGpt.vue rename to src/renderer/src/components/Book/MJReverse/DatatableHeaderAfterGpt.vue diff --git a/src/renderer/src/components/Book/Components/DatatableHeaderAfterGptSelectAndReplace.vue b/src/renderer/src/components/Book/MJReverse/DatatableHeaderAfterGptSelectAndReplace.vue similarity index 100% rename from src/renderer/src/components/Book/Components/DatatableHeaderAfterGptSelectAndReplace.vue rename to src/renderer/src/components/Book/MJReverse/DatatableHeaderAfterGptSelectAndReplace.vue diff --git a/src/renderer/src/components/Book/Components/DatatableHeaderGenerateImage.vue b/src/renderer/src/components/Book/MJReverse/DatatableHeaderGenerateImage.vue similarity index 100% rename from src/renderer/src/components/Book/Components/DatatableHeaderGenerateImage.vue rename to src/renderer/src/components/Book/MJReverse/DatatableHeaderGenerateImage.vue diff --git a/src/renderer/src/components/Book/MJReverse/DatatableHeaderGptPrompt.vue b/src/renderer/src/components/Book/MJReverse/DatatableHeaderGptPrompt.vue new file mode 100644 index 0000000..289f65e --- /dev/null +++ b/src/renderer/src/components/Book/MJReverse/DatatableHeaderGptPrompt.vue @@ -0,0 +1,235 @@ + + + diff --git a/src/renderer/src/components/Book/Components/DatatableHeaderImage.vue b/src/renderer/src/components/Book/MJReverse/DatatableHeaderImage.vue similarity index 88% rename from src/renderer/src/components/Book/Components/DatatableHeaderImage.vue rename to src/renderer/src/components/Book/MJReverse/DatatableHeaderImage.vue index 867fbf4..7edd42d 100644 --- a/src/renderer/src/components/Book/Components/DatatableHeaderImage.vue +++ b/src/renderer/src/components/Book/MJReverse/DatatableHeaderImage.vue @@ -2,7 +2,7 @@
出图
{}) + let options = ref([]) + + onMounted(async () => { + // 获取生图下拉列表的方式 + await window.api.GetImageGenerateCategory((value) => { + if (value.code == 0) { + message.error(value.message) + return + } + options.value = value.data + }) + }) async function handleUpdateValue(value) { - if (value == BookImageCategory.D3) { - message.error('暂不支持该出图模式') - return - } // 将值的修改保存进数据库中 let res = await window.db.UpdateBookTaskData(reverseManageStore.selectBookTask.id, { imageCategory: value @@ -124,20 +131,7 @@ export default defineComponent({ key: 'unlock' } ], - options: [ - { - label: 'MJ', - value: BookImageCategory.MJ - }, - { - label: 'SD', - value: BookImageCategory.SD - }, - { - label: 'D3', - value: BookImageCategory.D3 - } - ] + options } } }) diff --git a/src/renderer/src/components/Book/Components/DatatableHeaderPrompt.vue b/src/renderer/src/components/Book/MJReverse/DatatableHeaderPrompt.vue similarity index 100% rename from src/renderer/src/components/Book/Components/DatatableHeaderPrompt.vue rename to src/renderer/src/components/Book/MJReverse/DatatableHeaderPrompt.vue diff --git a/src/renderer/src/components/Book/MJReverse/DatatablePrompt.vue b/src/renderer/src/components/Book/MJReverse/DatatablePrompt.vue new file mode 100644 index 0000000..372a408 --- /dev/null +++ b/src/renderer/src/components/Book/MJReverse/DatatablePrompt.vue @@ -0,0 +1,305 @@ + + + diff --git a/src/renderer/src/components/Book/Components/GetVideoFrame.vue b/src/renderer/src/components/Book/MJReverse/GetVideoFrame.vue similarity index 100% rename from src/renderer/src/components/Book/Components/GetVideoFrame.vue rename to src/renderer/src/components/Book/MJReverse/GetVideoFrame.vue diff --git a/src/renderer/src/components/Book/MJReverse/JianyingFrameOption.vue b/src/renderer/src/components/Book/MJReverse/JianyingFrameOption.vue new file mode 100644 index 0000000..e2886ea --- /dev/null +++ b/src/renderer/src/components/Book/MJReverse/JianyingFrameOption.vue @@ -0,0 +1,110 @@ + + + diff --git a/src/renderer/src/components/Book/MJReverse/MJReversePrompt.vue b/src/renderer/src/components/Book/MJReverse/MJReversePrompt.vue index 441574e..ced1234 100644 --- a/src/renderer/src/components/Book/MJReverse/MJReversePrompt.vue +++ b/src/renderer/src/components/Book/MJReverse/MJReversePrompt.vue @@ -37,7 +37,7 @@ import { ref, onMounted, defineComponent, onUnmounted, toRaw, computed } from 'vue' import { useMessage, NButton, NInput, NSpin, NDropdown } from 'naive-ui' import { useReverseManageStore } from '../../../../../stores/reverseManage' -import DatatableGptPromptButton from '../Components/DatatableGptPromptButton.vue' +import DatatableGptPromptButton from './DatatableGptPromptButton.vue' import { debounce } from 'lodash' export default defineComponent({ diff --git a/src/renderer/src/components/Book/MJReverse/ManageBookDetailButton.vue b/src/renderer/src/components/Book/MJReverse/ManageBookDetailButton.vue new file mode 100644 index 0000000..ea8663b --- /dev/null +++ b/src/renderer/src/components/Book/MJReverse/ManageBookDetailButton.vue @@ -0,0 +1,1007 @@ + + + + + diff --git a/src/renderer/src/components/Book/Components/ManageBookOldImage.vue b/src/renderer/src/components/Book/MJReverse/ManageBookOldImage.vue similarity index 100% rename from src/renderer/src/components/Book/Components/ManageBookOldImage.vue rename to src/renderer/src/components/Book/MJReverse/ManageBookOldImage.vue diff --git a/src/renderer/src/components/Book/MJReverse/ManageBookReverseTable.vue b/src/renderer/src/components/Book/MJReverse/ManageBookReverseTable.vue index f666811..77bb288 100644 --- a/src/renderer/src/components/Book/MJReverse/ManageBookReverseTable.vue +++ b/src/renderer/src/components/Book/MJReverse/ManageBookReverseTable.vue @@ -10,164 +10,162 @@
- diff --git a/src/renderer/src/components/Book/Components/MonitorStatus.vue b/src/renderer/src/components/Book/MJReverse/MonitorStatus.vue similarity index 100% rename from src/renderer/src/components/Book/Components/MonitorStatus.vue rename to src/renderer/src/components/Book/MJReverse/MonitorStatus.vue diff --git a/src/renderer/src/components/Book/MJReverse/SDReversePrompt.vue b/src/renderer/src/components/Book/MJReverse/SDReversePrompt.vue new file mode 100644 index 0000000..402a929 --- /dev/null +++ b/src/renderer/src/components/Book/MJReverse/SDReversePrompt.vue @@ -0,0 +1,44 @@ + + + diff --git a/src/renderer/src/components/Book/Components/Setting.vue b/src/renderer/src/components/Book/MJReverse/Setting.vue similarity index 100% rename from src/renderer/src/components/Book/Components/Setting.vue rename to src/renderer/src/components/Book/MJReverse/Setting.vue diff --git a/src/renderer/src/components/Book/ManageBookDetail.vue b/src/renderer/src/components/Book/ManageBookDetail.vue index 51dd7ce..db056d1 100644 --- a/src/renderer/src/components/Book/ManageBookDetail.vue +++ b/src/renderer/src/components/Book/ManageBookDetail.vue @@ -1,7 +1,12 @@ -