feat(view): serve downscaled image previews via the assets system

This commit is contained in:
Terry Jia 2026-07-17 14:39:17 -04:00
parent 285a98944c
commit 676a6df4f8
3 changed files with 300 additions and 48 deletions

View File

@ -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)

View File

@ -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 "<format>[;<quality>]", e.g. "webp;90".
Format may be webp or jpeg; requests that need the alpha channel (channel
containing 'a') are forced to an alpha-capable format. Quality is an integer
from 1 to 100 (defaults to 90); values outside that range are ignored.
Takes precedence over 'res' and 'channel' processing.
in: query
name: preview
schema:
example: webp;90
type: string
- description: |
Requests a downscaled compressed preview, giving the maximum dimension
(width or height) while preserving aspect ratio. Larger images are
downscaled to fit; smaller images are never upscaled. Any integer is
accepted and clamped to 512-8192; omit the parameter for no downscale.
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
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:

191
server.py
View File

@ -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,68 @@ 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():
parsed_quality = int(preview_info[-1])
if 1 <= parsed_quality <= 100:
quality = parsed_quality
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)
# A max_size (any integer) requests downscaling and is
# clamped into the permitted range; absent or non-integer
# means no downscale.
max_size = None
max_size_param = request.rel_url.query.get('max_size', '')
if max_size_param:
try:
max_size = min(8192, max(512, int(max_size_param)))
except ValueError:
max_size = None
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'