ausboss
f601467eef
With Resize on, Image Crop + Rotate + Pad and both Video Crop + Rotate + Pad nodes resized the padded canvas as one picture. Along every seam that mixed fill into the picture and left the mask part on, part off. A hard paste (stitch blend 0) kept that column as a faint line, and a band under a pixel wide after the resize came out as nothing but that column: 1 px of padding at a 0.62 scale left a line in the stitched result whatever the model painted. resize_then_pad (nodes/_transform_engine.py) resizes only the picture, with the chosen method, into the box it has on the resized canvas, and puts the fill around it. Bands are whole pixels, the picture's edge is picture, and a band that comes to under half a pixel is left out. The mask is 1 on the fill; inside the box it is the transform's own mask there (the feather's ramp, the corners of a turn, see-through parts), resized bilinear with the picture. The prompt view is placed the same way. The box is transform_picture_box (nodes/_inpaint_crop_helpers.py), which the stitcher now uses for its source bbox too, so the paste and the bbox cannot disagree. It is never empty. The output size is unchanged. With Resize off, or with no padding, image and mask are bit for bit what they were. With padding, the picture can sit up to half a pixel from where it did.
587 lines
26 KiB
Python
587 lines
26 KiB
Python
"""Pure rotate, crop, and pad processing shared by the AusBoss loaders."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
from ._execution_helpers import advance_progress, frame_progress, raise_if_interrupted
|
|
from PIL import Image, ImageFilter
|
|
import torch
|
|
|
|
from ._color_helpers import parse_fill_color
|
|
from ._inpaint_crop_helpers import transform_picture_box
|
|
from ._pad_helpers import _axis_layout
|
|
|
|
|
|
MAX_DIMENSION = 65536
|
|
MAX_PADDING = 32768
|
|
|
|
# A source pixel stays picture when its alpha is at least this share, in
|
|
# percent, of the most solid pixel's; anything more see-through is area to
|
|
# paint, like the padding. Measured on stranger-test cutouts: an oval photo
|
|
# whose edge ramp stores colours darkened by their alpha keeps a rim at most
|
|
# 10% dark at 90 (half dark at 50, the ring the model painted back), while
|
|
# matting noise inside solid subjects sits above 90, so no holes open there.
|
|
# Relative to the most solid pixel, a picture saved at 50% opacity throughout
|
|
# is kept whole instead of painted over. A picture at least this solid
|
|
# everywhere has no see-through part and goes through exactly as before
|
|
# (generated pictures often carry alpha 249-254 in places).
|
|
SEE_THROUGH_KEEP_PERCENT = 90
|
|
# How see-through parts are shown to a prompt writer (prompt_image) and in
|
|
# the untransformed `original`: on white, as a picture viewer shows them.
|
|
SEE_THROUGH_BACKDROP = (255, 255, 255)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TransformSpec:
|
|
rotation_degrees: float = 0.0
|
|
crop_aspect_ratio: str = "free"
|
|
crop_x: int = 0
|
|
crop_y: int = 0
|
|
crop_width: int = 0
|
|
crop_height: int = 0
|
|
pad_left: int = 0
|
|
pad_top: int = 0
|
|
pad_right: int = 0
|
|
pad_bottom: int = 0
|
|
feather: int = 0
|
|
canvas_multiple: int = 1
|
|
fill_color: str = "#808080"
|
|
|
|
def normalized(self) -> "TransformSpec":
|
|
angle = ((float(self.rotation_degrees) + 180.0) % 360.0) - 180.0
|
|
if abs(angle) < 0.00005:
|
|
angle = 0.0
|
|
return TransformSpec(
|
|
rotation_degrees=round(angle, 4),
|
|
crop_aspect_ratio=normalize_aspect_ratio(self.crop_aspect_ratio),
|
|
crop_x=max(0, int(self.crop_x)),
|
|
crop_y=max(0, int(self.crop_y)),
|
|
crop_width=max(0, int(self.crop_width)),
|
|
crop_height=max(0, int(self.crop_height)),
|
|
pad_left=_bounded_padding(self.pad_left, "pad_left"),
|
|
pad_top=_bounded_padding(self.pad_top, "pad_top"),
|
|
pad_right=_bounded_padding(self.pad_right, "pad_right"),
|
|
pad_bottom=_bounded_padding(self.pad_bottom, "pad_bottom"),
|
|
feather=max(0, min(int(self.feather), 4096)),
|
|
canvas_multiple=max(1, min(int(self.canvas_multiple), 4096)),
|
|
fill_color=normalize_fill_color(self.fill_color),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TransformGeometry:
|
|
rotated_width: int
|
|
rotated_height: int
|
|
crop_x: int
|
|
crop_y: int
|
|
crop_width: int
|
|
crop_height: int
|
|
pad_left: int
|
|
pad_top: int
|
|
pad_right: int
|
|
pad_bottom: int
|
|
output_width: int
|
|
output_height: int
|
|
|
|
|
|
def _bounded_padding(value: int, name: str) -> int:
|
|
value = int(value)
|
|
if value < 0 or value > MAX_PADDING:
|
|
raise ValueError(
|
|
f"Transform: input '{name}' expected 0..{MAX_PADDING}, received {value}."
|
|
)
|
|
return value
|
|
|
|
|
|
def normalize_fill_color(value: str) -> str:
|
|
return "#" + "".join(f"{channel:02x}" for channel in parse_fill_color(value))
|
|
|
|
|
|
def fill_rgb(value: str) -> tuple[int, int, int]:
|
|
return parse_fill_color(value)
|
|
|
|
|
|
def normalize_aspect_ratio(value: str) -> str:
|
|
text = str(value or "free").strip().lower()
|
|
if text in {"free", "source"}:
|
|
return text
|
|
try:
|
|
width, height = (int(part) for part in text.split(":", 1))
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(
|
|
"Transform: input 'crop_aspect_ratio' expected free, source, or W:H."
|
|
) from exc
|
|
if width < 1 or height < 1:
|
|
raise ValueError("Transform: crop aspect ratio values must be positive.")
|
|
common = math.gcd(width, height)
|
|
return f"{width // common}:{height // common}"
|
|
|
|
|
|
def _validate_source(image: Image.Image) -> None:
|
|
width, height = image.size
|
|
if width < 1 or height < 1:
|
|
raise ValueError("Transform: source image must contain at least one pixel.")
|
|
if width > MAX_DIMENSION or height > MAX_DIMENSION:
|
|
raise ValueError(
|
|
f"Transform: source dimensions must be at most {MAX_DIMENSION} pixels per side."
|
|
)
|
|
|
|
|
|
def see_through_kept(image: Image.Image) -> np.ndarray | None:
|
|
"""[H, W] bool map of the pixels that stay picture; the rest are see-through.
|
|
|
|
None when every pixel is at least SEE_THROUGH_KEEP_PERCENT solid, so an
|
|
opaque source takes exactly the path it always took. Otherwise a pixel
|
|
stays when its alpha is at least SEE_THROUGH_KEEP_PERCENT of the most
|
|
solid pixel's; a source with no solid pixel at all is see-through
|
|
everywhere.
|
|
"""
|
|
if "A" not in image.getbands():
|
|
return None
|
|
alpha = np.asarray(image.getchannel("A"))
|
|
if int(alpha.min()) * 100 >= 255 * SEE_THROUGH_KEEP_PERCENT:
|
|
return None
|
|
top = int(alpha.max())
|
|
return alpha.astype(np.int32) * 100 >= top * SEE_THROUGH_KEEP_PERCENT if top else np.zeros(alpha.shape, bool)
|
|
|
|
|
|
def _empty_see_through(rgba: Image.Image, kept: np.ndarray) -> Image.Image:
|
|
"""The source with its see-through pixels made empty, alpha 0 like a
|
|
rotation corner, and every kept pixel fully solid in its own colour.
|
|
|
|
The stored colour of a see-through pixel is often black or the old
|
|
background, and blending it over the fill left a dark ring that the model
|
|
painted back; with alpha 0 it never reaches the canvas (Pillow's turn
|
|
resamples with the alpha applied), and the mask marks it as area to paint.
|
|
"""
|
|
array = np.array(rgba, dtype=np.uint8)
|
|
array[..., 3] = np.where(kept, 255, 0).astype(np.uint8)
|
|
return Image.fromarray(array, "RGBA")
|
|
|
|
|
|
def _solid(image: Image.Image) -> Image.Image:
|
|
"""The source fully solid: every pixel in its stored colour."""
|
|
return image.convert("RGB").convert("RGBA")
|
|
|
|
|
|
def _rotate_rgba(image: Image.Image, spec: TransformSpec) -> Image.Image:
|
|
rgba = image.convert("RGBA")
|
|
if spec.rotation_degrees == 0.0:
|
|
return rgba.copy()
|
|
color = fill_rgb(spec.fill_color)
|
|
return rgba.rotate(
|
|
-spec.rotation_degrees,
|
|
resample=Image.Resampling.BICUBIC,
|
|
expand=True,
|
|
fillcolor=(*color, 0),
|
|
)
|
|
|
|
|
|
def _ratio_box(width: int, height: int, ratio_width: int, ratio_height: int) -> tuple[int, int]:
|
|
"""The largest ratio_width:ratio_height box inside width x height.
|
|
|
|
Integer maths: the float version lost a pixel on the limiting side
|
|
(21 * (3813 / 21) is 3812.99...). A box already within a pixel of the
|
|
ratio on either side is kept, so a resolved crop written back resolves
|
|
to itself. transform_geometry.mjs ratioBox is the same rule, which keeps
|
|
the editor's size readout equal to the run.
|
|
"""
|
|
fit_height = width * ratio_height // ratio_width
|
|
fit_width = height * ratio_width // ratio_height
|
|
if height == fit_height or width == fit_width:
|
|
return width, height
|
|
if width * ratio_height > height * ratio_width:
|
|
return max(1, fit_width), height
|
|
return width, max(1, fit_height)
|
|
|
|
|
|
def _geometry(rotated: Image.Image, spec: TransformSpec) -> TransformGeometry:
|
|
rotated_width, rotated_height = rotated.size
|
|
crop_x = min(spec.crop_x, rotated_width - 1)
|
|
crop_y = min(spec.crop_y, rotated_height - 1)
|
|
available_width = rotated_width - crop_x
|
|
available_height = rotated_height - crop_y
|
|
crop_width = available_width if spec.crop_width <= 0 else min(spec.crop_width, available_width)
|
|
crop_height = available_height if spec.crop_height <= 0 else min(spec.crop_height, available_height)
|
|
crop_width = max(1, crop_width)
|
|
crop_height = max(1, crop_height)
|
|
if spec.crop_aspect_ratio != "free":
|
|
if spec.crop_aspect_ratio == "source":
|
|
ratio_width, ratio_height = rotated_width, rotated_height
|
|
else:
|
|
ratio_width, ratio_height = (int(part) for part in spec.crop_aspect_ratio.split(":"))
|
|
crop_width, crop_height = _ratio_box(crop_width, crop_height, ratio_width, ratio_height)
|
|
|
|
# The canvas rounds to canvas_multiple one axis at a time, by Load Image +
|
|
# Pad's rule (_pad_helpers._axis_layout): the leftover joins a side the
|
|
# user padded, and an axis nobody padded has its crop trimmed evenly to
|
|
# the multiple below instead. The leftover used to go to the right and
|
|
# bottom whatever was padded, so an unpadded edge got a strip of fill a
|
|
# pixel or more wide; a resize then smeared it into the picture's last
|
|
# column, and the stitch kept that as a light line.
|
|
# transform_geometry.mjs resolvePadding is the same rule, which keeps
|
|
# the editor's picture and size readout equal to the run.
|
|
pad_left, pad_right, trim_left, _, output_width = _axis_layout(
|
|
crop_width, spec.pad_left, spec.pad_right, spec.canvas_multiple
|
|
)
|
|
pad_top, pad_bottom, trim_top, _, output_height = _axis_layout(
|
|
crop_height, spec.pad_top, spec.pad_bottom, spec.canvas_multiple
|
|
)
|
|
if output_width > MAX_DIMENSION or output_height > MAX_DIMENSION:
|
|
raise ValueError(
|
|
"Transform: requested crop and padding exceed the 65536-pixel output limit."
|
|
)
|
|
|
|
return TransformGeometry(
|
|
rotated_width=rotated_width,
|
|
rotated_height=rotated_height,
|
|
crop_x=crop_x + trim_left,
|
|
crop_y=crop_y + trim_top,
|
|
crop_width=output_width - pad_left - pad_right,
|
|
crop_height=output_height - pad_top - pad_bottom,
|
|
pad_left=pad_left,
|
|
pad_top=pad_top,
|
|
pad_right=pad_right,
|
|
pad_bottom=pad_bottom,
|
|
output_width=output_width,
|
|
output_height=output_height,
|
|
)
|
|
|
|
|
|
def transform_pil(image: Image.Image, spec: TransformSpec) -> tuple[Image.Image, Image.Image, TransformGeometry]:
|
|
"""Apply rotate -> crop -> pad and return opaque RGB, BHW-style mask image, and geometry."""
|
|
output, mask, geometry, _ = _transform_frame(image, spec, False)
|
|
return output, mask, geometry
|
|
|
|
|
|
def _transform_frame(
|
|
image: Image.Image, spec: TransformSpec, view: bool, keep_see_through: bool = False
|
|
) -> tuple[Image.Image, Image.Image, TransformGeometry, Image.Image | None]:
|
|
"""transform_pil, plus the prompt view (:func:`_see_through_view`) when
|
|
``view`` is set and the source has see-through parts; None otherwise.
|
|
|
|
With ``keep_see_through`` the stored colour under a see-through part
|
|
stays in the picture and only the mask marks it, so there is no prompt
|
|
view either: the picture is whole. That is what a mask painted in the
|
|
MaskEditor needs when a later step reads what is under the paint.
|
|
"""
|
|
_validate_source(image)
|
|
spec = spec.normalized()
|
|
rgba = image if image.mode == "RGBA" else image.convert("RGBA")
|
|
kept = see_through_kept(rgba)
|
|
if kept is not None:
|
|
rgba = _empty_see_through(rgba, kept)
|
|
rotated = _rotate_rgba(rgba, spec)
|
|
geometry = _geometry(rotated, spec)
|
|
crop_box = (
|
|
geometry.crop_x,
|
|
geometry.crop_y,
|
|
geometry.crop_x + geometry.crop_width,
|
|
geometry.crop_y + geometry.crop_height,
|
|
)
|
|
cropped = rotated.crop(crop_box)
|
|
|
|
# What the mask is cut from: the kept picture, see-through parts empty.
|
|
alpha = cropped.getchannel("A")
|
|
if kept is not None and keep_see_through:
|
|
# The picture itself goes through whole, turned and cropped the same
|
|
# way; its own alpha is only the corners a turn leaves empty.
|
|
cropped = _rotate_rgba(_solid(image), spec).crop(crop_box)
|
|
fill = fill_rgb(spec.fill_color)
|
|
filled_crop = Image.new("RGB", cropped.size, fill)
|
|
filled_crop.paste(cropped.convert("RGB"), mask=cropped.getchannel("A"))
|
|
|
|
output = Image.new("RGB", (geometry.output_width, geometry.output_height), fill)
|
|
output.paste(filled_crop, (geometry.pad_left, geometry.pad_top))
|
|
|
|
generated_crop = Image.fromarray(255 - np.asarray(alpha, dtype=np.uint8))
|
|
mask = Image.new("L", output.size, 255)
|
|
mask.paste(generated_crop, (geometry.pad_left, geometry.pad_top))
|
|
if spec.feather > 0:
|
|
original = np.asarray(mask, dtype=np.uint8)
|
|
blurred = np.asarray(
|
|
mask.filter(ImageFilter.GaussianBlur(spec.feather)), dtype=np.uint16
|
|
)
|
|
# A blurred step edge sits at ~50% exactly on the boundary, so using
|
|
# the blur directly (or max() with it) leaves a visible 255->127 seam.
|
|
# Doubling and clipping pins the generated side at 255 and starts a
|
|
# smooth ramp exactly at the edge; max() keeps thin generated slivers
|
|
# fully masked.
|
|
feathered = np.maximum(original, np.minimum(blurred * 2, 255).astype(np.uint8))
|
|
mask = Image.fromarray(feathered)
|
|
# Only the mask is feathered; the picture keeps its hard edge against
|
|
# the fill, as Load Image + Pad's does. Outpaint models key on a solid
|
|
# hard-edged fill and read a faded ramp as content (a darker band, or
|
|
# LTX's IC-LoRA leaving the bars unpainted), and the stitcher built
|
|
# from this canvas must hold real pixels for its color match.
|
|
|
|
prompt_view = None
|
|
if view and kept is not None and not keep_see_through:
|
|
prompt_view = _see_through_view(output, alpha, rgba.size, crop_box, spec, geometry)
|
|
return output, mask, geometry, prompt_view
|
|
|
|
|
|
def _see_through_view(
|
|
output: Image.Image,
|
|
picture_alpha: Image.Image,
|
|
source_size: tuple[int, int],
|
|
crop_box: tuple[int, int, int, int],
|
|
spec: TransformSpec,
|
|
geometry: TransformGeometry,
|
|
) -> Image.Image:
|
|
"""The canvas with its see-through parts on SEE_THROUGH_BACKDROP.
|
|
|
|
``picture_alpha`` is the kept picture's alpha, turned and cropped. The
|
|
source's whole rectangle, turned and cropped the same way, covers the
|
|
picture plus its see-through parts, so what it covers beyond the picture
|
|
is see-through; padding and the corners a turn leaves keep the fill.
|
|
A partly covered edge pixel keeps each share: picture, backdrop, fill.
|
|
"""
|
|
area = Image.new("L", source_size, 255)
|
|
if spec.rotation_degrees != 0.0:
|
|
area = area.rotate(
|
|
-spec.rotation_degrees, resample=Image.Resampling.BICUBIC, expand=True, fillcolor=0
|
|
)
|
|
area = np.asarray(area.crop(crop_box), dtype=np.int16)
|
|
share = np.clip(area - np.asarray(picture_alpha, dtype=np.int16), 0, 255)
|
|
canvas = np.array(output, dtype=np.uint8)
|
|
rows = np.flatnonzero(share.any(axis=1))
|
|
cols = np.flatnonzero(share.any(axis=0))
|
|
if rows.size == 0:
|
|
return Image.fromarray(canvas, "RGB")
|
|
# Only the see-through part's bounding box changes; whole-number maths:
|
|
# (2 * share * lift + 255) // 510 is share * lift / 255 rounded.
|
|
y0, y1, x0, x1 = int(rows[0]), int(rows[-1]) + 1, int(cols[0]), int(cols[-1]) + 1
|
|
lift = np.asarray(SEE_THROUGH_BACKDROP, dtype=np.int32) - np.asarray(fill_rgb(spec.fill_color), dtype=np.int32)
|
|
top, left = geometry.pad_top + y0, geometry.pad_left + x0
|
|
window = canvas[top : top + (y1 - y0), left : left + (x1 - x0)]
|
|
rise = (2 * share[y0:y1, x0:x1, None].astype(np.int32) * lift + 255) // 510
|
|
window[...] = np.clip(window.astype(np.int32) + rise, 0, 255).astype(np.uint8)
|
|
return Image.fromarray(canvas, "RGB")
|
|
|
|
|
|
def transform_pil_batch(
|
|
images: Iterable[Image.Image], spec: TransformSpec, *, view: bool = False, keep_see_through: bool = False
|
|
) -> tuple:
|
|
"""transform_pil over every frame, as BHWC image and BHW mask batches.
|
|
|
|
With ``view`` a fourth item follows: the prompt view batch, each frame's
|
|
canvas with its see-through parts on SEE_THROUGH_BACKDROP, or None when
|
|
no frame has any (the view is then the image itself). With
|
|
``keep_see_through`` the picture keeps what is stored under its
|
|
see-through parts and only the mask marks them.
|
|
"""
|
|
frames: list[torch.Tensor] = []
|
|
masks: list[torch.Tensor] = []
|
|
views: list[Image.Image] = []
|
|
first_geometry: TransformGeometry | None = None
|
|
for index, image in enumerate(images):
|
|
output, mask, geometry, prompt_view = _transform_frame(image, spec, view, keep_see_through)
|
|
if first_geometry is None:
|
|
first_geometry = geometry
|
|
elif output.size != (first_geometry.output_width, first_geometry.output_height):
|
|
raise ValueError(
|
|
f"Transform: frame {index} produced dimensions that differ from frame 0."
|
|
)
|
|
image_array = np.asarray(output, dtype=np.float32) / 255.0
|
|
mask_array = np.asarray(mask, dtype=np.float32) / 255.0
|
|
frames.append(torch.from_numpy(image_array.copy()))
|
|
masks.append(torch.from_numpy(mask_array.copy()))
|
|
views.append(prompt_view)
|
|
|
|
if not frames or first_geometry is None:
|
|
raise ValueError("Transform: source contained no decodable frames.")
|
|
result = (torch.stack(frames, dim=0), torch.stack(masks, dim=0), first_geometry)
|
|
if not view:
|
|
return result
|
|
if all(item is None for item in views):
|
|
return (*result, None)
|
|
shown = [
|
|
frames[index] if item is None else torch.from_numpy(np.asarray(item, dtype=np.float32) / 255.0)
|
|
for index, item in enumerate(views)
|
|
]
|
|
return (*result, torch.stack(shown, dim=0))
|
|
|
|
|
|
def transform_tensor_batch(
|
|
image: torch.Tensor, spec: TransformSpec, source_mask: torch.Tensor | None = None
|
|
) -> tuple[torch.Tensor, torch.Tensor, TransformGeometry]:
|
|
if not isinstance(image, torch.Tensor) or image.ndim != 4 or image.shape[-1] not in (3, 4):
|
|
received = tuple(image.shape) if isinstance(image, torch.Tensor) else type(image).__name__
|
|
raise ValueError(
|
|
"Transform: input 'image' expected BHWC with 3 or 4 channels, "
|
|
f"received {received}."
|
|
)
|
|
if source_mask is not None:
|
|
if not isinstance(source_mask, torch.Tensor) or tuple(source_mask.shape) != tuple(image.shape[:3]):
|
|
raise ValueError(
|
|
"Transform: input 'source_mask' expected BHW matching the image batch."
|
|
)
|
|
|
|
device = image.device
|
|
dtype = image.dtype
|
|
pil_frames: list[Image.Image] = []
|
|
image_cpu = image.detach().to(device="cpu", dtype=torch.float32).clamp(0.0, 1.0)
|
|
mask_cpu = None if source_mask is None else source_mask.detach().to(device="cpu", dtype=torch.float32).clamp(0.0, 1.0)
|
|
for index in range(image_cpu.shape[0]):
|
|
rgb = (image_cpu[index, ..., :3].numpy() * 255.0).round().astype(np.uint8)
|
|
if image_cpu.shape[-1] == 4:
|
|
alpha = (image_cpu[index, ..., 3].numpy() * 255.0).round().astype(np.uint8)
|
|
else:
|
|
alpha = np.full(image_cpu.shape[1:3], 255, dtype=np.uint8)
|
|
if mask_cpu is not None:
|
|
alpha = np.minimum(alpha, ((1.0 - mask_cpu[index].numpy()) * 255.0).round().astype(np.uint8))
|
|
pil_frames.append(Image.fromarray(np.dstack((rgb, alpha))))
|
|
|
|
output, mask, geometry = transform_pil_batch(pil_frames, spec)
|
|
return output.to(device=device, dtype=dtype), mask.to(device=device, dtype=dtype), geometry
|
|
|
|
|
|
def transform_tensor_batch_chunked(
|
|
image: torch.Tensor, spec: TransformSpec, chunk_size: int = 16
|
|
) -> tuple[torch.Tensor, torch.Tensor, TransformGeometry]:
|
|
"""transform_tensor_batch over a long batch, one chunk of frames at a time.
|
|
|
|
Every chunk goes through the same PIL round trip, but peak memory stays
|
|
at the finished output plus one chunk instead of every frame twice, and
|
|
the queue's interrupt and progress bar are serviced between chunks - the
|
|
video node feeds whole clips through here.
|
|
"""
|
|
if not isinstance(image, torch.Tensor) or image.ndim != 4:
|
|
received = tuple(image.shape) if isinstance(image, torch.Tensor) else type(image).__name__
|
|
raise ValueError(f"Transform: input 'image' expected a BHWC batch, received {received}.")
|
|
total = int(image.shape[0])
|
|
size = max(1, int(chunk_size))
|
|
output = mask = geometry = None
|
|
progress = frame_progress(total)
|
|
for start in range(0, total, size):
|
|
raise_if_interrupted()
|
|
chunk_output, chunk_mask, chunk_geometry = transform_tensor_batch(image[start : start + size], spec)
|
|
if output is None:
|
|
geometry = chunk_geometry
|
|
output = torch.empty((total, *chunk_output.shape[1:]), dtype=chunk_output.dtype, device=chunk_output.device)
|
|
mask = torch.empty((total, *chunk_mask.shape[1:]), dtype=chunk_mask.dtype, device=chunk_mask.device)
|
|
elif chunk_output.shape[1:] != output.shape[1:]:
|
|
raise ValueError(f"Transform: frame {start} produced dimensions that differ from frame 0.")
|
|
count = int(chunk_output.shape[0])
|
|
output[start : start + count] = chunk_output
|
|
mask[start : start + count] = chunk_mask
|
|
advance_progress(progress, min(start + count, total), total)
|
|
if output is None or mask is None or geometry is None:
|
|
raise ValueError("Transform: source contained no decodable frames.")
|
|
return output, mask, geometry
|
|
|
|
|
|
def stable_file_fingerprint(path: str | os.PathLike[str], inputs: dict[str, object]) -> str:
|
|
normalized_path = str(Path(path).resolve()).casefold()
|
|
try:
|
|
stat = os.stat(path)
|
|
file_state: dict[str, object] = {
|
|
"path": normalized_path,
|
|
"mtime_ns": stat.st_mtime_ns,
|
|
"size": stat.st_size,
|
|
}
|
|
except OSError:
|
|
file_state = {"path": normalized_path, "missing": True}
|
|
payload = {"file": file_state, "inputs": inputs}
|
|
encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
|
|
return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def scale_to_megapixels(
|
|
width: int, height: int, megapixels: float, steps: int = 1
|
|
) -> tuple[int, int]:
|
|
"""Target size whose pixel count is ~megapixels, aspect preserved.
|
|
|
|
Mirrors core ImageScaleToTotalPixels semantics so the two stay
|
|
interchangeable in a workflow: the budget is megapixels * 1024 * 1024,
|
|
and each dimension rounds independently to the nearest multiple of
|
|
steps (never below one step)."""
|
|
source_width = max(1, int(width))
|
|
source_height = max(1, int(height))
|
|
total = max(1.0, float(megapixels) * 1024 * 1024)
|
|
scale = (total / (source_width * source_height)) ** 0.5
|
|
step = max(1, int(steps))
|
|
scaled_width = max(step, round(source_width * scale / step) * step)
|
|
scaled_height = max(step, round(source_height * scale / step) * step)
|
|
return int(scaled_width), int(scaled_height)
|
|
|
|
|
|
def resize_then_pad(output, mask, geometry, fill_color, megapixels, method, steps, view=None):
|
|
"""A transform's canvas at a pixel budget: the picture resized, then padded.
|
|
|
|
``output`` and ``mask`` are a transform's BHWC canvas and BHW mask,
|
|
``view`` its prompt view or None. Returns the three at the size
|
|
:func:`scale_to_megapixels` gives the canvas.
|
|
|
|
Resizing the padded canvas as one picture mixed fill into the picture
|
|
along every seam: a column part fill with a part mask, which a hard
|
|
paste kept as a line, and a band under a pixel wide came out as nothing
|
|
else. So only the picture is resized, with the chosen method, into the
|
|
box it has on the resized canvas (transform_picture_box, the stitcher's
|
|
own source bbox), and the fill goes around it. Bands are whole pixels,
|
|
the picture's edge is picture, and a band that comes to under half a
|
|
pixel is gone.
|
|
|
|
The mask is 1 on the fill. Inside the box it is the transform's own
|
|
mask there, resized with the picture: the feather's ramp from each
|
|
seam, the corners of a turn and see-through parts. It resizes bilinear:
|
|
it is a soft coverage map, and a ringing method would push it outside
|
|
0..1. Without padding the box is the whole canvas, and this is a plain
|
|
resize of image and mask.
|
|
"""
|
|
import comfy.utils
|
|
|
|
height, width = int(output.shape[1]), int(output.shape[2])
|
|
target_width, target_height = scale_to_megapixels(width, height, megapixels, steps)
|
|
if (target_width, target_height) == (width, height):
|
|
return output, mask, view
|
|
x0, y0, x1, y1 = transform_picture_box(geometry, target_width, target_height)
|
|
rows = slice(geometry.pad_top, geometry.pad_top + geometry.crop_height)
|
|
cols = slice(geometry.pad_left, geometry.pad_left + geometry.crop_width)
|
|
fill = torch.tensor([channel / 255.0 for channel in fill_rgb(fill_color)])
|
|
|
|
def place(canvas):
|
|
picture = comfy.utils.common_upscale(
|
|
canvas[:, rows, cols].movedim(-1, 1), x1 - x0, y1 - y0, str(method), "disabled"
|
|
)
|
|
placed = fill.to(canvas).expand(canvas.shape[0], target_height, target_width, -1).clone()
|
|
placed[:, y0:y1, x0:x1] = picture.movedim(1, -1).clamp(0.0, 1.0)
|
|
return placed
|
|
|
|
resized = mask.new_ones((mask.shape[0], target_height, target_width))
|
|
inside = mask[:, rows, cols]
|
|
if bool(inside.any()):
|
|
inside = comfy.utils.common_upscale(inside.unsqueeze(1), x1 - x0, y1 - y0, "bilinear", "disabled")
|
|
resized[:, y0:y1, x0:x1] = inside.squeeze(1).clamp(0.0, 1.0)
|
|
else:
|
|
resized[:, y0:y1, x0:x1] = 0.0
|
|
return place(output), resized, None if view is None else place(view)
|
|
|
|
|
|
def original_image_batch(images: Iterable[Image.Image], keep_see_through: bool = False) -> torch.Tensor:
|
|
"""Untransformed, EXIF-oriented RGB source frames for reference outputs.
|
|
|
|
See-through parts (:func:`see_through_kept`) show on
|
|
SEE_THROUGH_BACKDROP, as a picture viewer shows them. Dropping the alpha
|
|
showed their stored colour instead, black for most cutouts, and a prompt
|
|
writer reading this output described a black backdrop. With
|
|
``keep_see_through`` the stored colour is what the caller asked for.
|
|
"""
|
|
frames = []
|
|
for image in images:
|
|
array = np.asarray(image.convert("RGB"), dtype=np.float32).copy()
|
|
kept = None if keep_see_through else see_through_kept(image if image.mode == "RGBA" else image.convert("RGBA"))
|
|
if kept is not None:
|
|
array[~kept] = SEE_THROUGH_BACKDROP
|
|
frames.append(torch.from_numpy(array / 255.0))
|
|
return torch.stack(frames)
|