Coverage for open_webui/routers/images.py: 37%

1from __future__ import annotations 

2 

3import asyncio 

4import base64 

5import io 

6import logging 

7import mimetypes 

8import re 

9import uuid 

10from pathlib import Path 

11from types import SimpleNamespace 

12from typing import Optional 

13from urllib.parse import quote, urlparse 

14 

15import aiofiles 

16import aiohttp 

17from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile 

18from fastapi.responses import FileResponse 

19from open_webui.config import ( 

20 CACHE_DIR, 

21 ENABLE_OPENAI_IMAGE_EDIT_NORMALIZATION, 

22 IMAGE_AUTO_SIZE_MODELS_REGEX_PATTERN, 

23 IMAGE_URL_RESPONSE_MODELS_REGEX_PATTERN, 

24) 

25from open_webui.constants import ERROR_MESSAGES 

26from open_webui.env import AIOHTTP_CLIENT_ALLOW_REDIRECTS, AIOHTTP_CLIENT_SESSION_SSL, ENABLE_FORWARD_USER_INFO_HEADERS 

27from open_webui.events import EVENTS, publish_event 

28from open_webui.internal.db import get_async_session 

29from open_webui.models.chats import Chats 

30from open_webui.models.config import Config 

31from open_webui.retrieval.web.utils import get_ssrf_safe_session, validate_url 

32from open_webui.routers.files import get_file_content_by_id, upload_file_handler 

33from open_webui.utils.access_control import has_permission 

34from open_webui.utils.auth import get_admin_user, get_verified_user 

35from open_webui.utils.headers import include_user_info_headers 

36from open_webui.utils.images.comfyui import ( 

37 ComfyUICreateImageForm, 

38 ComfyUIEditImageForm, 

39 ComfyUIWorkflow, 

40 comfyui_create_image, 

41 comfyui_edit_image, 

42 comfyui_upload_image, 

43) 

44from open_webui.utils.json_codec import JSONCodec 

45from open_webui.utils.session_pool import get_session 

46from PIL import Image, ImageOps 

47from pydantic import BaseModel 

48from sqlalchemy.ext.asyncio import AsyncSession 

49 

50log = logging.getLogger(__name__) 

51 

52# An image can lie as easily as it can illuminate. Let what 

53# is generated here be honest about what it shows. 

54IMAGE_CACHE_DIR = CACHE_DIR / 'image' / 'generations' 

55IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True) 

56 

57router = APIRouter() 

58 

59IMAGE_FILE_EXTENSIONS = { 

60 'image/jpeg': '.jpg', 

61 'image/jpg': '.jpg', 

62 'image/mpo': '.jpg', 

63 'image/png': '.png', 

64 'image/webp': '.webp', 

65} 

66 

67IMAGE_CONFIG_KEYS = { 

68 'ENABLE_IMAGE_GENERATION': 'image_generation.enable', 

69 'ENABLE_IMAGE_PROMPT_GENERATION': 'image_generation.prompt.enable', 

70 'IMAGE_GENERATION_ENGINE': 'image_generation.engine', 

71 'IMAGE_GENERATION_MODEL': 'image_generation.model', 

72 'IMAGE_SIZE': 'image_generation.size', 

73 'IMAGE_STEPS': 'image_generation.steps', 

74 'IMAGES_OPENAI_API_BASE_URL': 'image_generation.openai.api_base_url', 

75 'IMAGES_OPENAI_API_KEY': 'image_generation.openai.api_key', 

76 'IMAGES_OPENAI_API_VERSION': 'image_generation.openai.api_version', 

77 'IMAGES_OPENAI_API_PARAMS': 'image_generation.openai.params', 

78 'AUTOMATIC1111_BASE_URL': 'image_generation.automatic1111.base_url', 

79 'AUTOMATIC1111_API_AUTH': 'image_generation.automatic1111.api_auth', 

80 'AUTOMATIC1111_PARAMS': 'image_generation.automatic1111.api_params', 

81 'COMFYUI_BASE_URL': 'image_generation.comfyui.base_url', 

82 'COMFYUI_API_KEY': 'image_generation.comfyui.api_key', 

83 'COMFYUI_WORKFLOW': 'image_generation.comfyui.workflow', 

84 'COMFYUI_WORKFLOW_NODES': 'image_generation.comfyui.nodes', 

85 'IMAGES_GEMINI_API_BASE_URL': 'image_generation.gemini.api_base_url', 

86 'IMAGES_GEMINI_API_KEY': 'image_generation.gemini.api_key', 

87 'IMAGES_GEMINI_ENDPOINT_METHOD': 'image_generation.gemini.endpoint_method', 

88 'ENABLE_IMAGE_EDIT': 'images.edit.enable', 

89 'IMAGE_EDIT_ENGINE': 'images.edit.engine', 

90 'IMAGE_EDIT_MODEL': 'images.edit.model', 

91 'IMAGE_EDIT_SIZE': 'images.edit.size', 

92 'IMAGES_EDIT_OPENAI_API_BASE_URL': 'images.edit.openai.api_base_url', 

93 'IMAGES_EDIT_OPENAI_API_KEY': 'images.edit.openai.api_key', 

94 'IMAGES_EDIT_OPENAI_API_VERSION': 'images.edit.openai.api_version', 

95 'IMAGES_EDIT_GEMINI_API_BASE_URL': 'images.edit.gemini.api_base_url', 

96 'IMAGES_EDIT_GEMINI_API_KEY': 'images.edit.gemini.api_key', 

97 'IMAGES_EDIT_COMFYUI_BASE_URL': 'images.edit.comfyui.base_url', 

98 'IMAGES_EDIT_COMFYUI_API_KEY': 'images.edit.comfyui.api_key', 

99 'IMAGES_EDIT_COMFYUI_WORKFLOW': 'images.edit.comfyui.workflow', 

100 'IMAGES_EDIT_COMFYUI_WORKFLOW_NODES': 'images.edit.comfyui.nodes', 

101 'USER_PERMISSIONS': 'user.permissions', 

102} 

