REFINE 重新调整302ai图像生成逻辑,优化任务查询和结果处理

This commit is contained in:
iHeyTang 2025-08-12 15:33:45 +08:00
parent 7630db1be3
commit 1ed33f9bbe
1 changed files with 23 additions and 40 deletions

View File

@ -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