mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-16 11:57:51 +01:00
feat: image edit support
This commit is contained in:
Timothy Jaeryang Baek
@@ -47,7 +47,8 @@ from open_webui.routers.retrieval import (
|
||||
from open_webui.routers.images import (
|
||||
image_generations,
|
||||
CreateImageForm,
|
||||
upload_image,
|
||||
image_edits,
|
||||
EditImageForm,
|
||||
)
|
||||
from open_webui.routers.pipelines import (
|
||||
process_pipeline_inlet_filter,
|
||||
@@ -717,9 +718,31 @@ async def chat_web_search_handler(
|
||||
return form_data
|
||||
|
||||
|
||||
def get_last_images(message_list):
|
||||
images = []
|
||||
for message in reversed(message_list):
|
||||
images_flag = False
|
||||
for file in message.get("files", []):
|
||||
if file.get("type") == "image":
|
||||
images.append(file.get("url"))
|
||||
images_flag = True
|
||||
|
||||
if images_flag:
|
||||
break
|
||||
|
||||
return images
|
||||
|
||||
|
||||
async def chat_image_generation_handler(
|
||||
request: Request, form_data: dict, extra_params: dict, user
|
||||
):
|
||||
metadata = extra_params.get("__metadata__", {})
|
||||
chat_id = metadata.get("chat_id", None)
|
||||
if not chat_id:
|
||||
return form_data
|
||||
|
||||
chat = Chats.get_chat_by_id_and_user_id(chat_id, user.id)
|
||||
|
||||
__event_emitter__ = extra_params["__event_emitter__"]
|
||||
await __event_emitter__(
|
||||
{
|
||||
@@ -728,87 +751,151 @@ async def chat_image_generation_handler(
|
||||
}
|
||||
)
|
||||
|
||||
messages = form_data["messages"]
|
||||
user_message = get_last_user_message(messages)
|
||||
messages_map = chat.chat.get("history", {}).get("messages", {})
|
||||
message_id = chat.chat.get("history", {}).get("currentId")
|
||||
message_list = get_message_list(messages_map, message_id)
|
||||
user_message = get_last_user_message(message_list)
|
||||
|
||||
prompt = user_message
|
||||
negative_prompt = ""
|
||||
|
||||
if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION:
|
||||
try:
|
||||
res = await generate_image_prompt(
|
||||
request,
|
||||
{
|
||||
"model": form_data["model"],
|
||||
"messages": messages,
|
||||
},
|
||||
user,
|
||||
)
|
||||
|
||||
response = res["choices"][0]["message"]["content"]
|
||||
|
||||
try:
|
||||
bracket_start = response.find("{")
|
||||
bracket_end = response.rfind("}") + 1
|
||||
|
||||
if bracket_start == -1 or bracket_end == -1:
|
||||
raise Exception("No JSON object found in the response")
|
||||
|
||||
response = response[bracket_start:bracket_end]
|
||||
response = json.loads(response)
|
||||
prompt = response.get("prompt", [])
|
||||
except Exception as e:
|
||||
prompt = user_message
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
prompt = user_message
|
||||
input_images = get_last_images(message_list)
|
||||
|
||||
system_message_content = ""
|
||||
if len(input_images) == 0:
|
||||
# Create image(s)
|
||||
if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION:
|
||||
try:
|
||||
res = await generate_image_prompt(
|
||||
request,
|
||||
{
|
||||
"model": form_data["model"],
|
||||
"messages": form_data["messages"],
|
||||
},
|
||||
user,
|
||||
)
|
||||
|
||||
try:
|
||||
images = await image_generations(
|
||||
request=request,
|
||||
form_data=CreateImageForm(**{"prompt": prompt}),
|
||||
user=user,
|
||||
)
|
||||
response = res["choices"][0]["message"]["content"]
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {"description": "Image created", "done": True},
|
||||
}
|
||||
)
|
||||
try:
|
||||
bracket_start = response.find("{")
|
||||
bracket_end = response.rfind("}") + 1
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "files",
|
||||
"data": {
|
||||
"files": [
|
||||
{
|
||||
"type": "image",
|
||||
"url": image["url"],
|
||||
}
|
||||
for image in images
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
if bracket_start == -1 or bracket_end == -1:
|
||||
raise Exception("No JSON object found in the response")
|
||||
|
||||
system_message_content = "<context>User is shown the generated image, tell the user that the image has been generated</context>"
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {
|
||||
"description": f"An error occurred while generating an image",
|
||||
"done": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
response = response[bracket_start:bracket_end]
|
||||
response = json.loads(response)
|
||||
prompt = response.get("prompt", [])
|
||||
except Exception as e:
|
||||
prompt = user_message
|
||||
|
||||
system_message_content = "<context>Unable to generate an image, tell the user that an error occurred</context>"
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
prompt = user_message
|
||||
|
||||
try:
|
||||
images = await image_generations(
|
||||
request=request,
|
||||
form_data=CreateImageForm(**{"prompt": prompt}),
|
||||
user=user,
|
||||
)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {"description": "Image created", "done": True},
|
||||
}
|
||||
)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "files",
|
||||
"data": {
|
||||
"files": [
|
||||
{
|
||||
"type": "image",
|
||||
"url": image["url"],
|
||||
}
|
||||
for image in images
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
system_message_content = "<context>The requested image has been created and is now being shown to the user. Let them know that it has been generated.</context>"
|
||||
except Exception as e:
|
||||
log.debug(e)
|
||||
|
||||
error_message = ""
|
||||
if isinstance(e, HTTPException):
|
||||
if e.detail and isinstance(e.detail, dict):
|
||||
error_message = e.detail.get("message", str(e.detail))
|
||||
else:
|
||||
error_message = str(e.detail)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {
|
||||
"description": f"An error occurred while generating an image",
|
||||
"done": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
system_message_content = f"<context>Image generation was attempted but failed. The system is currently unable to generate the image. Tell the user that an error occurred: {error_message}</context>"
|
||||
else:
|
||||
# Edit image(s)
|
||||
try:
|
||||
images = await image_edits(
|
||||
request=request,
|
||||
form_data=EditImageForm(**{"prompt": prompt, "image": input_images}),
|
||||
user=user,
|
||||
)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {"description": "Image created", "done": True},
|
||||
}
|
||||
)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "files",
|
||||
"data": {
|
||||
"files": [
|
||||
{
|
||||
"type": "image",
|
||||
"url": image["url"],
|
||||
}
|
||||
for image in images
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
system_message_content = "<context>The requested image has been created and is now being shown to the user. Let them know that it has been generated.</context>"
|
||||
except Exception as e:
|
||||
log.debug(e)
|
||||
|
||||
error_message = ""
|
||||
if isinstance(e, HTTPException):
|
||||
if e.detail and isinstance(e.detail, dict):
|
||||
error_message = e.detail.get("message", str(e.detail))
|
||||
else:
|
||||
error_message = str(e.detail)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {
|
||||
"description": f"An error occurred while generating an image",
|
||||
"done": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
system_message_content = f"<context>Image generation was attempted but failed. The system is currently unable to generate the image. Tell the user that an error occurred: {error_message}</context>"
|
||||
|
||||
if system_message_content:
|
||||
form_data["messages"] = add_or_update_system_message(
|
||||
|
||||
Reference in New Issue
Block a user