103 

104 

105async def get_config_values(key_map: dict[str, str]) -> dict: 

106 values = await Config.get_many(*key_map.values()) 

107 return {field: values[storage_key] for field, storage_key in key_map.items() if storage_key in values} 

108 

109 

110async def get_image_config() -> SimpleNamespace: 

111 return SimpleNamespace(**await get_config_values(IMAGE_CONFIG_KEYS)) 

112 

113 

114def config_updates(data: dict, key_map: dict[str, str]) -> dict: 

115 return {key_map[field]: value for field, value in data.items() if field in key_map} 

116 

117 

118def normalize_openai_edit_image_data_url(data_url: str) -> str: 

119 if not data_url.startswith('data:') or ',' not in data_url: 

120 return data_url 

121 

122 header, encoded = data_url.split(',', 1) 

123 mime_type = header.split(';')[0].lstrip('data:').lower() 

124 if mime_type not in {'image/jpeg', 'image/jpg', 'image/mpo'}: 

125 return data_url 

126 

127 try: 

128 image_bytes = base64.b64decode(encoded) 

129 with Image.open(io.BytesIO(image_bytes)) as image: 

130 orientation = image.getexif().get(274) 

131 needs_normalization = ( 

132 mime_type == 'image/mpo' 

133 or image.format == 'MPO' 

134 or getattr(image, 'n_frames', 1) > 1 

135 or orientation not in (None, 1) 

136 or image.mode not in ('RGB', 'L') 

137 ) 

138 

139 if not needs_normalization: 

140 return data_url 

141 

142 image.seek(0) 

143 image = ImageOps.exif_transpose(image) 

144 if image.mode != 'RGB': 

145 image = image.convert('RGB') 

146 

147 output = io.BytesIO() 

148 image.save(output, format='JPEG', quality=95) 

149 normalized_image = base64.b64encode(output.getvalue()).decode('utf-8') 

150 return f'data:image/jpeg;base64,{normalized_image}' 

151 except Exception as e: 

152 log.debug('Image edit normalization skipped: %s', e) 

153 

154 return data_url 

155 

156 

157def get_image_file_item(base64_string, param_name='image'): 

158 header, encoded = base64_string.split(',', 1) 

159 mime_type = header.split(';')[0].lstrip('data:') or 'image/png' 

160 image_data = base64.b64decode(encoded) 

161 extension = IMAGE_FILE_EXTENSIONS.get(mime_type.lower()) or mimetypes.guess_extension(mime_type) or '.png' 

162 return ( 

163 param_name, 

164 ( 

165 f'{uuid.uuid4()}{extension}', 

166 io.BytesIO(image_data), 

167 mime_type, 

168 ), 

169 ) 

170 

171 

172async def set_image_model(request: Request, model: str): 

173 log.info('Setting image model to %s', model) 

174 await Config.upsert({'image_generation.model': model}) 

175 image_config = await get_image_config() 

