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()}')}
else:
files = None
if provider == "302ai":
job_resp = send_request("post", f"https://{endpoint}/api/union/img/sync/generate/image",
headers={'Authorization': 'Bearer bowong7777'},
data={"prompt": prompt},
interval = 3
job_resp = send_request("post", f"https://{endpoint}/api/custom/image/submit/task",
data={"model_name":"302ai/mj", "prompt": prompt, "mode": "turbo"},
files=files,
timeout=timeout)
job_resp.raise_for_status()
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"]:
raise Exception("生成失败, 可能因为风控")
job_id = job_resp["data"]
for _ in range(0, timeout // interval, interval):
logger.info("等待" + str(interval) + "")
sleep(interval)
logger.info("查询结果")
resp = send_request("get", f"https://{endpoint}/api/mj/async/query/status?task_id={job_id}",
headers={'Authorization': 'Bearer bowong7777'}, timeout=30)
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"]:
if "fail" in resp.json()["msg"]:
raise Exception("生成失败,可能因为风控")
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 isinstance(result_url, list):
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,)
logger.success("img_url: " + result_url)
return (url_to_tensor(result_url),)
sleep(interval)
raise Exception("等待超时")
except Exception as e:
raise e