REFINE 重新调整302ai图像生成逻辑,优化任务查询和结果处理
This commit is contained in:
parent
7630db1be3
commit
1ed33f9bbe
|
|
@ -235,56 +235,39 @@ class ModalMidJourneyGenerateImage:
|
||||||
f'image/{format.lower()}')}
|
f'image/{format.lower()}')}
|
||||||
else:
|
else:
|
||||||
files = None
|
files = None
|
||||||
|
|
||||||
if provider == "302ai":
|
if provider == "302ai":
|
||||||
job_resp = send_request("post", f"https://{endpoint}/api/union/img/sync/generate/image",
|
interval = 3
|
||||||
headers={'Authorization': 'Bearer bowong7777'},
|
job_resp = send_request("post", f"https://{endpoint}/api/custom/image/submit/task",
|
||||||
data={"prompt": prompt},
|
data={"model_name":"302ai/mj", "prompt": prompt, "mode": "turbo"},
|
||||||
files=files,
|
files=files,
|
||||||
timeout=timeout)
|
timeout=timeout)
|
||||||
job_resp.raise_for_status()
|
job_resp.raise_for_status()
|
||||||
job_resp = job_resp.json()
|
job_resp = job_resp.json()
|
||||||
if not job_resp["status"]:
|
|
||||||
raise Exception("生成失败, 可能因为风控")
|
|
||||||
result_url = job_resp["data"]
|
|
||||||
if isinstance(result_url, list):
|
|
||||||
result_list = []
|
|
||||||
for url in result_url:
|
|
||||||
logger.success("img_url: " + url)
|
|
||||||
result_list.append(url_to_tensor(url).squeeze(0))
|
|
||||||
result_list = torch.stack(result_list, dim=0)
|
|
||||||
return (result_list,)
|
|
||||||
logger.success("img_url: " + result_url)
|
|
||||||
return (url_to_tensor(result_url),)
|
|
||||||
elif provider == "ttapi":
|
|
||||||
interval = 3
|
|
||||||
job_resp = send_request("post", f"https://{endpoint}/api/mj/async/generate/image?prompt={prompt}",
|
|
||||||
headers={'Authorization': 'Bearer bowong7777'},
|
|
||||||
files=files,
|
|
||||||
timeout=150)
|
|
||||||
job_resp.raise_for_status()
|
|
||||||
job_resp = job_resp.json()
|
|
||||||
if not job_resp["status"]:
|
if not job_resp["status"]:
|
||||||
raise Exception("生成失败, 可能因为风控")
|
raise Exception("生成失败, 可能因为风控")
|
||||||
job_id = job_resp["data"]
|
job_id = job_resp["data"]
|
||||||
for _ in range(0, timeout // interval, interval):
|
for _ in range(0, timeout // interval, interval):
|
||||||
logger.info("查询结果")
|
logger.info("等待" + str(interval) + "秒")
|
||||||
resp = send_request("get", f"https://{endpoint}/api/mj/async/query/status?task_id={job_id}",
|
|
||||||
headers={'Authorization': 'Bearer bowong7777'}, timeout=30)
|
|
||||||
resp.raise_for_status()
|
|
||||||
if resp.json()["status"]:
|
|
||||||
if "fail" in resp.json()["msg"]:
|
|
||||||
raise Exception("生成失败,可能因为风控")
|
|
||||||
result_url = resp.json()["data"]
|
|
||||||
if isinstance(result_url, list):
|
|
||||||
result_list = []
|
|
||||||
for url in result_url:
|
|
||||||
logger.success("img_url: " + url)
|
|
||||||
result_list.append(url_to_tensor(url).squeeze(0))
|
|
||||||
result_list = torch.stack(result_list, dim=0)
|
|
||||||
return (result_list,)
|
|
||||||
logger.success("img_url: " + result_url)
|
|
||||||
return (url_to_tensor(result_url),)
|
|
||||||
sleep(interval)
|
sleep(interval)
|
||||||
|
logger.info("查询结果")
|
||||||
|
resp = send_request("get", f"https://{endpoint}/api/custom/task/status?task_id={job_id}", timeout=30)
|
||||||
|
resp.raise_for_status()
|
||||||
|
if resp.json()["status"] == "running":
|
||||||
|
logger.info("任务正在运行")
|
||||||
|
continue
|
||||||
|
if resp.json()["status"] == "failed":
|
||||||
|
raise Exception(f"生成失败: {resp.json()['msg']}")
|
||||||
|
if resp.json()["status"] == "success":
|
||||||
|
result_url = resp.json()["data"]
|
||||||
|
if not isinstance(result_url, list):
|
||||||
|
raise Exception("生成失败,返回结果为空")
|
||||||
|
result_list = []
|
||||||
|
for url in result_url:
|
||||||
|
logger.success("img_url: " + url)
|
||||||
|
result_list.append(url_to_tensor(url).squeeze(0))
|
||||||
|
result_list = torch.stack(result_list, dim=0)
|
||||||
|
return (result_list,)
|
||||||
raise Exception("等待超时")
|
raise Exception("等待超时")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise e
|
raise e
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue