diff --git a/app/assets/services/preview.py b/app/assets/services/preview.py new file mode 100644 index 000000000..ffdf7e24b --- /dev/null +++ b/app/assets/services/preview.py @@ -0,0 +1,90 @@ +import logging +import os +import uuid + +from PIL import Image, ImageOps + +import folder_paths +from app.assets.database.queries.asset_reference import ( + get_reference_by_file_path, + get_reference_by_id, + set_reference_preview, +) +from app.assets.services.ingest import register_file_in_place +from app.database.db import create_session + +PREVIEW_TAG = "preview" + + +def get_or_create_preview_file( + source_abs_path: str, max_size: int, quality: int +) -> str: + source_hash, source_ref_id = _ensure_source_asset(source_abs_path) + hash_hex = source_hash.partition(":")[2] + preview_dir = folder_paths.get_system_user_directory("preview_cache") + preview_path = os.path.join( + preview_dir, f"{hash_hex}_{max_size}_q{quality}.webp" + ) + if os.path.isfile(preview_path): + return preview_path + + os.makedirs(preview_dir, exist_ok=True) + tmp_path = f"{preview_path}.{uuid.uuid4().hex}.tmp" + try: + with Image.open(source_abs_path) as img: + preview_img = img + if max(img.size) > max_size: + preview_img = ImageOps.contain( + img, (max_size, max_size), Image.Resampling.LANCZOS + ) + preview_img.save(tmp_path, format="webp", quality=quality) + os.replace(tmp_path, preview_path) + except Exception: + if os.path.exists(tmp_path): + try: + os.remove(tmp_path) + except OSError: + pass + raise + + _register_preview_asset(preview_path, source_ref_id) + return preview_path + + +def _ensure_source_asset(abs_path: str) -> tuple[str, str]: + abs_path = os.path.abspath(abs_path) + mtime_ns = os.stat(abs_path).st_mtime_ns + + with create_session() as session: + ref = get_reference_by_file_path(session, abs_path) + if ( + ref is not None + and ref.mtime_ns == mtime_ns + and ref.asset is not None + and ref.asset.hash + ): + return ref.asset.hash, ref.id + + result = register_file_in_place( + abs_path=abs_path, name=os.path.basename(abs_path), tags=[] + ) + if not result.asset.hash: + raise RuntimeError(f"asset registration produced no hash for {abs_path}") + return result.asset.hash, result.ref.id + + +def _register_preview_asset(preview_path: str, source_ref_id: str) -> None: + try: + result = register_file_in_place( + abs_path=preview_path, + name=os.path.basename(preview_path), + tags=[PREVIEW_TAG], + mime_type="image/webp", + ) + with create_session() as session: + source_ref = get_reference_by_id(session, source_ref_id) + if source_ref is not None and source_ref.preview_id is None: + set_reference_preview(session, source_ref_id, result.ref.id) + session.commit() + except Exception: + logging.warning("Failed to register preview image as asset", exc_info=True) diff --git a/openapi.yaml b/openapi.yaml index e00643bad..cf6b795cd 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -3808,7 +3808,24 @@ paths: /api/upload/mask: post: description: | - Upload a mask image to be applied to an existing image. + Upload a mask image to be composited onto an existing image server-side. + + The uploaded mask's alpha channel replaces the original image's alpha. + If the mask resolution differs from the original (e.g. the mask was + painted on a downscaled preview obtained via /api/view's preview and + max_size parameters), the mask alpha is upscaled to the original's + resolution before compositing, so the result keeps the original's full + resolution. + + An optional paint layer can be supplied to additionally produce painted + composites: 'painted' (original + paint) and 'painted_masked' + (original + paint + mask alpha), saved next to the main output under + the provided filenames. + + Clients that already have the fully composited result client-side + should upload it as a plain image via /api/upload/image instead; + this endpoint exists for server-side compositing at original + resolution. Image limits apply to both the uploaded mask and the referenced original image: @@ -3824,12 +3841,25 @@ paths: schema: properties: image: - description: The mask image file to upload + description: The mask image file to upload; its alpha channel is composited onto the original format: binary type: string original_ref: description: JSON string containing reference to the original image type: string + paint: + description: Optional RGBA paint-stroke layer to composite over the original + format: binary + type: string + paint_filename: + description: Optional filename to save the (upscaled) paint layer as, next to the main output + type: string + painted_filename: + description: Optional filename to save the original+paint composite as + type: string + painted_masked_filename: + description: Optional filename to save the original+paint+mask composite as + type: string required: - image - original_ref @@ -4450,6 +4480,32 @@ paths: maximum: 1024 minimum: 64 type: integer + - description: | + Compressed preview request in the form "[;]", e.g. "webp;90". + Format may be webp or jpeg (quality defaults to 90); requests that need the + alpha channel (channel containing 'a') are forced to an alpha-capable format. + Takes precedence over 'res' and 'channel' processing. + in: query + name: preview + schema: + example: webp;90 + type: string + - description: | + Maximum dimension (width or height) for the compressed preview, preserving + aspect ratio. Larger images are downscaled to fit; smaller images are never + upscaled. Values outside 512-8192 are clamped to that range. + Only used together with 'preview'. Requires the assets system + (--enable-assets): generated previews are keyed by the source's content + hash and registered as assets linked to the source via preview_id. When + assets are disabled, max_size is ignored and the image is only + recompressed at full resolution. + in: query + name: max_size + schema: + example: 4096 + maximum: 8192 + minimum: 512 + type: integer responses: "200": content: @@ -4463,7 +4519,12 @@ paths: description: Processed PNG image with extracted channel format: binary type: string - description: Success - File content returned (used when channel or res parameter is present) + image/webp: + schema: + description: Compressed preview (returned when the preview parameter is used) + format: binary + type: string + description: Success - File content returned (used when channel, res, or preview parameter is present) "302": description: Redirect to GCS signed URL headers: diff --git a/server.py b/server.py index e28fe2d22..0e438b3ee 100644 --- a/server.py +++ b/server.py @@ -49,6 +49,7 @@ from app.assets.api.routes import register_assets_routes from app.assets.services.ingest import register_file_in_place from app.assets.services.path_utils import get_known_subfolder_tags from app.assets.services.asset_management import resolve_hash_to_path +from app.assets.services.preview import get_or_create_preview_file from app.user_manager import UserManager from app.model_manager import ModelFileManager @@ -473,45 +474,101 @@ class PromptServer(): def image_save_function(image, post, filepath): original_ref = json.loads(post.get("original_ref")) - filename, output_dir = folder_paths.annotated_filepath(original_ref['filename']) + ref_filename = original_ref['filename'] - if not filename: - return web.Response(status=400) + if ref_filename.startswith("blake3:"): + owner_id = self.user_manager.get_request_user_id(request) + result = resolve_hash_to_path(ref_filename, owner_id=owner_id) + if result is None: + raise web.HTTPBadRequest() + file = result.abs_path + else: + filename, output_dir = folder_paths.annotated_filepath(ref_filename) - # validation for security: prevent accessing arbitrary path - if filename[0] == '/' or '..' in filename: - return web.Response(status=400) + if not filename: + raise web.HTTPBadRequest() - if output_dir is None: - type = original_ref.get("type", "output") - output_dir = folder_paths.get_directory_by_type(type) + # validation for security: prevent accessing arbitrary path + if filename[0] == '/' or '..' in filename: + raise web.HTTPBadRequest() - if output_dir is None: - return web.Response(status=400) + if output_dir is None: + type = original_ref.get("type", "output") + output_dir = folder_paths.get_directory_by_type(type) - if original_ref.get("subfolder", "") != "": - full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) - if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: - return web.Response(status=403) - output_dir = full_output_dir + if output_dir is None: + raise web.HTTPBadRequest() - file = os.path.join(output_dir, filename) + if original_ref.get("subfolder", "") != "": + full_output_dir = os.path.join(output_dir, original_ref["subfolder"]) + if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: + raise web.HTTPForbidden() + output_dir = full_output_dir - if os.path.isfile(file): - with Image.open(file) as original_pil: - metadata = PngInfo() - if hasattr(original_pil,'text'): - for key in original_pil.text: - metadata.add_text(key, original_pil.text[key]) - original_pil = original_pil.convert('RGBA') - mask_pil = Image.open(image.file).convert('RGBA') + file = os.path.join(output_dir, filename) - # alpha copy - new_alpha = mask_pil.getchannel('A') - original_pil.putalpha(new_alpha) - original_pil.save(filepath, compress_level=4, pnginfo=metadata) + if not os.path.isfile(file): + raise web.HTTPBadRequest() - return image_upload(post, image_save_function) + with Image.open(file) as original_pil: + metadata = PngInfo() + if hasattr(original_pil,'text'): + for key in original_pil.text: + metadata.add_text(key, original_pil.text[key]) + original_pil = original_pil.convert('RGBA') + original_pil.putalpha(Image.new('L', original_pil.size, 255)) + mask_pil = Image.open(image.file).convert('RGBA') + + # alpha copy; the mask may come from a downscaled + # preview edit, so upscale it to the original size + new_alpha = mask_pil.getchannel('A') + if new_alpha.size != original_pil.size: + new_alpha = new_alpha.resize(original_pil.size, Image.Resampling.LANCZOS) + + masked_pil = original_pil.copy() + masked_pil.putalpha(new_alpha) + masked_pil.save(filepath, compress_level=4, pnginfo=metadata) + + paint = post.get("paint") + if paint is not None and paint.file: + save_dir = os.path.dirname(filepath) + + paint_pil = Image.open(paint.file).convert('RGBA') + if paint_pil.size != original_pil.size: + paint_pil = paint_pil.resize(original_pil.size, Image.Resampling.LANCZOS) + painted_pil = Image.alpha_composite(original_pil, paint_pil) + painted_masked_pil = painted_pil.copy() + painted_masked_pil.putalpha(new_alpha) + + sibling_outputs = [ + ("paint_filename", paint_pil, {}), + ("painted_filename", painted_pil, {"pnginfo": metadata}), + ("painted_masked_filename", painted_masked_pil, {"pnginfo": metadata}), + ] + staged = [] + try: + for field, pil_image, save_kwargs in sibling_outputs: + name = post.get(field) + if not name: + continue + dest = os.path.join(save_dir, os.path.basename(name)) + tmp = f"{dest}.{uuid.uuid4().hex}.tmp" + pil_image.save(tmp, format="PNG", compress_level=4, **save_kwargs) + staged.append((tmp, dest)) + for tmp, dest in staged: + os.replace(tmp, dest) + except Exception: + for tmp, _ in staged: + if os.path.exists(tmp): + try: + os.remove(tmp) + except OSError: + pass + raise + + # Compositing large originals can take seconds; keep it off the event loop + return await asyncio.get_running_loop().run_in_executor( + None, image_upload, post, image_save_function) @routes.get("/view") async def view_image(request): @@ -557,24 +614,62 @@ class PromptServer(): if os.path.isfile(file): if 'preview' in request.rel_url.query: - with Image.open(file) as img: - preview_info = request.rel_url.query['preview'].split(';') - image_format = preview_info[0] - if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''): - image_format = 'webp' + preview_info = request.rel_url.query['preview'].split(';') + channel = request.rel_url.query.get('channel', '') + image_format = preview_info[0] + if image_format not in ['webp', 'jpeg'] or 'a' in channel: + image_format = 'webp' - quality = 90 - if preview_info[-1].isdigit(): - quality = int(preview_info[-1]) + quality = 90 + if preview_info[-1].isdigit(): + quality = int(preview_info[-1]) - buffer = BytesIO() - if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb': - img = img.convert("RGB") - img.save(buffer, format=image_format, quality=quality) - buffer.seek(0) + # Clamp any positive max_size into the permitted range so + # out-of-range callers still get a valid preview. + max_size = None + max_size_param = request.rel_url.query.get('max_size', '') + if max_size_param.isdigit() and int(max_size_param) > 0: + max_size = min(8192, max(512, int(max_size_param))) - return web.Response(body=buffer.read(), content_type=f'image/{image_format}', - headers={"Content-Disposition": f"filename=\"{filename}\""}) + safe_filename = filename.replace("\\", "\\\\").replace('"', '\\"') + preview_headers = {"Content-Disposition": f'filename="{safe_filename}"'} + loop = asyncio.get_running_loop() + + preview_file = None + preview_failed = False + if max_size is not None and args.enable_assets: + try: + preview_file = await loop.run_in_executor( + None, get_or_create_preview_file, file, max_size, quality) + except Exception: + logging.warning("Failed to generate preview asset, downscaling in memory without caching", exc_info=True) + preview_failed = True + + needs_convert = image_format == 'jpeg' or channel == 'rgb' + if preview_file is not None and not needs_convert: + return web.FileResponse(preview_file, headers={ + **preview_headers, + "Content-Type": "image/webp", + }) + + render_source = preview_file or file + + def render_preview(): + with Image.open(render_source) as img: + preview_img = img + + if preview_failed and max_size is not None and max(img.size) > max_size: + preview_img = ImageOps.contain( + img, (max_size, max_size), Image.Resampling.LANCZOS) + if needs_convert: + preview_img = preview_img.convert("RGB") + buffer = BytesIO() + preview_img.save(buffer, format=image_format, quality=quality) + return buffer.getvalue() + + body = await loop.run_in_executor(None, render_preview) + return web.Response(body=body, content_type=f'image/{image_format}', + headers=preview_headers) if 'channel' not in request.rel_url.query: channel = 'rgba'