428 lines
15 KiB
Python
428 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from ..repair.repair_utils import (
|
|
build_canvas_info,
|
|
build_canvas_repair_mask,
|
|
compute_crop_box,
|
|
crop_image_to,
|
|
crop_mask_to,
|
|
expand_canvas_image,
|
|
expand_canvas_mask,
|
|
get_image_size,
|
|
get_mask_metrics,
|
|
make_noise_mask,
|
|
merge_masks,
|
|
preprocess_mask,
|
|
recommend_denoise,
|
|
resize_image_to,
|
|
resize_mask_to,
|
|
resolve_repair_scope,
|
|
resolve_work_size,
|
|
)
|
|
from ..utils.model_info import parse_jsonish_info
|
|
|
|
try:
|
|
import node_helpers
|
|
except Exception:
|
|
node_helpers = None
|
|
|
|
try:
|
|
import torch
|
|
except Exception:
|
|
torch = None
|
|
|
|
|
|
EDIT_INFO_TYPE = "LLS_EDIT_INFO"
|
|
BACKEND_MODE_CHOICES = ["auto", "sdxl", "flux"]
|
|
EDIT_SCOPE_CHOICES = ["auto", "region", "crop", "canvas"]
|
|
PREVIEW_MODE_CHOICES = ["final", "compare", "mask", "before_after"]
|
|
AUTO_RECOMMEND_CHOICES = ["enabled", "disabled"]
|
|
CANVAS_FILL_CHOICES = ["edge", "blur", "black", "white", "neutral"]
|
|
COLOR_MATCH_CHOICES = ["disabled", "mean_std", "histogram_simple"]
|
|
BRIGHTNESS_MATCH_CHOICES = ["disabled", "enabled"]
|
|
EDGE_FIX_CHOICES = ["none", "soft", "strong"]
|
|
|
|
|
|
def normalize_edit_info(edit_info):
|
|
info = parse_jsonish_info(edit_info)
|
|
info.setdefault("backend_name", "")
|
|
info.setdefault("routing_reason", "")
|
|
info.setdefault("execution_path", "")
|
|
info.setdefault("model_family", "SD1.5")
|
|
info.setdefault("model_role", "base")
|
|
info.setdefault("profile_id", "")
|
|
info.setdefault("backend_type", "none")
|
|
info.setdefault("sampler_strategy", "standard_k")
|
|
info.setdefault("loader_strategy", "")
|
|
info.setdefault("edit_scope", "region")
|
|
info.setdefault("edit_payload_version", "1.0")
|
|
return info
|
|
|
|
|
|
def set_conditioning_values(conditioning, values: dict[str, Any]):
|
|
if node_helpers is not None:
|
|
try:
|
|
return node_helpers.conditioning_set_values(conditioning, values)
|
|
except Exception:
|
|
pass
|
|
|
|
updated = []
|
|
for entry in conditioning:
|
|
if not isinstance(entry, (list, tuple)) or len(entry) != 2:
|
|
updated.append(entry)
|
|
continue
|
|
token, meta = entry
|
|
merged = dict(meta or {})
|
|
merged.update(values)
|
|
updated.append([token, merged])
|
|
return updated
|
|
|
|
|
|
def build_masked_pixel_image(image, mask, *, fill_value: float):
|
|
if torch is not None and isinstance(image, torch.Tensor):
|
|
mask_4d = mask.unsqueeze(-1).clamp(0.0, 1.0)
|
|
return (image * (1.0 - mask_4d)) + (float(fill_value) * mask_4d)
|
|
if hasattr(image, "masked_fill"):
|
|
return image.masked_fill(mask, fill_value)
|
|
raise RuntimeError("[LLS] image object does not support masked pixel preprocessing.")
|
|
|
|
|
|
def build_native_conditioning_payload(vae, work_image, work_mask, *, latent_source: str, masked_fill_value: float):
|
|
latent_samples = vae.encode(work_image)
|
|
masked_pixels = build_masked_pixel_image(work_image, work_mask, fill_value=masked_fill_value)
|
|
concat_latent_image = vae.encode(masked_pixels)
|
|
concat_mask = make_noise_mask(work_mask, latent_samples)
|
|
latent = {
|
|
"samples": latent_samples,
|
|
"noise_mask": concat_mask,
|
|
"source": latent_source,
|
|
}
|
|
return latent, concat_latent_image, concat_mask
|
|
|
|
|
|
def build_fallback_latent_payload(vae, work_image, work_mask, *, latent_source: str):
|
|
latent_samples = vae.encode(work_image)
|
|
return {
|
|
"samples": latent_samples,
|
|
"noise_mask": make_noise_mask(work_mask, latent_samples),
|
|
"source": latent_source,
|
|
}
|
|
|
|
|
|
def build_workspace(
|
|
image,
|
|
mask,
|
|
*,
|
|
edit_scope,
|
|
mask_grow,
|
|
mask_blur,
|
|
mask_threshold,
|
|
invert_mask,
|
|
crop_context,
|
|
crop_context_factor,
|
|
min_size,
|
|
max_size,
|
|
resize_mode,
|
|
expand_left,
|
|
expand_right,
|
|
expand_top,
|
|
expand_bottom,
|
|
canvas_fill,
|
|
auto_recommend,
|
|
):
|
|
original_width, original_height = get_image_size(image)
|
|
original_size = (original_width, original_height)
|
|
processed_mask = preprocess_mask(
|
|
mask,
|
|
original_size,
|
|
invert_mask=bool(invert_mask),
|
|
mask_threshold=float(mask_threshold),
|
|
mask_grow=int(mask_grow),
|
|
mask_blur=float(mask_blur),
|
|
)
|
|
mask_bbox, mask_area_ratio = get_mask_metrics(processed_mask, original_size)
|
|
canvas_expand = (
|
|
max(0, int(expand_left)),
|
|
max(0, int(expand_right)),
|
|
max(0, int(expand_top)),
|
|
max(0, int(expand_bottom)),
|
|
)
|
|
scope = resolve_repair_scope(
|
|
edit_scope,
|
|
mask_area_ratio=mask_area_ratio,
|
|
mask_bbox=mask_bbox,
|
|
image_size=original_size,
|
|
canvas_expand=canvas_expand,
|
|
)
|
|
|
|
if scope == "region":
|
|
return {
|
|
"edit_scope": "region",
|
|
"work_image": image,
|
|
"work_mask": resize_mask_to(processed_mask, original_width, original_height),
|
|
"original_size": original_size,
|
|
"work_size": original_size,
|
|
"crop_box": None,
|
|
"crop_scale": 1.0,
|
|
"canvas_expand": list(canvas_expand),
|
|
"original_box_in_canvas": None,
|
|
"mask_bbox": list(mask_bbox) if mask_bbox is not None else None,
|
|
"mask_area_ratio": mask_area_ratio,
|
|
"recommended_denoise": recommend_denoise("replace", "region", "native_fill", auto_recommend),
|
|
"canvas_fill": canvas_fill,
|
|
}
|
|
|
|
if scope == "crop":
|
|
if mask_bbox is None or mask_area_ratio <= 0.0:
|
|
raise RuntimeError("[LLS] crop edit scope requires a non-empty mask.")
|
|
crop_box = compute_crop_box(mask_bbox, original_size, int(crop_context), float(crop_context_factor))
|
|
crop_width = max(1, crop_box[2] - crop_box[0])
|
|
crop_height = max(1, crop_box[3] - crop_box[1])
|
|
work_width, work_height, crop_scale = resolve_work_size(
|
|
(crop_width, crop_height),
|
|
int(min_size),
|
|
int(max_size),
|
|
resize_mode,
|
|
)
|
|
work_image = resize_image_to(crop_image_to(image, crop_box), work_width, work_height)
|
|
work_mask = resize_mask_to(crop_mask_to(processed_mask, crop_box), work_width, work_height)
|
|
return {
|
|
"edit_scope": "crop",
|
|
"work_image": work_image,
|
|
"work_mask": work_mask,
|
|
"original_size": original_size,
|
|
"work_size": (work_width, work_height),
|
|
"crop_box": list(crop_box),
|
|
"crop_scale": crop_scale,
|
|
"canvas_expand": list(canvas_expand),
|
|
"original_box_in_canvas": None,
|
|
"mask_bbox": None,
|
|
"mask_area_ratio": mask_area_ratio,
|
|
"recommended_denoise": recommend_denoise("replace", "crop", "native_fill", auto_recommend),
|
|
"canvas_fill": canvas_fill,
|
|
}
|
|
|
|
canvas_info = build_canvas_info(
|
|
original_size,
|
|
int(expand_left),
|
|
int(expand_right),
|
|
int(expand_top),
|
|
int(expand_bottom),
|
|
)
|
|
work_width, work_height = canvas_info["work_size"]
|
|
original_box = canvas_info["original_box"]
|
|
user_mask = expand_canvas_mask(processed_mask, work_width, work_height, original_box=original_box)
|
|
canvas_mask = build_canvas_repair_mask(user_mask, work_width, work_height, original_box=original_box)
|
|
work_mask = merge_masks(user_mask, canvas_mask, (work_width, work_height))
|
|
return {
|
|
"edit_scope": "canvas",
|
|
"work_image": expand_canvas_image(
|
|
image,
|
|
work_width,
|
|
work_height,
|
|
fill_mode=canvas_fill,
|
|
original_box=original_box,
|
|
),
|
|
"work_mask": work_mask,
|
|
"original_size": original_size,
|
|
"work_size": (work_width, work_height),
|
|
"crop_box": None,
|
|
"crop_scale": 1.0,
|
|
"canvas_expand": list(canvas_expand),
|
|
"original_box_in_canvas": list(original_box),
|
|
"mask_bbox": None,
|
|
"mask_area_ratio": mask_area_ratio,
|
|
"recommended_denoise": recommend_denoise("fill", "canvas", "native_fill", auto_recommend),
|
|
"canvas_fill": canvas_fill,
|
|
}
|
|
|
|
|
|
def build_edit_info(workspace: dict[str, Any], routing, *, backend_hints: dict[str, Any] | None = None) -> dict[str, Any]:
|
|
profile = routing.profile
|
|
info = {
|
|
"backend_name": routing.backend_name,
|
|
"routing_reason": routing.routing_reason,
|
|
"execution_path": routing.execution_path,
|
|
"model_family": profile["family"],
|
|
"model_role": profile["role"],
|
|
"profile_id": profile["profile_id"],
|
|
"backend_type": profile["backend_type"],
|
|
"sampler_strategy": profile["sampler_strategy"],
|
|
"loader_strategy": profile["loader_strategy"],
|
|
"supports_inpaint_native": profile["supports_inpaint_native"],
|
|
"supports_image_edit_native": profile["supports_image_edit_native"],
|
|
"preferred_edit_backend": profile["preferred_edit_backend"],
|
|
"edit_scope": workspace["edit_scope"],
|
|
"original_size": list(workspace["original_size"]),
|
|
"work_size": list(workspace["work_size"]),
|
|
"crop_box": workspace["crop_box"],
|
|
"crop_scale": workspace["crop_scale"],
|
|
"canvas_expand": workspace["canvas_expand"],
|
|
"original_box_in_canvas": workspace["original_box_in_canvas"],
|
|
"recommended_denoise": float(workspace["recommended_denoise"]),
|
|
"edit_payload_version": "1.0",
|
|
}
|
|
if backend_hints:
|
|
info.update(backend_hints)
|
|
return info
|
|
|
|
|
|
def _ensure_torch_available():
|
|
if torch is None:
|
|
raise RuntimeError("[LLS] Pro image edit compositing requires PyTorch.")
|
|
|
|
|
|
def _resize_mask_for_image(mask, image):
|
|
width, height = get_image_size(image)
|
|
return resize_mask_to(mask, width, height)
|
|
|
|
|
|
def _build_canvas_base(original_image, canvas_image, original_box):
|
|
_ensure_torch_available()
|
|
x1, y1, x2, y2 = original_box
|
|
out = torch.zeros_like(canvas_image)
|
|
out[:, y1:y2, x1:x2, :] = resize_image_to(original_image, x2 - x1, y2 - y1)
|
|
return out
|
|
|
|
|
|
def project_work_mask(work_mask, edit_info, reference_image):
|
|
_ensure_torch_available()
|
|
scope = str(edit_info.get("edit_scope") or "region")
|
|
if scope == "region":
|
|
return _resize_mask_for_image(work_mask, reference_image)
|
|
|
|
if scope == "crop":
|
|
crop_box = edit_info.get("crop_box")
|
|
if not isinstance(crop_box, (list, tuple)) or len(crop_box) != 4:
|
|
raise RuntimeError("[LLS] crop edit_info must include crop_box.")
|
|
x1, y1, x2, y2 = [int(value) for value in crop_box]
|
|
projected = torch.zeros(
|
|
(work_mask.shape[0], reference_image.shape[1], reference_image.shape[2]),
|
|
dtype=work_mask.dtype,
|
|
device=work_mask.device,
|
|
)
|
|
projected[:, y1:y2, x1:x2] = resize_mask_to(work_mask, x2 - x1, y2 - y1)
|
|
return projected
|
|
|
|
return _resize_mask_for_image(work_mask, reference_image)
|
|
|
|
|
|
def _blend_torch_images(original, generated, mask, blend_strength: float):
|
|
_ensure_torch_available()
|
|
mask_4d = _resize_mask_for_image(mask, original).unsqueeze(-1).clamp(0.0, 1.0) * float(blend_strength)
|
|
return (original * (1.0 - mask_4d)) + (generated * mask_4d)
|
|
|
|
|
|
def overlay_mask_preview(image, mask, *, alpha: float = 0.4):
|
|
if torch is None or not isinstance(image, torch.Tensor):
|
|
return image
|
|
mask_4d = _resize_mask_for_image(mask, image).unsqueeze(-1).clamp(0.0, 1.0) * float(alpha)
|
|
red = torch.zeros_like(image)
|
|
red[..., 0] = 1.0
|
|
return (image * (1.0 - mask_4d)) + (red * mask_4d)
|
|
|
|
|
|
def _paste_torch_patch(base, patch, box):
|
|
_ensure_torch_available()
|
|
x1, y1, x2, y2 = box
|
|
out = base.clone()
|
|
out[:, y1:y2, x1:x2, :] = patch
|
|
return out
|
|
|
|
|
|
def compose_region_result(
|
|
original_image,
|
|
generated_image,
|
|
work_mask,
|
|
edit_info,
|
|
feather,
|
|
color_match,
|
|
brightness_match,
|
|
blend_strength,
|
|
restore_unmasked_area,
|
|
edge_fix,
|
|
):
|
|
del edit_info, feather, color_match, brightness_match, restore_unmasked_area, edge_fix
|
|
original_width, original_height = get_image_size(original_image)
|
|
base_image = resize_image_to(original_image, original_width, original_height)
|
|
edited_image = resize_image_to(generated_image, original_width, original_height)
|
|
return _blend_torch_images(base_image, edited_image, _resize_mask_for_image(work_mask, base_image), blend_strength)
|
|
|
|
|
|
def compose_crop_result(
|
|
original_image,
|
|
generated_image,
|
|
work_mask,
|
|
edit_info,
|
|
feather,
|
|
color_match,
|
|
brightness_match,
|
|
blend_strength,
|
|
restore_unmasked_area,
|
|
edge_fix,
|
|
):
|
|
del feather, color_match, brightness_match, restore_unmasked_area, edge_fix
|
|
crop_box = edit_info.get("crop_box")
|
|
if not isinstance(crop_box, (list, tuple)) or len(crop_box) != 4:
|
|
raise RuntimeError("[LLS] crop edit_info must include crop_box.")
|
|
crop_width = max(1, int(crop_box[2]) - int(crop_box[0]))
|
|
crop_height = max(1, int(crop_box[3]) - int(crop_box[1]))
|
|
resized_patch = resize_image_to(generated_image, crop_width, crop_height)
|
|
resized_mask = resize_mask_to(work_mask, crop_width, crop_height)
|
|
original_patch = crop_image_to(original_image, tuple(int(v) for v in crop_box))
|
|
blended_patch = _blend_torch_images(original_patch, resized_patch, resized_mask, blend_strength)
|
|
return _paste_torch_patch(original_image, blended_patch, tuple(int(v) for v in crop_box))
|
|
|
|
|
|
def compose_canvas_result(
|
|
original_image,
|
|
generated_image,
|
|
work_mask,
|
|
edit_info,
|
|
feather,
|
|
color_match,
|
|
brightness_match,
|
|
blend_strength,
|
|
restore_unmasked_area,
|
|
edge_fix,
|
|
):
|
|
del feather, color_match, brightness_match, restore_unmasked_area, edge_fix
|
|
original_box = edit_info.get("original_box_in_canvas")
|
|
if not isinstance(original_box, (list, tuple)) or len(original_box) != 4:
|
|
raise RuntimeError("[LLS] canvas edit_info must include original_box_in_canvas.")
|
|
base_canvas = _build_canvas_base(
|
|
original_image,
|
|
generated_image,
|
|
tuple(int(value) for value in original_box),
|
|
)
|
|
return _blend_torch_images(base_canvas, generated_image, _resize_mask_for_image(work_mask, generated_image), blend_strength)
|
|
|
|
|
|
def build_preview_image(original_image, final_image, work_mask, edit_info, preview_mode):
|
|
_ensure_torch_available()
|
|
scope = str(edit_info.get("edit_scope") or "region")
|
|
if preview_mode == "final":
|
|
return final_image
|
|
if preview_mode == "mask":
|
|
if scope == "canvas":
|
|
preview_base = _build_canvas_base(
|
|
original_image,
|
|
final_image,
|
|
tuple(int(value) for value in edit_info["original_box_in_canvas"]),
|
|
)
|
|
projected_mask = _resize_mask_for_image(work_mask, final_image)
|
|
return overlay_mask_preview(preview_base, projected_mask)
|
|
projected_mask = project_work_mask(work_mask, edit_info, original_image)
|
|
return overlay_mask_preview(original_image, projected_mask)
|
|
if scope == "canvas":
|
|
compare_left = _build_canvas_base(
|
|
original_image,
|
|
final_image,
|
|
tuple(int(value) for value in edit_info["original_box_in_canvas"]),
|
|
)
|
|
else:
|
|
compare_left = resize_image_to(original_image, get_image_size(final_image)[0], get_image_size(final_image)[1])
|
|
return torch.cat([compare_left, final_image], dim=2)
|