DemoGPT / modules /models /midjourney.py
JohnSmith9982's picture
Upload 98 files
0cc999a
raw
history blame
16.4 kB
import base64
import io
import json
import logging
import os
import pathlib
import tempfile
import time
from datetime import datetime
import requests
import tiktoken
from PIL import Image
from modules.config import retrieve_proxy
from modules.models.XMChat import XMChat
mj_proxy_api_base = os.getenv("MIDJOURNEY_PROXY_API_BASE")
mj_discord_proxy_url = os.getenv("MIDJOURNEY_DISCORD_PROXY_URL")
mj_temp_folder = os.getenv("MIDJOURNEY_TEMP_FOLDER")
class Midjourney_Client(XMChat):
class FetchDataPack:
"""
A class to store data for current fetching data from Midjourney API
"""
action: str # current action, e.g. "IMAGINE", "UPSCALE", "VARIATION"
prefix_content: str # prefix content, task description and process hint
task_id: str # task id
start_time: float # task start timestamp
timeout: int # task timeout in seconds
finished: bool # whether the task is finished
prompt: str # prompt for the task
def __init__(self, action, prefix_content, task_id, timeout=900):
self.action = action
self.prefix_content = prefix_content
self.task_id = task_id
self.start_time = time.time()
self.timeout = timeout
self.finished = False
def __init__(self, model_name, api_key, user_name=""):
super().__init__(api_key, user_name)
self.model_name = model_name
self.history = []
self.api_key = api_key
self.headers = {
"Content-Type": "application/json",
"mj-api-secret": f"{api_key}"
}
self.proxy_url = mj_proxy_api_base
self.command_splitter = "::"
if mj_temp_folder:
temp = "./tmp"
if user_name:
temp = os.path.join(temp, user_name)
if not os.path.exists(temp):
os.makedirs(temp)
self.temp_path = tempfile.mkdtemp(dir=temp)
logging.info("mj temp folder: " + self.temp_path)
else:
self.temp_path = None
def use_mj_self_proxy_url(self, img_url):
"""
replace discord cdn url with mj self proxy url
"""
return img_url.replace(
"https://cdn.discordapp.com/",
mj_discord_proxy_url and mj_discord_proxy_url or "https://cdn.discordapp.com/"
)
def split_image(self, image_url):
"""
when enabling temp dir, split image into 4 parts
"""
with retrieve_proxy():
image_bytes = requests.get(image_url).content
img = Image.open(io.BytesIO(image_bytes))
width, height = img.size
# calculate half width and height
half_width = width // 2
half_height = height // 2
# create coordinates (top-left x, top-left y, bottom-right x, bottom-right y)
coordinates = [(0, 0, half_width, half_height),
(half_width, 0, width, half_height),
(0, half_height, half_width, height),
(half_width, half_height, width, height)]
images = [img.crop(c) for c in coordinates]
return images
def auth_mj(self):
"""
auth midjourney api
"""
# TODO: check if secret is valid
return {'status': 'ok'}
def request_mj(self, path: str, action: str, data: str, retries=3):
"""
request midjourney api
"""
mj_proxy_url = self.proxy_url
if mj_proxy_url is None or not (mj_proxy_url.startswith("http://") or mj_proxy_url.startswith("https://")):
raise Exception('please set MIDJOURNEY_PROXY_API_BASE in ENV or in config.json')
auth_ = self.auth_mj()
if auth_.get('error'):
raise Exception('auth not set')
fetch_url = f"{mj_proxy_url}/{path}"
# logging.info(f"[MJ Proxy] {action} {fetch_url} params: {data}")
for _ in range(retries):
try:
with retrieve_proxy():
res = requests.request(method=action, url=fetch_url, headers=self.headers, data=data)
break
except Exception as e:
print(e)
if res.status_code != 200:
raise Exception(f'{res.status_code} - {res.content}')
return res
def fetch_status(self, fetch_data: FetchDataPack):
"""
fetch status of current task
"""
if fetch_data.start_time + fetch_data.timeout < time.time():
fetch_data.finished = True
return "任务超时,请检查 dc 输出。描述:" + fetch_data.prompt
time.sleep(3)
status_res = self.request_mj(f"task/{fetch_data.task_id}/fetch", "GET", '')
status_res_json = status_res.json()
if not (200 <= status_res.status_code < 300):
raise Exception("任务状态获取失败:" + status_res_json.get(
'error') or status_res_json.get('description') or '未知错误')
else:
fetch_data.finished = False
if status_res_json['status'] == "SUCCESS":
content = status_res_json['imageUrl']
fetch_data.finished = True
elif status_res_json['status'] == "FAILED":
content = status_res_json['failReason'] or '未知原因'
fetch_data.finished = True
elif status_res_json['status'] == "NOT_START":
content = f'任务未开始,已等待 {time.time() - fetch_data.start_time:.2f} 秒'
elif status_res_json['status'] == "IN_PROGRESS":
content = '任务正在运行'
if status_res_json.get('progress'):
content += f",进度:{status_res_json['progress']}"
elif status_res_json['status'] == "SUBMITTED":
content = '任务已提交处理'
elif status_res_json['status'] == "FAILURE":
fetch_data.finished = True
return "任务处理失败,原因:" + status_res_json['failReason'] or '未知原因'
else:
content = status_res_json['status']
if fetch_data.finished:
img_url = self.use_mj_self_proxy_url(status_res_json['imageUrl'])
if fetch_data.action == "DESCRIBE":
return f"\n{status_res_json['prompt']}"
time_cost_str = f"\n\n{fetch_data.action} 花费时间:{time.time() - fetch_data.start_time:.2f} 秒"
upscale_str = ""
variation_str = ""
if fetch_data.action in ["IMAGINE", "UPSCALE", "VARIATION"]:
upscale = [f'/mj UPSCALE{self.command_splitter}{i+1}{self.command_splitter}{fetch_data.task_id}'
for i in range(4)]
upscale_str = '\n放大图片:\n\n' + '\n\n'.join(upscale)
variation = [f'/mj VARIATION{self.command_splitter}{i+1}{self.command_splitter}{fetch_data.task_id}'
for i in range(4)]
variation_str = '\n图片变体:\n\n' + '\n\n'.join(variation)
if self.temp_path and fetch_data.action in ["IMAGINE", "VARIATION"]:
try:
images = self.split_image(img_url)
# save images to temp path
for i in range(4):
images[i].save(pathlib.Path(self.temp_path) / f"{fetch_data.task_id}_{i}.png")
img_str = '\n'.join(
[f"![{fetch_data.task_id}](/file={self.temp_path}/{fetch_data.task_id}_{i}.png)"
for i in range(4)])
return fetch_data.prefix_content + f"{time_cost_str}\n\n{img_str}{upscale_str}{variation_str}"
except Exception as e:
logging.error(e)
return fetch_data.prefix_content + \
f"{time_cost_str}[![{fetch_data.task_id}]({img_url})]({img_url}){upscale_str}{variation_str}"
else:
content = f"**任务状态:** [{(datetime.now()).strftime('%Y-%m-%d %H:%M:%S')}] - {content}"
content += f"\n\n花费时间:{time.time() - fetch_data.start_time:.2f} 秒"
if status_res_json['status'] == 'IN_PROGRESS' and status_res_json.get('imageUrl'):
img_url = status_res_json.get('imageUrl')
return f"{content}\n[![{fetch_data.task_id}]({img_url})]({img_url})"
return content
return None
def handle_file_upload(self, files, chatbot, language):
"""
handle file upload
"""
if files:
for file in files:
if file.name:
logging.info(f"尝试读取图像: {file.name}")
self.try_read_image(file.name)
if self.image_path is not None:
chatbot = chatbot + [((self.image_path,), None)]
if self.image_bytes is not None:
logging.info("使用图片作为输入")
return None, chatbot, None
def reset(self, remain_system_prompt=False):
self.image_bytes = None
self.image_path = None
return super().reset()
def get_answer_at_once(self):
content = self.history[-1]['content']
answer = self.get_help()
if not content.lower().startswith("/mj"):
return answer, len(content)
prompt = content[3:].strip()
action = "IMAGINE"
first_split_index = prompt.find(self.command_splitter)
if first_split_index > 0:
action = prompt[:first_split_index]
if action not in ["IMAGINE", "DESCRIBE", "UPSCALE",
# "VARIATION", "BLEND", "REROLL"
]:
raise Exception("任务提交失败:未知的任务类型")
else:
action_index = None
action_use_task_id = None
if action in ["VARIATION", "UPSCALE", "REROLL"]:
action_index = int(prompt[first_split_index + 2:first_split_index + 3])
action_use_task_id = prompt[first_split_index + 5:]
try:
res = None
if action == "IMAGINE":
data = {
"prompt": prompt
}
if self.image_bytes is not None:
data["base64"] = 'data:image/png;base64,' + self.image_bytes
res = self.request_mj("submit/imagine", "POST",
json.dumps(data))
elif action == "DESCRIBE":
res = self.request_mj("submit/describe", "POST",
json.dumps({"base64": 'data:image/png;base64,' + self.image_bytes}))
elif action == "BLEND":
res = self.request_mj("submit/blend", "POST", json.dumps(
{"base64Array": [self.image_bytes, self.image_bytes]}))
elif action in ["UPSCALE", "VARIATION", "REROLL"]:
res = self.request_mj(
"submit/change", "POST",
json.dumps({"action": action, "index": action_index, "taskId": action_use_task_id}))
res_json = res.json()
if not (200 <= res.status_code < 300) or (res_json['code'] not in [1, 22]):
answer = "任务提交失败:" + res_json.get('error', res_json.get('description', '未知错误'))
else:
task_id = res_json['result']
prefix_content = f"**画面描述:** {prompt}\n**任务ID:** {task_id}\n"
fetch_data = Midjourney_Client.FetchDataPack(
action=action,
prefix_content=prefix_content,
task_id=task_id,
)
fetch_data.prompt = prompt
while not fetch_data.finished:
answer = self.fetch_status(fetch_data)
except Exception as e:
logging.error("submit failed", e)
answer = "任务提交错误:" + str(e.args[0]) if e.args else '未知错误'
return answer, tiktoken.get_encoding("cl100k_base").encode(content)
def get_answer_stream_iter(self):
content = self.history[-1]['content']
answer = self.get_help()
if not content.lower().startswith("/mj"):
yield answer
return
prompt = content[3:].strip()
action = "IMAGINE"
first_split_index = prompt.find(self.command_splitter)
if first_split_index > 0:
action = prompt[:first_split_index]
if action not in ["IMAGINE", "DESCRIBE", "UPSCALE",
"VARIATION", "BLEND", "REROLL"
]:
yield "任务提交失败:未知的任务类型"
return
action_index = None
action_use_task_id = None
if action in ["VARIATION", "UPSCALE", "REROLL"]:
action_index = int(prompt[first_split_index + 2:first_split_index + 3])
action_use_task_id = prompt[first_split_index + 5:]
try:
res = None
if action == "IMAGINE":
data = {
"prompt": prompt
}
if self.image_bytes is not None:
data["base64"] = 'data:image/png;base64,' + self.image_bytes
res = self.request_mj("submit/imagine", "POST",
json.dumps(data))
elif action == "DESCRIBE":
res = self.request_mj("submit/describe", "POST", json.dumps(
{"base64": 'data:image/png;base64,' + self.image_bytes}))
elif action == "BLEND":
res = self.request_mj("submit/blend", "POST", json.dumps(
{"base64Array": [self.image_bytes, self.image_bytes]}))
elif action in ["UPSCALE", "VARIATION", "REROLL"]:
res = self.request_mj(
"submit/change", "POST",
json.dumps({"action": action, "index": action_index, "taskId": action_use_task_id}))
res_json = res.json()
if not (200 <= res.status_code < 300) or (res_json['code'] not in [1, 22]):
yield "任务提交失败:" + res_json.get('error', res_json.get('description', '未知错误'))
else:
task_id = res_json['result']
prefix_content = f"**画面描述:** {prompt}\n**任务ID:** {task_id}\n"
content = f"[{(datetime.now()).strftime('%Y-%m-%d %H:%M:%S')}] - 任务提交成功:" + \
res_json.get('description') or '请稍等片刻'
yield content
fetch_data = Midjourney_Client.FetchDataPack(
action=action,
prefix_content=prefix_content,
task_id=task_id,
)
while not fetch_data.finished:
yield self.fetch_status(fetch_data)
except Exception as e:
logging.error('submit failed', e)
yield "任务提交错误:" + str(e.args[0]) if e.args else '未知错误'
def get_help(self):
return """```
【绘图帮助】
所有命令都需要以 /mj 开头,如:/mj a dog
IMAGINE - 绘图,可以省略该命令,后面跟上绘图内容
/mj a dog
/mj IMAGINE::a cat
DESCRIBE - 描述图片,需要在右下角上传需要描述的图片内容
/mj DESCRIBE::
UPSCALE - 确认后放大图片,第一个数值为需要放大的图片(1~4),第二参数为任务ID
/mj UPSCALE::1::123456789
请使用SD进行UPSCALE
VARIATION - 图片变体,第一个数值为需要放大的图片(1~4),第二参数为任务ID
/mj VARIATION::1::123456789
【绘图参数】
所有命令默认会带上参数--v 5.2
其他参数参照 https://docs.midjourney.com/docs/parameter-list
长宽比 --aspect/--ar
--ar 1:2
--ar 16:9
负面tag --no
--no plants
--no hands
随机种子 --seed
--seed 1
生成动漫风格(NijiJourney) --niji
--niji
```
"""