USDU parity: progressive tiles, deterministic blend, crop/seed/model fixes

utils/crop_model_patch.py

@@ -0,0 +1,114 @@
+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

utils/usdu_utils.py

@@ -466,9 +466,13 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
     x1_px, y1_px, x2_px, y2_px = region
 
     new_latents = []
-    for t in latents:  # (B,C,H_lat_in,W_lat_in)
+    for t in latents:  # (B,C,H_lat_in,W_lat_in) or (B,C,1,H_lat_in,W_lat_in)
+        has_5d = False
+        if t.ndim == 5:
+            has_5d = True
+            t = t.squeeze(2)
         if t.ndim != 4:
-            raise ValueError(f"expected BCHW, got {t.shape}")
+            raise ValueError(f"expected BCHW or BC1HW, got {t.shape}")
 
         # 1. Resize to canvas resolution in latent units only if needed
         if t.shape[-2:] != (H_can_lat, W_can_lat):
@@ -491,6 +495,8 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
                                 mode="bilinear",
                                 align_corners=False)
 
+        if has_5d:
+            cropped = cropped.unsqueeze(2)
         new_latents.append(cropped)
 
     cond_dict["reference_latents"] = new_latents
@@ -508,4 +514,4 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
         crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
         cropped.append(n)
-    return cropped
+    return cropped