This commit is contained in:
Timothy Jaeryang Baek
2026-08-30 23:46:02 -04:00
parent df495a7945
commit aeb126b95d
5 changed files with 78 additions and 33 deletions

View File

@@ -552,7 +552,12 @@ async def upload_image(request, image_data, content_type, metadata, user, db=Non
)
url = request.app.url_path_for('get_file_content_by_id', id=file_item.id)
return file_item, url
return file_item, {
'id': file_item.id,
'url': url,
'name': (file_item.meta or {}).get('name') or file_item.filename,
'content_type': (file_item.meta or {}).get('content_type'),
}
@router.post('/generations')
@@ -668,8 +673,8 @@ async def image_generations(
else:
image_data, content_type = await get_image_data(image['b64_json'])
_, url = await upload_image(request, image_data, content_type, {**data, **metadata}, user)
images.append({'url': url})
_, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user)
images.append(image_file)
return images
elif image_config.IMAGE_GENERATION_ENGINE == 'gemini':
@@ -712,21 +717,21 @@ async def image_generations(
if model.endswith(':predict'):
for image in res['predictions']:
image_data, content_type = await get_image_data(image['bytesBase64Encoded'])
_, url = await upload_image(request, image_data, content_type, {**data, **metadata}, user)
images.append({'url': url})
_, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user)
images.append(image_file)
elif model.endswith(':generateContent'):
for image in res['candidates']:
for part in image['content']['parts']:
if part.get('inlineData', {}).get('data'):
image_data, content_type = await get_image_data(part['inlineData']['data'])
_, url = await upload_image(
_, image_file = await upload_image(
request,
image_data,
content_type,
{**data, **metadata},
user,
)
images.append({'url': url})
images.append(image_file)
return images
@@ -776,14 +781,14 @@ async def image_generations(
headers,
trusted_base_url=image_config.COMFYUI_BASE_URL,
)
_, url = await upload_image(
_, image_file = await upload_image(
request,
image_data,
content_type,
{**form_data.model_dump(exclude_none=True), **metadata},
user,
)
images.append({'url': url})
images.append(image_file)
return images
elif image_config.IMAGE_GENERATION_ENGINE == 'automatic1111' or image_config.IMAGE_GENERATION_ENGINE == '':
# Automatic1111 holds one checkpoint instance-wide, so set_image_model
@@ -823,14 +828,14 @@ async def image_generations(
for image in res['images']:
image_data, content_type = await get_image_data(image)
_, url = await upload_image(
_, image_file = await upload_image(
request,
image_data,
content_type,
{**data, 'info': res['info'], **metadata},
user,
)
images.append({'url': url})
images.append(image_file)
return images
except Exception as e:
error = e
@@ -1039,8 +1044,8 @@ async def image_edits(
else:
image_data, content_type = await get_image_data(image['b64_json'])
_, url = await upload_image(request, image_data, content_type, {**data, **metadata}, user)
images.append({'url': url})
_, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user)
images.append(image_file)
return images
elif image_config.IMAGE_EDIT_ENGINE == 'gemini':
@@ -1089,14 +1094,14 @@ async def image_edits(
for part in image['content']['parts']:
if part.get('inlineData', {}).get('data'):
image_data, content_type = await get_image_data(part['inlineData']['data'])
_, url = await upload_image(
_, image_file = await upload_image(
request,
image_data,
content_type,
{**data, **metadata},
user,
)
images.append({'url': url})
images.append(image_file)
return images
@@ -1173,14 +1178,14 @@ async def image_edits(
headers,
trusted_base_url=image_config.IMAGES_EDIT_COMFYUI_BASE_URL,
)
_, url = await upload_image(
_, image_file = await upload_image(
request,
image_data,
content_type,
{**form_data.model_dump(exclude_none=True), **metadata},
user,
)
images.append({'url': url})
images.append(image_file)
return images
except Exception as e:

View File

@@ -938,20 +938,26 @@ async def _make_channel_emitter(request_info):
state = {'last_emit_at': 0.0, 'output': []}
THROTTLE_INTERVAL = 0.15 # ~6 updates/sec
async def _emit_channel_update(content: str, done: bool = False, output: list | None = None):
async def _emit_channel_update(
content: str,
done: bool = False,
output: list | None = None,
data: dict | None = None,
):
from open_webui.models.messages import MessageForm, Messages
msg = await Messages.get_message_by_id(message_id)
if not msg or msg.channel_id != channel_id:
return
update_form = MessageForm(content=content, data={'output': output} if output else None)
update_data = data or ({'output': output} if output else None)
update_form = MessageForm(content=content, data=update_data)
if done:
# Merge done flag into existing meta (preserve model_id etc.)
existing_meta = msg.meta or {}
update_form = MessageForm(
content=content,
data={'output': output} if output else None,
data=update_data,
meta={**existing_meta, 'done': True},
)
@@ -1000,6 +1006,29 @@ async def _make_channel_emitter(request_info):
state['last_emit_at'] = now
await _emit_channel_update(content, False, state['output'])
elif event_type in ('files', 'chat:message:files'):
from open_webui.models.messages import Messages
files = event_data.get('data', {}).get('files', [])
if not files:
return
msg = await Messages.get_message_by_id(message_id)
if not msg or msg.channel_id != channel_id:
return
existing_files = (msg.data or {}).get('files')
for file in files:
if isinstance(file, dict) and file.get('id'):
file['url'] = file['id']
await Channels.add_file_to_channel_by_id(channel_id, file['id'], msg.user_id)
await Channels.set_file_message_id_in_channel_by_id(channel_id, file['id'], message_id)
if isinstance(existing_files, list):
files.extend(existing_files)
await _emit_channel_update(msg.content, data={'files': files})
elif event_type == 'chat:message:error':
error = event_data.get('data', {}).get('error', {})
error_content = error.get('content', 'An error occurred') if isinstance(error, dict) else str(error)

View File

@@ -389,11 +389,16 @@ async def generate_image(
images = await image_generations(
request=__request__,
form_data=CreateImageForm(prompt=prompt),
metadata=(
{'channel_id': __chat_id__.removeprefix('channel:'), 'message_id': __message_id__}
if isinstance(__chat_id__, str) and __chat_id__.startswith('channel:')
else None
),
user=user,
)
# Prepare file entries for the images
image_files = [{'type': 'image', 'url': img['url']} for img in images]
image_files = [{'type': 'image', **img} for img in images]
# Persist files to DB if chat context is available
if is_saved_chat_id(__chat_id__) and __message_id__ and images:
@@ -457,11 +462,16 @@ async def edit_image(
images = await image_edits(
request=__request__,
form_data=EditImageForm(prompt=prompt, image=image_urls),
metadata=(
{'channel_id': __chat_id__.removeprefix('channel:'), 'message_id': __message_id__}
if isinstance(__chat_id__, str) and __chat_id__.startswith('channel:')
else None
),
user=user,
)
# Prepare file entries for the images
image_files = [{'type': 'image', 'url': img['url']} for img in images]
image_files = [{'type': 'image', **img} for img in images]
# Persist files to DB if chat context is available
if is_saved_chat_id(__chat_id__) and __message_id__ and images:

View File

@@ -86,13 +86,14 @@ async def get_image_url_from_base64(request, base64_image_string, metadata, user
# Extract base64 image data from the line
image_data, content_type = await get_image_data(base64_image_string)
if image_data is not None:
_, image_url = await upload_image(
_, image_file = await upload_image(
request,
image_data,
content_type,
metadata,
user,
)
image_url = image_file['url']
return image_url
return None

View File

@@ -1771,6 +1771,12 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra
if not chat_id or not isinstance(chat_id, str) or not __event_emitter__:
return form_data
is_channel_chat = chat_id.startswith('channel:')
image_metadata = {
'message_id': metadata.get('message_id', None),
**({'channel_id': chat_id.removeprefix('channel:')} if is_channel_chat else {'chat_id': chat_id}),
}
if not is_saved_chat_id(chat_id):
message_list = form_data.get('messages', [])
else:
@@ -1815,10 +1821,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra
images = await image_edits(
request=request,
form_data=EditImageForm(**{'prompt': prompt, 'image': input_images}),
metadata={
'chat_id': metadata.get('chat_id', None),
'message_id': metadata.get('message_id', None),
},
metadata=image_metadata,
user=user,
)
@@ -1836,7 +1839,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra
'files': [
{
'type': 'image',
'url': image['url'],
**image,
}
for image in images
]
@@ -1926,10 +1929,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra
images = await image_generations(
request=request,
form_data=CreateImageForm(**{'prompt': prompt}),
metadata={
'chat_id': metadata.get('chat_id', None),
'message_id': metadata.get('message_id', None),
},
metadata=image_metadata,
user=user,
)
@@ -1947,7 +1947,7 @@ async def chat_image_generation_handler(request: Request, form_data: dict, extra
'files': [
{
'type': 'image',
'url': image['url'],
**image,
}
for image in images
]