fix: better trigger handle for zh prompt

LeafYeeXYZ · Dec 1, 2024 · 9f1c228 · 9f1c228
1 parent 0c553f7
commit 9f1c228
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 17 deletions.
diff --git a/app/components/Prompt.tsx b/app/components/Prompt.tsx
@@ -1,6 +1,6 @@
 'use client'
 
-import { Models, type Model } from '../lib/config'
+import { Models } from '../lib/config'
 import { Form, Button, Select, Input, Space, Upload, Popover } from 'antd'
 import { FileImageOutlined, FileAddOutlined } from '@ant-design/icons'
 import { useState } from 'react'
@@ -11,7 +11,7 @@ import { useZustand } from '../lib/useZustand'
 
 type FormValues = {
   prompt: string
-  model: Model['value']
+  model: string
 }
 
 export default function Prompt() {
@@ -22,7 +22,8 @@ export default function Prompt() {
     flushSync(() => setDisabled(true))
     const model = Models.find((m) => m.value === value.model)!
     const task = {
-      prompt: model.trigger ? `${model.trigger}, ${value.prompt}` : value.prompt,
+      prompt: value.prompt,
+      trigger: model.trigger,
       model: model.value,
       promptLanguage: getPromptLanguage(),
       status: 'waiting' as Task['status'],

diff --git a/app/lib/config.ts b/app/lib/config.ts
@@ -1,11 +1,6 @@
 import type { ThemeConfig } from 'antd'
 import { theme } from 'antd'
-
-export type Model = {
-  value: string
-  label: string
-  trigger?: string
-}
+import type { Model } from './types'
 
 export const Models: Model[] = [
   // 由于这个模型和 Cloudflare 的其他模型返回值不同, 故在 image/route.ts 中进行了两处特殊处理

diff --git a/app/lib/tasks.ts b/app/lib/tasks.ts
@@ -29,7 +29,7 @@ export async function handleTasks(tasks: Task[], setTasks: (action: SetAction<Ta
           star: false,
           hash,
           data,
-          prompt: task.prompt,
+          prompt: `${task.trigger ? `${task.trigger}, ` : ''}${task.prompt}`,
           model: task.model,
         }, ...prev])
         setTasks((prev) => prev.map((t) => t.createTimestamp === task.createTimestamp ? { ...t, status: 'success' } : t))
@@ -42,9 +42,9 @@ export async function handleTasks(tasks: Task[], setTasks: (action: SetAction<Ta
 
 /** Return whether the task is successfully generated and either base64 data or error message */
 async function generateImage(task: Task): Promise<{ success: boolean, data: string }> {
-  const { prompt, model, promptLanguage } = task
+  const { prompt, model, promptLanguage, trigger } = task
   try {
-    const promptEN = promptLanguage === 'zh' ? await translate(prompt) : prompt
+    const promptEN = `${trigger ? `${trigger}, ` : ''}${promptLanguage === 'en' ? prompt : await translate(prompt)}`
     let res: Response | undefined
     if (process.env.NEXT_PUBLIC_WORKERS_SERVER) {
       res = await fetch(`${process.env.NEXT_PUBLIC_WORKERS_SERVER}/painter/generate`, {

diff --git a/app/lib/types.ts b/app/lib/types.ts
@@ -16,6 +16,8 @@ export type Task = {
   createTimestamp: number
   /** The task's prompt */
   prompt: string
+  /** The trigger of lora model */
+  trigger?: string
   /** The task's model value */
   model: string
   /** The task's prompt language */
@@ -25,3 +27,9 @@ export type Task = {
   /** The task's error message */
   error?: string
 }
+
+export type Model = {
+  value: string
+  label: string
+  trigger?: string
+}
diff --git a/app/widgets/Tasks.tsx b/app/widgets/Tasks.tsx
@@ -39,11 +39,7 @@ export default function Tasks() {
               }
             >
               <p>
-                {task.error ? (
-                  <span>{task.error}</span>
-                ) : (
-                  <span><Tag>{task.promptLanguage} prompt</Tag> {task.prompt}</span>
-                )}
+                {task.error || `${task.trigger ? `${task.trigger}, ` : ''}${task.prompt}`}
               </p>
             </Card>
           ))}