处理完成后修改任务状态

This commit is contained in:
2023-02-18 20:26:08 +08:00
parent 3125c59d2a
commit 12c62c5af9
2 changed files with 13 additions and 4 deletions

View File

@@ -213,6 +213,12 @@ def put_watermark(img, wm_encoder=None):
import time import time
import requests import requests
# 对任务状态的修改
def update_task_status(task:dict, status:str, progress:int):
task["status"] = status
task["progress"] = progress
requests.put(f"http://localhost:3000/api/drawing/{task['id']}", data=task)
def main_dev(opt): def main_dev(opt):
model_name = '' # 默认模型 model_name = '' # 默认模型
model = None # 默认模型 model = None # 默认模型
@@ -303,7 +309,7 @@ def main_dev(opt):
start_code = None start_code = None
if opt.fixed_code: if opt.fixed_code:
start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device) start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
# x # 切换模型
if opt.torchscript or opt.ipex: if opt.torchscript or opt.ipex:
transformer = model.cond_stage_model.model transformer = model.cond_stage_model.model
unet = model.model.diffusion_model unet = model.model.diffusion_model
@@ -410,8 +416,9 @@ def main_dev(opt):
grid = put_watermark(grid, wm_encoder) grid = put_watermark(grid, wm_encoder)
grid.save(os.path.join(outpath, f'grid-{grid_count:04}.png')) grid.save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
grid_count += 1 grid_count += 1
print(f"Your samples are ready and waiting for you here: \n{outpath} \n" print(f"Your samples are ready and waiting for you here: \n{outpath} \n", f" \nEnjoy.")
f" \nEnjoy.") # 修改任务状态为完成
update_task_status(task=item, status='done', progress=1)
# 任务结束, 等待20s后退出 # 任务结束, 等待20s后退出
print("任务结束, 等待20s后退出..") print("任务结束, 等待20s后退出..")
time.sleep(20) time.sleep(20)

View File

@@ -38,7 +38,6 @@ export default defineEventHandler(async event => {
const body = await readBody(event) // 读取body数据到 const body = await readBody(event) // 读取body数据到
body.tid = Math.random().toString(36).substring(2, 34) // 生成随机任务ID body.tid = Math.random().toString(36).substring(2, 34) // 生成随机任务ID
body.uid = 1234567890 // 用户ID body.uid = 1234567890 // 用户ID
body.status = 'waiting' // 任务状态
body.createdAt = new Date().getTime() // 设定创建时间 body.createdAt = new Date().getTime() // 设定创建时间
// 验证需要输入的参数 // 验证需要输入的参数
body.model = 'SD2' // 模型 body.model = 'SD2' // 模型
@@ -52,6 +51,9 @@ export default defineEventHandler(async event => {
body.sampler = 'pndm' // 扩散采样器 body.sampler = 'pndm' // 扩散采样器
body.prompt_guidance = 0.5 // 提示词权重 body.prompt_guidance = 0.5 // 提示词权重
body.quality_details = 25 // 质量和细节(步数) body.quality_details = 25 // 质量和细节(步数)
// 输出参数
body.status = 'waiting' // 任务状态
body.progress = 0.00 // 任务进度
body.image = '' // 图片数据 body.image = '' // 图片数据
body.data = null // 生成结果 body.data = null // 生成结果
console.log('task:', body) console.log('task:', body)