176 if image_config.IMAGE_GENERATION_ENGINE in ['', 'automatic1111']: 176 ↛ 177line 176 didn't jump to line 177 because the condition on line 176 was never true

177 api_auth = get_automatic1111_api_auth(image_config) 

178 

179 try: 

180 session = await get_session() 

181 async with session.get( 

182 url=f'{image_config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options', 

183 headers={'authorization': api_auth}, 

184 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

185 ) as r: 

186 options = await r.json() 

187 if model != options['sd_model_checkpoint']: 

188 options['sd_model_checkpoint'] = model 

189 async with session.post( 

190 url=f'{image_config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options', 

191 json=options, 

192 headers={'authorization': api_auth}, 

193 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

194 ) as r: 

195 r.raise_for_status() 

196 except Exception as e: 

197 log.debug('%s', e) 

198 

199 return image_config.IMAGE_GENERATION_MODEL 

200 

201 

202async def get_image_model(request): 

203 image_config = await get_image_config() 

204 if image_config.IMAGE_GENERATION_ENGINE == 'openai': 

205 return image_config.IMAGE_GENERATION_MODEL if image_config.IMAGE_GENERATION_MODEL else 'dall-e-2' 

206 elif image_config.IMAGE_GENERATION_ENGINE == 'gemini': 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true

207 return image_config.IMAGE_GENERATION_MODEL if image_config.IMAGE_GENERATION_MODEL else 'imagen-3.0-generate-002' 

208 elif image_config.IMAGE_GENERATION_ENGINE == 'comfyui': 208 ↛ 209line 208 didn't jump to line 209 because the condition on line 208 was never true

209 return image_config.IMAGE_GENERATION_MODEL if image_config.IMAGE_GENERATION_MODEL else '' 

210 elif image_config.IMAGE_GENERATION_ENGINE == 'automatic1111' or image_config.IMAGE_GENERATION_ENGINE == '': 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true

211 try: 

212 session = await get_session() 

213 async with session.get( 

214 url=f'{image_config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options', 

215 headers={'authorization': get_automatic1111_api_auth(image_config)}, 

216 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

217 ) as r: 

218 options = await r.json() 

219 return options['sd_model_checkpoint'] 

220 except Exception as e: 

221 log.exception(f'Failed to get default model from automatic1111: {e}') 

222 raise HTTPException( 

223 status_code=400, 

224 detail=ERROR_MESSAGES.DEFAULT(e, 'Failed to connect to the image generation engine'), 

225 ) 

226 

227 

228class ImagesConfig(BaseModel): 

229 ENABLE_IMAGE_GENERATION: bool 

230 ENABLE_IMAGE_PROMPT_GENERATION: bool 

231 

232 IMAGE_GENERATION_ENGINE: str 

233 IMAGE_GENERATION_MODEL: str 

234 IMAGE_SIZE: str | None 

235 IMAGE_STEPS: int | None 

236 

237 IMAGES_OPENAI_API_BASE_URL: str 

238 IMAGES_OPENAI_API_KEY: str 

239 IMAGES_OPENAI_API_VERSION: str 

240 IMAGES_OPENAI_API_PARAMS: dict | str | None 

241 

242 AUTOMATIC1111_BASE_URL: str 

243 AUTOMATIC1111_API_AUTH: dict | str | None 

244 AUTOMATIC1111_PARAMS: dict | str | None 

245 

246 COMFYUI_BASE_URL: str 

247 COMFYUI_API_KEY: str 

248 COMFYUI_WORKFLOW: str 

249 COMFYUI_WORKFLOW_NODES: list[dict] 

250 

251 IMAGES_GEMINI_API_BASE_URL: str 

252 IMAGES_GEMINI_API_KEY: str 

253 IMAGES_GEMINI_ENDPOINT_METHOD: str 

254 

255 ENABLE_IMAGE_EDIT: bool 

256 IMAGE_EDIT_ENGINE: str 

257 IMAGE_EDIT_MODEL: str 

258 IMAGE_EDIT_SIZE: str | None 

259 

260 IMAGES_EDIT_OPENAI_API_BASE_URL: str 

261 IMAGES_EDIT_OPENAI_API_KEY: str 

262 IMAGES_EDIT_OPENAI_API_VERSION: str 

263 IMAGES_EDIT_GEMINI_API_BASE_URL: str 

264 IMAGES_EDIT_GEMINI_API_KEY: str 

265 IMAGES_EDIT_COMFYUI_BASE_URL: str 

266 IMAGES_EDIT_COMFYUI_API_KEY: str 

267 IMAGES_EDIT_COMFYUI_WORKFLOW: str 

268 IMAGES_EDIT_COMFYUI_WORKFLOW_NODES: list[dict] 

269 

270 

271@router.get('/config', response_model=ImagesConfig) 

272async def get_config(request: Request, user=Depends(get_admin_user)): 

273 return await get_config_values(IMAGE_CONFIG_KEYS) 

274 

275 

276@router.post('/config/update') 

277async def update_config(request: Request, form_data: ImagesConfig, user=Depends(get_admin_user)): 

278 if form_data.IMAGE_SIZE == 'auto' and not re.match( 278 ↛ 281line 278 didn't jump to line 281 because the condition on line 278 was never true

279 IMAGE_AUTO_SIZE_MODELS_REGEX_PATTERN, form_data.IMAGE_GENERATION_MODEL 

280 ): 

281 raise HTTPException( 

282 status_code=400, 

283 detail=ERROR_MESSAGES.INCORRECT_FORMAT( 

284 f' (auto is only allowed with models matching {IMAGE_AUTO_SIZE_MODELS_REGEX_PATTERN}).' 

285 ), 

286 ) 

287 

288 pattern = r'^\d+x\d+$' 

289 if not (form_data.IMAGE_SIZE == 'auto' or form_data.IMAGE_SIZE == '' or re.match(pattern, form_data.IMAGE_SIZE)): 

290 raise HTTPException( 

291 status_code=400, 

292 detail=ERROR_MESSAGES.INCORRECT_FORMAT(' (e.g., 512x512).'), 

293 ) 

294 

295 if form_data.IMAGE_STEPS < 0: 

296 raise HTTPException( 

297 status_code=400, 

298 detail=ERROR_MESSAGES.INCORRECT_FORMAT(' (e.g., 50).'), 

299 ) 

300 

301 updates = config_updates(form_data.model_dump(), IMAGE_CONFIG_KEYS) 

302 updates['image_generation.comfyui.base_url'] = form_data.COMFYUI_BASE_URL.strip('/') 

303 updates['images.edit.comfyui.base_url'] = form_data.IMAGES_EDIT_COMFYUI_BASE_URL.strip('/') 

304 await Config.upsert(updates) 

305 await set_image_model(request, form_data.IMAGE_GENERATION_MODEL) 

306 values = await get_config_values(IMAGE_CONFIG_KEYS) 

307 await publish_event( 

308 request, 

309 EVENTS.CONFIG_UPDATED, 

310 actor=user, 

311 subject_id='images', 

312 data={ 

313 'image_generation_enabled': values.get('ENABLE_IMAGE_GENERATION'), 

314 'image_edit_enabled': values.get('ENABLE_IMAGE_EDIT'), 

315 'image_generation_engine': values.get('IMAGE_GENERATION_ENGINE'), 

316 'image_edit_engine': values.get('IMAGE_EDIT_ENGINE'), 

317 }, 

318 ) 

319 return values 

320 

321 

322def get_automatic1111_api_auth(image_config): 

323 if image_config.AUTOMATIC1111_API_AUTH is None: 

324 return '' 

325 else: 

326 auth1111_byte_string = image_config.AUTOMATIC1111_API_AUTH.encode('utf-8') 

327 auth1111_base64_encoded_bytes = base64.b64encode(auth1111_byte_string) 

328 auth1111_base64_encoded_string = auth1111_base64_encoded_bytes.decode('utf-8') 

329 return f'Basic {auth1111_base64_encoded_string}' 

330 

331 

332class ConnectionVerificationForm(BaseModel): 

333 engine: str 

334 url: str 

335 key: str | None = None 

336 

337 

338@router.post('/verify') 

339async def verify_connection(form_data: ConnectionVerificationForm, user=Depends(get_admin_user)): 

340 url = form_data.url.rstrip('/') 

341 headers = {} 

342 if form_data.engine == 'automatic1111': 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true

343 url = f'{url}/sdapi/v1/options' 

344 if form_data.key is not None: 

345 headers['Authorization'] = f'Basic {base64.b64encode(form_data.key.encode("utf-8")).decode("utf-8")}' 

346 elif form_data.engine == 'comfyui': 346 ↛ 347line 346 didn't jump to line 347 because the condition on line 346 was never true

347 url = f'{url}/object_info' 

348 if form_data.key: 

349 headers['Authorization'] = f'Bearer {form_data.key}' 

350 else: 

351 raise HTTPException(status_code=400, detail='Unsupported image engine') 

352 

353 try: 

354 session = await get_session() 

355 async with session.get(url=url, headers=headers, ssl=AIOHTTP_CLIENT_SESSION_SSL) as r: 

356 r.raise_for_status() 

357 return True 

358 except Exception: 

359 raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL) 

360 

361 

362@router.get('/models') 

363async def get_models(request: Request, user=Depends(get_verified_user)): 

364 image_config = await get_image_config() 

365 try: 

366 if image_config.IMAGE_GENERATION_ENGINE == 'openai': 

367 return [ 

368 {'id': 'dall-e-2', 'name': 'DALL·E 2'}, 

369 {'id': 'dall-e-3', 'name': 'DALL·E 3'}, 

370 {'id': 'gpt-image-1', 'name': 'GPT-IMAGE 1'}, 

371 {'id': 'gpt-image-1.5', 'name': 'GPT-IMAGE 1.5'}, 

372 ] 

373 elif image_config.IMAGE_GENERATION_ENGINE == 'gemini': 373 ↛ 374line 373 didn't jump to line 374 because the condition on line 373 was never true

374 return [ 

375 {'id': 'imagen-3.0-generate-002', 'name': 'imagen-3.0 generate-002'}, 

376 ] 

377 elif image_config.IMAGE_GENERATION_ENGINE == 'comfyui': 377 ↛ 379line 377 didn't jump to line 379 because the condition on line 377 was never true

378 # TODO - get models from comfyui 

379 headers = {'Authorization': f'Bearer {image_config.COMFYUI_API_KEY}'} 

380 session = await get_session() 

381 async with session.get( 

382 url=f'{image_config.COMFYUI_BASE_URL}/object_info', 

383 headers=headers, 

384 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

385 ) as r: 

386 info = await r.json() 

387 

388 workflow = JSONCodec.loads(image_config.COMFYUI_WORKFLOW) 

389 model_node_id = None 

390 

391 for node in image_config.COMFYUI_WORKFLOW_NODES: 

392 if node['type'] == 'model': 

393 if node['node_ids']: 

394 model_node_id = node['node_ids'][0] 

395 break 

396 

397 if model_node_id: 

398 model_list_key = None 

399 

400 log.info(workflow[model_node_id]['class_type']) 

401 for key in info[workflow[model_node_id]['class_type']]['input']['required']: 

402 if '_name' in key: 

403 model_list_key = key 

404 break 

405 

406 if model_list_key: 

407 return list( 

408 map( 

409 lambda model: {'id': model, 'name': model}, 

410 info[workflow[model_node_id]['class_type']]['input']['required'][model_list_key][0], 

411 ) 

412 ) 

413 else: 

414 return list( 

415 map( 

416 lambda model: {'id': model, 'name': model}, 

417 info['CheckpointLoaderSimple']['input']['required']['ckpt_name'][0], 

418 ) 

419 ) 

420 elif image_config.IMAGE_GENERATION_ENGINE == 'automatic1111' or image_config.IMAGE_GENERATION_ENGINE == '': 420 ↛ 421line 420 didn't jump to line 421 because the condition on line 420 was never true

421 session = await get_session() 

422 async with session.get( 

423 url=f'{image_config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models', 

424 headers={'authorization': get_automatic1111_api_auth(image_config)}, 

425 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

426 ) as r: 

427 models = await r.json() 

428 return list( 

429 map( 

430 lambda model: {'id': model['title'], 'name': model['model_name']}, 

431 models, 

432 ) 

433 ) 

434 except Exception as e: 

435 log.error( 

436 f'Failed to list image generation models: {str(e) or type(e).__name__}', 

437 exc_info=not isinstance(e, (aiohttp.ClientConnectionError, TimeoutError)), 

438 ) 

439 raise HTTPException( 

440 status_code=400, 

441 detail=ERROR_MESSAGES.DEFAULT(e, 'Failed to retrieve image generation models'), 

442 ) 

443 

444 

445class CreateImageForm(BaseModel): 

446 model: str | None = None 

447 prompt: str 

448 size: str | None = None 

449 n: int = 1 

450 steps: int | None = None 

451 negative_prompt: str | None = None 

452 

453 

454GenerateImageForm = CreateImageForm # Alias for backward compatibility 

455 

456 

457def _is_same_origin(url: str, base_url: str) -> bool: 

458 """Compare scheme + hostname + port of two URLs. 

459 

460 Pure string-prefix matching (``startswith``) is vulnerable to 

461 userinfo injection (``http://host:port@evil.com/``) and suffix 

462 confusion (``http://host:portevil.com/``). Parsing both URLs 

463 and comparing the three origin components eliminates those 

464 attack vectors. 

465 """ 

466 

467 def _default_port(scheme: str) -> int: 

468 return 443 if scheme == 'https' else 80 

469 

470 parsed = urlparse(url) 

471 trusted = urlparse(base_url) 

472 return ( 

473 parsed.scheme == trusted.scheme 

474 and parsed.hostname == trusted.hostname 

475 and (parsed.port or _default_port(parsed.scheme)) == (trusted.port or _default_port(trusted.scheme)) 

476 ) 

477 

478 

479async def get_image_data(data: str, headers=None, trusted_base_url: str | None = None): 

480 try: 

481 if data.startswith('http://') or data.startswith('https://'): 

482 # Defense-in-depth: gate before fetch (mirrors load_url_image). 

483 # For URLs originating from an admin-configured backend (e.g. 

484 # ComfyUI on a private network), skip SSRF validation only when 

485 # the URL shares the exact same origin (scheme + host + port) 

486 # as the admin-configured base. This avoids both the global 

487 # ENABLE_LOCAL_WEB_FETCH hammer and a blanket trust flag 

488 # that would follow arbitrary redirects. 

489 if trusted_base_url and _is_same_origin(data, trusted_base_url): 

490 log.debug('Skipping URL validation for trusted backend: %s', data) 

491 else: 

492 await asyncio.to_thread(validate_url, data) 

493 session = await get_session() 

494 async with session.get( 

495 data, 

496 headers=headers, 

497 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

498 ) as r: 

499 r.raise_for_status() 

500 content_type = r.headers.get('content-type', '') 

501 if content_type.split('/')[0] == 'image': 

502 return await r.read(), content_type 

503 else: 

504 log.error('Url does not point to an image.') 

505 return None, None 

506 else: 

507 if ',' in data: 

508 header, encoded = data.split(',', 1) 

509 mime_type = header.split(';')[0].lstrip('data:') 

510 img_data = base64.b64decode(encoded) 

511 else: 

512 mime_type = 'image/png' 

513 img_data = base64.b64decode(data) 

514 return img_data, mime_type 

515 except Exception as e: 

516 log.exception(f'Error loading image data: {e}') 

517 return None, None 

518 

519 

520async def upload_image(request, image_data, content_type, metadata, user, db=None): 

521 if image_data is None or content_type is None: 521 ↛ anywhereline 521 didn't jump anywhere: it always raised an exception.

522 raise ValueError('Failed to retrieve image data from the generation backend') 

523 image_format = mimetypes.guess_extension(content_type) 

524 file = UploadFile( 

525 file=io.BytesIO(image_data), 

526 filename=f'generated-image{image_format}', # will be converted to a unique ID on upload_file 

527 headers={ 

528 'content-type': content_type, 

529 }, 

530 ) 

531 file_item = await upload_file_handler( 

532 request, 

533 file=file, 

534 metadata=metadata, 

535 process=False, 

536 user=user, 

537 ) 

538 

539 if file_item and file_item.id: 

540 # If chat_id and message_id are provided in metadata, link the file to the chat message 

541 chat_id = metadata.get('chat_id') 

542 message_id = metadata.get('message_id') 

543 

544 if chat_id and message_id: 

545 await Chats.insert_chat_files( 

546 chat_id=chat_id, 

547 message_id=message_id, 

548 file_ids=[file_item.id], 

549 user_id=user.id, 

550 db=db, 

551 ) 

552 

553 url = request.app.url_path_for('get_file_content_by_id', id=file_item.id) 

554 return file_item, { 

555 'id': file_item.id, 

556 'url': url, 

557 'name': (file_item.meta or {}).get('name') or file_item.filename, 

558 'content_type': (file_item.meta or {}).get('content_type'), 

559 } 

560 

561 

562@router.post('/generations') 

563async def generate_images(request: Request, form_data: CreateImageForm, user=Depends(get_verified_user)): 

564 image_config = await get_image_config() 

565 if not image_config.ENABLE_IMAGE_GENERATION: 

566 raise HTTPException( 

567 status_code=403, 

568 detail=ERROR_MESSAGES.ACCESS_PROHIBITED, 

569 ) 

570 

571 if user.role != 'admin' and not await has_permission( 571 ↛ 574line 571 didn't jump to line 574 because the condition on line 571 was never true

572 user.id, 'features.image_generation', image_config.USER_PERMISSIONS 

573 ): 

574 raise HTTPException( 

575 status_code=403, 

576 detail=ERROR_MESSAGES.ACCESS_PROHIBITED, 

577 ) 

578 

579 result = await image_generations(request, form_data, user=user) 

580 await publish_event( 

581 request, 

582 EVENTS.IMAGE_GENERATED, 

583 actor=user, 

584 subject_id=None, 

585 subject_type='image', 

586 data={ 

587 'model': form_data.model, 

588 'size': form_data.size, 

589 'n': form_data.n, 

590 'prompt_preview': form_data.prompt[:300], 

591 }, 

592 ) 

593 return result 

594 

595 

596async def image_generations( 

597 request: Request, 

598 form_data: CreateImageForm, 

599 metadata: dict | None = None, 

600 user=None, 

601): 

602 image_config = await get_image_config() 

603 # if IMAGE_SIZE = 'auto', default WidthxHeight to the 512x512 default 

604 # This is only relevant when the user has set IMAGE_SIZE to 'auto' with an 

605 # image model other than gpt-image-1, which is warned about on settings save 

606 

607 size = '512x512' 

608 if image_config.IMAGE_SIZE and 'x' in image_config.IMAGE_SIZE: 608 ↛ 611line 608 didn't jump to line 611 because the condition on line 608 was always true

609 size = image_config.IMAGE_SIZE 

610 

611 if form_data.size and 'x' in form_data.size: 611 ↛ 612line 611 didn't jump to line 612 because the condition on line 611 was never true

612 size = form_data.size 

613 

614 width, height = tuple(map(int, size.split('x'))) 

615 

616 metadata = metadata or {} 

617 

618 model = await get_image_model(request) 

619 

620 try: 

621 if image_config.IMAGE_GENERATION_ENGINE == 'openai': 

622 headers = { 

623 'Authorization': f'Bearer {image_config.IMAGES_OPENAI_API_KEY}', 

624 'Content-Type': 'application/json', 

625 } 

626 

627 if ENABLE_FORWARD_USER_INFO_HEADERS: 627 ↛ 628line 627 didn't jump to line 628 because the condition on line 627 was never true

628 headers = include_user_info_headers(headers, user) 

629 

630 url = f'{image_config.IMAGES_OPENAI_API_BASE_URL}/images/generations' 

631 if image_config.IMAGES_OPENAI_API_VERSION: 631 ↛ 632line 631 didn't jump to line 632 because the condition on line 631 was never true

632 url = f'{url}?api-version={image_config.IMAGES_OPENAI_API_VERSION}' 

633 

634 data = { 

635 'model': model, 

636 'prompt': form_data.prompt, 

637 'n': form_data.n, 

638 **( 

639 {'size': form_data.size or image_config.IMAGE_SIZE} 

640 if (form_data.size or image_config.IMAGE_SIZE) 

641 else {} 

642 ), 

643 **( 

644 {} 

645 if re.match( 

646 IMAGE_URL_RESPONSE_MODELS_REGEX_PATTERN, 

647 image_config.IMAGE_GENERATION_MODEL, 

648 ) 

649 else {'response_format': 'b64_json'} 

650 ), 

651 **({} if not image_config.IMAGES_OPENAI_API_PARAMS else image_config.IMAGES_OPENAI_API_PARAMS), 

652 } 

653 

654 session = await get_session() 

655 async with session.post( 

656 url=url, 

657 json=data, 

658 headers=headers, 

659 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

660 ) as r: 

661 r.raise_for_status() 

662 res = await r.json(content_type=None) 

663 

664 images = [] 

665 

666 for image in res['data']: 

667 if image_url := image.get('url', None): 

668 image_data, content_type = await get_image_data( 

669 image_url, 

670 {k: v for k, v in headers.items() if k != 'Content-Type'} 

671 if _is_same_origin(image_url, image_config.IMAGES_OPENAI_API_BASE_URL) 

672 else None, 

673 ) 

674 else: 

675 image_data, content_type = await get_image_data(image['b64_json']) 

676 

677 _, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user) 

678 images.append(image_file) 

679 return images 

680 

681 elif image_config.IMAGE_GENERATION_ENGINE == 'gemini': 681 ↛ 682line 681 didn't jump to line 682 because the condition on line 681 was never true

682 headers = { 

683 'Content-Type': 'application/json', 

684 'x-goog-api-key': image_config.IMAGES_GEMINI_API_KEY, 

685 } 

