115 lines
4.1 KiB
Python
115 lines
4.1 KiB
Python
from contextlib import contextmanager
|
|
|
|
import torch
|
|
|
|
from .logging import debug_log
|
|
from .usdu_utils import resize_region
|
|
|
|
|
|
@contextmanager
|
|
def crop_model_cond(model, crop_regions, init_size, canvas_size, tile_size, latent_crop=False):
|
|
"""Clone model and crop compatible model patches for tile-local sampling."""
|
|
try:
|
|
patched_model = model.clone()
|
|
except Exception:
|
|
# Fallback to original model when clone/patch access is unavailable.
|
|
yield model
|
|
return
|
|
|
|
patches = (
|
|
patched_model
|
|
.model_options
|
|
.get("transformer_options", {})
|
|
.get("patches", {})
|
|
)
|
|
applied_croppers = {}
|
|
for _module, module_patches in patches.items():
|
|
for patch in module_patches:
|
|
if id(patch) in applied_croppers:
|
|
continue
|
|
if type(patch).__name__ not in ("DiffSynthCnetPatch", "ZImageControlPatch"):
|
|
continue
|
|
try:
|
|
cropper = ModelPatchCropper(patch).crop(crop_regions, canvas_size, latent_crop)
|
|
applied_croppers[id(patch)] = cropper
|
|
except Exception as exc:
|
|
debug_log(f"crop_model_cond: patch crop skipped for {type(patch).__name__}: {exc}")
|
|
try:
|
|
yield patched_model
|
|
finally:
|
|
for cropper in applied_croppers.values():
|
|
del cropper
|
|
|
|
|
|
class ModelPatchCropper:
|
|
"""Stateful crop helper that restores model patch tensors on cleanup."""
|
|
|
|
def __init__(self, patch):
|
|
self.patch = patch
|
|
self.original_state = {
|
|
"image": patch.image.clone() if isinstance(patch.image, torch.Tensor) else patch.image,
|
|
"encoded_image": patch.encoded_image.clone() if isinstance(patch.encoded_image, torch.Tensor) else patch.encoded_image,
|
|
"encoded_image_size": patch.encoded_image_size,
|
|
}
|
|
self.patch_class = type(patch).__name__
|
|
required_attrs = (
|
|
"image",
|
|
"model_patch",
|
|
"vae",
|
|
"strength",
|
|
"encoded_image",
|
|
"encoded_image_size",
|
|
)
|
|
missing = [attr for attr in required_attrs if not hasattr(patch, attr)]
|
|
if missing:
|
|
raise AttributeError(
|
|
f"{self.patch_class} missing required attrs: {', '.join(missing)}"
|
|
)
|
|
|
|
def __del__(self):
|
|
self.patch.image = self.original_state["image"]
|
|
self.patch.encoded_image = self.original_state["encoded_image"]
|
|
self.patch.encoded_image_size = self.original_state["encoded_image_size"]
|
|
|
|
def crop(self, crop_regions, canvas_size, latent_crop=True):
|
|
patch = self.patch
|
|
|
|
if not isinstance(crop_regions, list):
|
|
crop_regions = [crop_regions]
|
|
|
|
image_size = (patch.image.shape[2], patch.image.shape[1]) # (W,H)
|
|
|
|
cropped_images = []
|
|
for crop_region in crop_regions:
|
|
resized_crop = resize_region(crop_region, canvas_size, image_size)
|
|
x1, y1, x2, y2 = resized_crop
|
|
cropped_image = patch.image[:, y1:y2, x1:x2, :]
|
|
cropped_images.append(cropped_image)
|
|
|
|
concatenated_image = torch.cat(cropped_images, dim=0)
|
|
patch.image = concatenated_image
|
|
patch.encoded_image_size = (
|
|
concatenated_image.shape[1],
|
|
concatenated_image.shape[2],
|
|
)
|
|
|
|
if latent_crop:
|
|
downscale_ratio = patch.vae.spacial_compression_encode()
|
|
cropped_latents = []
|
|
for crop_region in crop_regions:
|
|
resized_crop = resize_region(crop_region, canvas_size, image_size)
|
|
x1, y1, x2, y2 = tuple(x // downscale_ratio for x in resized_crop)
|
|
cropped_latent = patch.encoded_image[:, :, y1:y2, x1:x2]
|
|
cropped_latents.append(cropped_latent)
|
|
patch.encoded_image = torch.cat(cropped_latents, dim=0)
|
|
else:
|
|
patch.__init__(
|
|
patch.model_patch,
|
|
patch.vae,
|
|
concatenated_image,
|
|
patch.strength,
|
|
inpaint_image=patch.inpaint_image,
|
|
mask=patch.mask,
|
|
)
|
|
return self
|