686 

687 data = {} 

688 

689 if ( 

690 image_config.IMAGES_GEMINI_ENDPOINT_METHOD == '' 

691 or image_config.IMAGES_GEMINI_ENDPOINT_METHOD == 'predict' 

692 ): 

693 model = f'{model}:predict' 

694 data = { 

695 'instances': {'prompt': form_data.prompt}, 

696 'parameters': { 

697 'sampleCount': form_data.n, 

698 'outputOptions': {'mimeType': 'image/png'}, 

699 }, 

700 } 

701 

702 elif image_config.IMAGES_GEMINI_ENDPOINT_METHOD == 'generateContent': 

703 model = f'{model}:generateContent' 

704 data = {'contents': [{'parts': [{'text': form_data.prompt}]}]} 

705 

706 session = await get_session() 

707 async with session.post( 

708 url=f'{image_config.IMAGES_GEMINI_API_BASE_URL}/models/{model}', 

709 json=data, 

710 headers=headers, 

711 ssl=AIOHTTP_CLIENT_SESSION_SSL, 

712 ) as r: 

713 r.raise_for_status() 

714 res = await r.json(content_type=None) 

715 

716 images = [] 

717 

718 if model.endswith(':predict'): 

719 for image in res['predictions']: 

720 image_data, content_type = await get_image_data(image['bytesBase64Encoded']) 

721 _, image_file = await upload_image(request, image_data, content_type, {**data, **metadata}, user) 

722 images.append(image_file) 

723 elif model.endswith(':generateContent'): 

724 for image in res['candidates']: 

725 for part in image['content']['parts']: 

726 if part.get('inlineData', {}).get('data'): 

727 image_data, content_type = await get_image_data(part['inlineData']['data']) 

728 _, image_file = await upload_image( 

729 request, 

730 image_data, 

731 content_type, 

732 {**data, **metadata}, 

733 user, 

734 ) 

735 images.append(image_file) 

736 

737 return images 

738 

739 elif image_config.IMAGE_GENERATION_ENGINE == 'comfyui': 739 ↛ 740line 739 didn't jump to line 740 because the condition on line 739 was never true

740 data = { 

741 'prompt': form_data.prompt, 

742 'width': width, 

743 'height': height, 

744 'n': form_data.n, 

745 } 

746 

747 if image_config.IMAGE_STEPS is not None or form_data.steps is not None: 

748 data['steps'] = form_data.steps if form_data.steps is not None else image_config.IMAGE_STEPS 

749 

750 if form_data.negative_prompt is not None: 

751 data['negative_prompt'] = form_data.negative_prompt 

752 

753 form_data = ComfyUICreateImageForm( 

754 **{ 

755 'workflow': ComfyUIWorkflow( 

756 **{ 

757 'workflow': image_config.COMFYUI_WORKFLOW, 

758 'nodes': image_config.COMFYUI_WORKFLOW_NODES, 

759 } 

760 ), 

761 **data, 

762 } 

763 ) 

764 res = await comfyui_create_image( 

765 model, 

766 form_data, 

767 str(uuid.uuid4()), 

768 image_config.COMFYUI_BASE_URL, 

769 image_config.COMFYUI_API_KEY, 

770 ) 

771 log.debug('res: %s', res) 

772 

773 images = [] 

774 

775 for image in res['data']: 

776 headers = None 

777 if image_config.COMFYUI_API_KEY: 

778 headers = {'Authorization': f'Bearer {image_config.COMFYUI_API_KEY}'} 

779 

780 image_data, content_type = await get_image_data( 

781 image['url'], 

782 headers, 

783 trusted_base_url=image_config.COMFYUI_BASE_URL, 

784 ) 

785 _, image_file = await upload_image( 

786 request, 

787 image_data, 

788 content_type, 

789 {**form_data.model_dump(exclude_none=True), **metadata}, 

790 user, 

791 ) 

792 images.append(image_file) 

793 return images 

794 elif image_config.IMAGE_GENERATION_ENGINE == 'automatic1111' or image_config.IMAGE_GENERATION_ENGINE == '': 794 ↛ 799line 794 didn't jump to line 799 because the condition on line 794 was never true

795 # Automatic1111 holds one checkpoint instance-wide, so set_image_model 

796 # persists the global default and switches the shared backend. Only an 

797 # admin may do that; a non-admin generates on the currently configured 

798 # checkpoint. The model field is not a per-user selection on this backend.