diff --git a/.github/workflows/github-release.yml b/.github/workflows/github-release.yml new file mode 100644 index 0000000..b9f9df6 --- /dev/null +++ b/.github/workflows/github-release.yml @@ -0,0 +1,40 @@ +name: Publish GitHub release + +on: + push: + branches: [main] + paths: + - "pyproject.toml" + workflow_dispatch: + +permissions: + contents: write + +concurrency: + group: github-release-${{ github.ref }} + cancel-in-progress: false + +jobs: + github-release: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Read package version + id: package + run: | + version="$(python3 -c 'import tomllib; print(tomllib.load(open("pyproject.toml", "rb"))["project"]["version"])')" + echo "version=${version}" >> "${GITHUB_OUTPUT}" + - name: Create annotated release notes and tag + env: + GH_TOKEN: ${{ github.token }} + VERSION: ${{ steps.package.outputs.version }} + run: | + set -euo pipefail + tag="v${VERSION}" + notes="docs/releases/${tag}.md" + test -s "${notes}" + if gh release view "${tag}" >/dev/null 2>&1; then + echo "GitHub Release ${tag} already exists; leaving it unchanged" + exit 0 + fi + gh release create "${tag}" --target "${GITHUB_SHA}" --title "ComfyUI Scale-Locked Residual Diffusion ${tag}" --notes-file "${notes}" --latest diff --git a/.github/workflows/krea2-tests.yml b/.github/workflows/krea2-tests.yml new file mode 100644 index 0000000..775aebe --- /dev/null +++ b/.github/workflows/krea2-tests.yml @@ -0,0 +1,44 @@ +name: Krea 2 flow regeneration numerical tests + +on: + pull_request: + paths: + - "krea2_flow.py" + - "slrd_runtime.py" + - "slrd_core.py" + - "nodes.py" + - "resolution.py" + - "tests/test_krea2_flow.py" + - "tests/test_krea2_composition_lock.py" + - ".github/workflows/krea2-tests.yml" + push: + branches: [main] + paths: + - "krea2_flow.py" + - "slrd_runtime.py" + - "slrd_core.py" + - "nodes.py" + - "resolution.py" + - "tests/test_krea2_flow.py" + - "tests/test_krea2_composition_lock.py" + - ".github/workflows/krea2-tests.yml" + +jobs: + cpu-noise-invariants: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install CPU test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install pytest + - name: Compile changed Python + run: python -m py_compile krea2_flow.py nodes.py slrd_runtime.py slrd_core.py resolution.py + - name: Validate nested-noise coupling + working-directory: tests + run: python -m pytest -q --rootdir=. --confcutdir=. test_krea2_flow.py test_krea2_composition_lock.py diff --git a/README.md b/README.md index e8b789c..f4441aa 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,109 @@ In practice, this is meant to reduce: - scene re-planning between 1 MP and 4 MP, - the "looks like a different sample entirely" problem. +## Krea 2 Turbo: experimental flow-aligned regeneration + +The `Krea 2 Turbo Flow-Aligned Regenerate [Experimental]` node is a compact +Krea-2-specific integration of the existing scale-locked image sampling path. +It uses the native ComfyUI Krea2 model, the Qwen Image VAE and a matching +Krea2 text encoder. It is not a port of MiniMax H3's audio/video/Continuum +runtime and does not load weights independently. + +- Input `latent_image` defines the final pixel canvas (8-pixel VAE stride, + two-pixel latent DiT patch; use multiples of 16 pixels). ComfyUI's standard + four-channel empty latent is expanded to Krea 2's 16 channels at preflight. + Native single-frame `B×16×1×H×W` image latents are preserved through + planning, patch-coupled noise and spatial locking; multi-frame video latents + are not supported by this image node. +- The node defaults to an eight-step, approximately 1 MP planner pass, then + eight high-resolution steps: the first ones follow the planner trajectory + completely, and the remaining ones add detail at the target resolution. + The optional `steps` input controls both passes (default 8). Connect the same + step count to a native sampler for matched target-pass comparisons. +- The Krea 2 Turbo sampling profile is ComfyUI CFG **1.0** (no + classifier-free guidance), Euler, `simple` scheduler, shift **1.15**. + This is for **Turbo**, not the undistilled RAW checkpoint. +- High-resolution noise is drawn from an exact *Gaussian conditional coupling*: + within every 2x2 transformer-patch phase the normalized grouped projection + matches the source Gaussian field, including noninteger upscaling ratios. + This is a stochastic consistency invariant, not a learned image upscaler. +- `composition_lock_sigma` (optional, default 0.80) sets where the + high-resolution pass takes over. Target steps that start above this noise + level use the planner's denoised prediction in every frequency band, so the + target state is exactly the resized planner state carrying the target noise. + At those noise levels the high-resolution model would otherwise commit to its + own composition; mixing only the planner's coarse bands into it can + superimpose two layouts, visible as duplicated or displaced contours. Lower values + keep more of the planner image and leave fewer steps for new detail; 1.0 + disables the full lock. With ComfyUI's Krea 2 schedule (Euler/simple, + shift 1.15), the default locks four of eight steps (six of twelve). + Locked steps cost no high-resolution model evaluations: the target state at + the release is computed in closed form from the planner predictions, and + the sampler starts there. With a target spatial mask or latent noise mask, + locked steps are evaluated instead, because the mask blends in the + target prediction. +- **Resolution stages.** Large jumps from the planner to the output are split + into a ladder of intermediate resolutions, so no single handoff upscales by + more than `max_stage_upscale`. The first stage starts at the + composition-lock release from the planner predictions, exactly like a + single-stage run. Each later stage takes the finished image of the stage + below, resized to its own size, re-noises it to the same release noise + level with its own patch-coupled noise, and runs the remaining steps. Every + stage therefore repeats the bounded planner-to-target handoff from a sharp + image. Handing over mid-schedule predictions instead was observed to fail: + they are blurry noise-conditional averages, and a later, lower-noise start + rendered that blur as defocus with speckle. + - `max_stage_upscale` (default 4.0, i.e. 2x per side) is the largest + pixel-area growth per stage. Jumps larger than that get geometrically + spaced intermediate stages; 0 always jumps straight to the output. + - `stage_megapixels` lists explicit intermediate sizes, e.g. `4, 8`, and + overrides `max_stage_upscale`. + - Every stage runs all steps after the release, so each intermediate + stage adds that many evaluations at its own size. With eight steps and + the default lock, that is four per stage. +- After the release, the high-resolution pass refines freely by default. + `coarse_lock_strength` and `mid_lock_strength` (default 0) optionally keep + blending the low and middle spatial frequencies of each denoised prediction + toward the planner. Their bands are defined relative to the target canvas, + so at large sizes they include fine structures such as fingers, jewelry and + seams. The planner version of those structures is upscaled and blurred. + Nonzero values were observed to double fine contours and soften texture, + and are kept only for experiments. `composition_lock_sigma` is a starting + value, not a calibrated optimum. +- The `output` is the high-resolution LATENT, `lowres_planner` is the planner + LATENT, and `denoised_output` is the available final x0 preview LATENT. + +For an initial experiment, use a **1024x1024 planner (1.0 MP)** and a +**2048x2048 final latent (4.19 MP)**. Use the same native model, prompt, +seed, LoRAs and reference conditioning in both stages, all supplied through +the single node. The source/target widths need not have an integral ratio. +The official Krea 2 Turbo range is roughly 1K–2K per axis; larger target +canvases are experimental and may run out of memory or generate artifacts. + +A meaningful A/B must compare the same high-resolution model, eight steps, +seed, prompt and conditioning both **with and without** the two-pass method. +Connect one step-count input to both nodes when comparing other step budgets. +Judge composition/subject preservation, duplicated objects, typography, +fine-texture detail and wall time/VRAM. Real Krea 2 checks so far are +limited to individual visual comparisons, not a systematic evaluation; +there is no claimed quality or speed improvement. Against a native eight-step target-resolution run, the +default replaces four target-resolution DiT evaluations with eight planner +evaluations; it is intended to improve structural consistency, not +throughput. For RAW, custom samplers and advanced research controls, use +the generic Scale-Locked Residual KSampler instead. + +## Resolution Selector (Uncapped) + +`Resolution Selector (Uncapped)` (category `utilities`) behaves like ComfyUI's +built-in Resolution Selector: the same aspect ratios, megapixel convention +(1 MP = 1024x1024 pixels), `multiple` rounding and live resolution preview. +The built-in node caps `megapixels` at 16; this one accepts up to 256, the +largest square canvas ComfyUI latents allow (16384 pixels per side). A ratio +and size that would exceed 16384 pixels on either side is rejected with an +error instead of failing later in the empty-latent node. On ComfyUI versions +without the preview widget it falls back to the same calculation without the +preview. + ## Nodes ### 1. Scale-Locked Residual KSampler @@ -160,7 +263,7 @@ What is not implemented yet: - explicit residual-only tiled model execution, - sigma-perfect trajectory matching for samplers that perform unusual extra model evaluations, - scheduled cutoff animation for coarse or mid bands, -- multi-stage 1 MP -> 2 MP -> 4 MP progressive ladder inside one node, +- multi-stage progressive ladders in the generic Scale-Locked samplers (the Krea 2 Turbo node supports resolution stages), - exact support tuning for every possible exotic custom sampler, - verified bindings for every historical Impact Pack hook/provider variant. diff --git a/docs/releases/v1.1.0.md b/docs/releases/v1.1.0.md new file mode 100644 index 0000000..b84bbde --- /dev/null +++ b/docs/releases/v1.1.0.md @@ -0,0 +1,27 @@ +# ComfyUI Scale-Locked Residual Diffusion v1.1.0 + +## Added + +- **Krea 2 Turbo Flow-Aligned Regenerate [Experimental]**: a low-resolution planner and full-prediction composition lock for native Krea 2 Turbo, using CFG 1, Euler/simple and flow shift 1.15. +- Closed-form target handoff for unmasked flow-Euler schedules: locked high-resolution steps do not run the target denoiser. +- Patch-phase Gaussian noise coupling, including noninteger planner-to-output scale factors. +- Restart-based resolution ladder for large output canvases, with automatic area-growth limits or explicit intermediate megapixels. Intermediate stages each finish their image before the next stage re-noises and refines it. +- **Resolution Selector (Uncapped)** with the built-in aspect ratios, rounding, preview where available and explicit per-axis maximum validation. + +## Changed + +- The experimental Krea 2 node now defaults both post-release frequency-band lock strengths to zero to avoid doubled contours and softened fine detail. Existing saved workflows retain their stored values; adjust them manually if necessary. +- The generic Scale-Locked Residual samplers and their defaults remain unchanged. + +## Validation + +- 28 CPU numerical tests passed on the release candidate, including patch-noise invariants, composition-lock trajectory algebra and stage parsing. Small random-weight ComfyUI integration exercises native execution, schedule handling and both Krea 2 image latent layouts. +- Limited real-checkpoint comparisons have shown clean examples near 4 MP and for a 4096² multi-stage image. These individual observations are not a systematic quality study. + +## Limitations + +- Krea 2 Turbo regeneration above its normal native resolution is **experimental**. Artifacts, excessive VRAM use and slower wall time are possible. There is no general image-quality or speedup claim. +- The resolution ladder requires native flow-Euler sampling without masks or manifold correction. Each intermediate stage incurs additional post-release denoiser evaluations. +- Other checkpoints, solvers and custom-sampling integrations should use the generic Scale-Locked nodes. + +Install the released repository through ComfyUI Patcher or the ComfyUI registry and restart ComfyUI. No manual repository operations are necessary. diff --git a/krea2_flow.py b/krea2_flow.py new file mode 100644 index 0000000..cb178da --- /dev/null +++ b/krea2_flow.py @@ -0,0 +1,87 @@ +from __future__ import annotations + +import torch + + +def build_patch_coupled_noise( + lowres_noise: torch.Tensor, + target_shape: tuple[int, ...], + seed: int, +) -> torch.Tensor: + """Lift low-grid Gaussian noise onto the finer Krea 2 patch lattice. + + Each source 2x2 latent-patch phase is refined independently. Within every + parent cell's target group, the normalized sum equals the source noise + component. For iid Gaussian source and innovation this produces iid + standard Gaussian target noise, including non-integral scale ratios. + + This is a *noise coupling*, not an upscaler or a guarantee of visual quality. + """ + if isinstance(lowres_noise, torch.Tensor) and lowres_noise.ndim == 5: + if lowres_noise.shape[2] != 1 or len(target_shape) != 5 or target_shape[2] != 1: + raise ValueError("Krea 2 image noise requires a single temporal slice") + return build_patch_coupled_noise( + lowres_noise.squeeze(2), tuple(target_shape[:2]) + tuple(target_shape[-2:]), seed + ).unsqueeze(2) + if not isinstance(lowres_noise, torch.Tensor) or lowres_noise.ndim != 4: + raise TypeError("source noise must be BxCxHxW") + if not lowres_noise.is_floating_point() or not bool(torch.isfinite(lowres_noise).all()): + raise ValueError("source noise must contain finite floating values") + + b, c, sh, sw = map(int, lowres_noise.shape) + tb, tc, th, tw = map(int, target_shape) + if (b, c) != (tb, tc): + raise ValueError("batch/channel dimensions must match") + if min(sh, sw, th, tw) < 2 or any(v % 2 for v in (sh, sw, th, tw)): + raise ValueError("Krea 2 patch geometry requires positive even latent H/W") + if th < sh or tw < sw or (th == sh and tw == sw): + raise ValueError("target must grow at least one spatial dimension") + + g = torch.Generator(device="cpu") + g.manual_seed(int(seed) & 0xFFFFFFFFFFFFFFFF) + innovation = torch.randn( + target_shape, generator=g, device="cpu", dtype=torch.float32 + ).to(lowres_noise.device) + result = torch.empty_like(innovation) + + source_patch_h, source_patch_w = sh // 2, sw // 2 + target_patch_h, target_patch_w = th // 2, tw // 2 + owner_y = ( + torch.arange(target_patch_h, device=lowres_noise.device, dtype=torch.long) + * source_patch_h // target_patch_h + ) + owner_x = ( + torch.arange(target_patch_w, device=lowres_noise.device, dtype=torch.long) + * source_patch_w // target_patch_w + ) + + def axis_means(x: torch.Tensor, groups: int) -> tuple[torch.Tensor, torch.Tensor]: + length = x.shape[-1] + bounds = torch.div( + torch.arange(groups + 1, device=x.device) * length + groups - 1, + groups, rounding_mode="floor", + ) + prefix = torch.cat((torch.zeros_like(x[..., :1]), x.cumsum(dim=-1)), dim=-1) + sums = prefix.index_select(-1, bounds[1:]) - prefix.index_select(-1, bounds[:-1]) + counts = bounds[1:] - bounds[:-1] + return sums / counts.to(dtype=x.dtype), counts + + for phase_y in range(2): + for phase_x in range(2): + z = innovation[..., phase_y::2, phase_x::2] + x_means, x_counts = axis_means(z, source_patch_w) + y_means_t, y_counts = axis_means( + x_means.transpose(-1, -2), source_patch_h + ) + means = y_means_t.transpose(-1, -2) + group_counts = ( + (y_counts[:, None] * x_counts[None, :])[owner_y[:, None], owner_x[None, :]] + .to(dtype=innovation.dtype) + ) + parent = lowres_noise[..., phase_y::2, phase_x::2].float() + owner = (..., owner_y[:, None], owner_x[None, :]) + result[..., phase_y::2, phase_x::2] = ( + z - means[owner] + parent[owner] / group_counts.sqrt() + ) + + return result.to(dtype=lowres_noise.dtype) diff --git a/nodes.py b/nodes.py index 0bf8595..c30ca76 100644 --- a/nodes.py +++ b/nodes.py @@ -8,14 +8,17 @@ from typing import Any import math import torch import comfy.samplers +import comfy.model_base -from .slrd_core import build_nested_noise, clone_latent, init_scale_lock_state, resize_4d_tensor, resize_mask +from .slrd_core import build_nested_noise, clone_latent, full_lock_prefix_length, init_scale_lock_state, latent_target_hw_from_megapixels, parse_stage_megapixels, plan_stage_megapixels, resize_4d_tensor, resize_mask +from .resolution import DISPLAY_NAME as RESOLUTION_DISPLAY_NAME, NODE_ID as RESOLUTION_NODE_ID, ResolutionSelectorUncapped from .slrd_runtime import ( LOCK_SCHEDULE_OPTIONS, MID_SCHEDULE_OPTIONS, ScaleLockConfig, _resolve_sampler_device, apply_scale_lock_to_noise_prediction, + calculate_sigmas, build_runtime_context_from_advanced, clean_latent, compat_sampler_names, @@ -208,6 +211,166 @@ class ScaleLockedResidualKSampler: return (result.output, result.lowres_planner, result.denoised_output) +class Krea2TurboFlowAlignedRegenerate: + """Experimental Krea 2 Turbo two-pass composition-lock sampler.""" + + CATEGORY = "sampling/scale_locked" + RETURN_TYPES = ("LATENT", "LATENT", "LATENT") + RETURN_NAMES = ("output", "lowres_planner", "denoised_output") + FUNCTION = "sample" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "latent_image": ("LATENT",), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "planner_megapixels": ("FLOAT", {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.05}), + "coarse_lock_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "Low-band pull toward the planner after the composition lock releases. " + "0 lets the high-resolution pass refine freely; nonzero values can double fine contours and soften detail."}), + "mid_lock_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "Mid-band pull toward the planner after the composition lock releases, relative to the low band. " + "0 is recommended; nonzero values can double fine contours and soften detail."}), + }, + "optional": { + "steps": ("INT", {"default": 8, "min": 1, "max": 10000, + "tooltip": "Steps in each pass. Connect the same value to a native sampler for matched target-pass comparisons."}), + "composition_lock_sigma": ("FLOAT", {"default": 0.80, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Target steps starting above this noise level reuse the complete planner prediction, " + "so the high-resolution pass cannot form a second, displaced composition. " + "Lower values keep more planner structure; 1.0 disables the full lock."}), + "max_stage_upscale": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 64.0, "step": 0.25, + "tooltip": "Largest pixel-area growth per resolution stage (4.0 = 2x per side). Larger jumps are " + "split into geometrically spaced intermediate stages. 0 always jumps from the planner " + "straight to the output size. Ignored when stage_megapixels is set."}), + "stage_megapixels": ("STRING", {"default": "", "multiline": False, + "tooltip": "Optional explicit intermediate stage sizes in megapixels, ascending and between the " + "planner and the output, e.g. '4, 8'. Empty uses max_stage_upscale."}), + }, + } + + def sample( + self, + model, + positive, + negative, + latent_image, + seed, + planner_megapixels, + coarse_lock_strength, + mid_lock_strength, + steps=8, + composition_lock_sigma=0.80, + max_stage_upscale=4.0, + stage_megapixels="", + ): + composition_lock_sigma = float(composition_lock_sigma) + if not math.isfinite(composition_lock_sigma) or not 0.0 <= composition_lock_sigma <= 1.0: + raise ValueError("composition_lock_sigma must be within [0, 1]") + if not isinstance(getattr(model, "model", None), comfy.model_base.Krea2): + raise ValueError("Krea 2 Turbo Flow-Aligned Regenerate requires a native ComfyUI Krea2 MODEL") + if not isinstance(latent_image, dict) or not isinstance(latent_image.get("samples"), torch.Tensor): + raise TypeError("Krea 2 requires a LATENT dictionary with tensor samples") + # Native ComfyUI uses four-channel EmptyLatentImage placeholders for Krea 2, + # then expands them to sixteen channels during normal sampler preflight. + prepared_latent = fix_latent_channels(model, latent_image) + latent = prepared_latent["samples"] + image_layout = latent.ndim == 4 or (latent.ndim == 5 and latent.shape[2] == 1) + if not image_layout or latent.shape[1] != 16: + raise ValueError("Krea 2 requires a Bx16xHxW or Bx16x1xHxW image LATENT after native empty-latent conversion") + target_h, target_w = map(int, latent.shape[-2:]) + if target_h % 2 or target_w % 2: + raise ValueError("Krea 2 requires even latent dimensions (16-pixel output alignment)") + low_h, low_w = latent_target_hw_from_megapixels(latent, planner_megapixels) + if low_h > target_h or low_w > target_w or (low_h, low_w) == (target_h, target_w): + raise ValueError("Final LATENT must be larger than the planner canvas") + + explicit_stages = parse_stage_megapixels(stage_megapixels) + if explicit_stages: + stage_mp = explicit_stages + else: + stage_mp = plan_stage_megapixels(low_h * low_w * 64.0, target_h * target_w * 64.0, max_stage_upscale) + stage_hw = [] + previous = (low_h, low_w) + for megapixels in stage_mp: + hw = latent_target_hw_from_megapixels(latent, megapixels) + if hw == previous and not explicit_stages: + continue # Rounding to the 16-pixel grid merged two automatic stages. + if hw[0] < previous[0] or hw[1] < previous[1] or hw == previous: + raise ValueError( + "stage_megapixels must increase from the planner size, and each stage must grow on the 16-pixel grid" + ) + stage_hw.append(hw) + previous = hw + if stage_hw and (target_h < previous[0] or target_w < previous[1] or (target_h, target_w) == previous): + if explicit_stages: + raise ValueError("stage_megapixels must stay below the output LATENT size") + stage_hw.pop() + if stage_hw: + # Fail before the planner pass when no step remains after the lock. + schedule = [float(v) for v in calculate_sigmas(model, scheduler="simple", steps=steps, denoise=1.0).tolist()] + if full_lock_prefix_length(schedule, composition_lock_sigma) >= len(schedule) - 1: + raise ValueError("resolution stages need at least one step after the composition lock; raise steps or composition_lock_sigma") + logger.info( + "Krea 2 resolution stages: planner %sx%s -> %s -> output %sx%s", + low_w * 8, low_h * 8, " -> ".join(f"{w * 8}x{h * 8}" for h, w in stage_hw), target_w * 8, target_h * 8, + ) + + model_sampling = model.get_model_object("model_sampling") + shift = float(getattr(model_sampling, "shift", float("nan"))) + if not math.isfinite(shift) or abs(shift - 1.15) > 1e-5: + raise ValueError("Krea 2 Turbo mode requires the native flow sampling shift 1.15") + if max(target_h, target_w) * 8 > 2048: + logger.warning( + "Krea 2 Turbo experimental target %sx%s exceeds the official " + "2048-pixel-axis inference range; visual fidelity and VRAM use are unvalidated", + target_w * 8, target_h * 8, + ) + + config = ScaleLockConfig( + lock_strength=coarse_lock_strength, + lock_strength_start=1.0, + lock_strength_end=0.20, + cutoff=0.33, + mid_cutoff=0.60, + mid_strength=mid_lock_strength, + schedule="hold_then_drop", + schedule_power=2.0, + schedule_hold=0.25, + mid_strength_start=1.0, + mid_strength_end=0.0, + mid_schedule="linked", + mid_schedule_power=2.0, + mid_schedule_hold=0.0, + full_lock_sigma=composition_lock_sigma, + ) + result = run_scale_locked_ksampler( + model=model, + positive=positive, + negative=negative, + latent_image=prepared_latent, + seed=seed, + steps=steps, + cfg=1.0, + sampler_name="euler", + scheduler="simple", + denoise=1.0, + target_megapixels=planner_megapixels, + nested_noise_strength=1.0, + add_noise=True, + pin_anchors=True, + sampler_guard="error", + config=config, + noise_mode="patch_coupled", + stage_hw=tuple(stage_hw), + ) + return (result.output, result.lowres_planner, result.denoised_output) + + class ScaleLockedRuntimeContextBuilder: CATEGORY = "sampling/scale_locked" RETURN_TYPES = ("SLRD_RUNTIME", "NOISE", "LATENT") @@ -1646,6 +1809,8 @@ def _snapshot_impact_latent(latent: dict[str, Any]) -> dict[str, Any]: NODE_CLASS_MAPPINGS = { "ScaleLockedResidualKSampler": ScaleLockedResidualKSampler, + "Krea2TurboFlowAlignedRegenerate": Krea2TurboFlowAlignedRegenerate, + RESOLUTION_NODE_ID: ResolutionSelectorUncapped, "ScaleLockedRuntimeContextBuilder": ScaleLockedRuntimeContextBuilder, "ScaleLockedCFGGuider": ScaleLockedCFGGuiderNode, "ScaleLockedResidualSamplerCustomAdvanced": ScaleLockedResidualSamplerCustomAdvanced, @@ -1655,6 +1820,8 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "ScaleLockedResidualKSampler": "Scale-Locked Residual KSampler", + "Krea2TurboFlowAlignedRegenerate": "Krea 2 Turbo Flow-Aligned Regenerate [Experimental]", + RESOLUTION_NODE_ID: RESOLUTION_DISPLAY_NAME, "ScaleLockedRuntimeContextBuilder": "Scale-Locked Runtime Context", "ScaleLockedCFGGuider": "Scale-Locked CFG Guider", "ScaleLockedResidualSamplerCustomAdvanced": "Scale-Locked Residual SamplerCustomAdvanced", diff --git a/pyproject.toml b/pyproject.toml index 54b584d..f651776 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "scale-locked-residual-diffusion" description = "A ComfyUI custom node pack implementing Scale-Locked Residual Diffusion for high-resolution composition and anatomy stability." -version = "1.0.32" +version = "1.1.0" license = { file = "LICENSE" } [project.urls] diff --git a/resolution.py b/resolution.py new file mode 100644 index 0000000..2981c79 --- /dev/null +++ b/resolution.py @@ -0,0 +1,127 @@ +"""Resolution Selector without the core 16-megapixel input cap. + +Same aspect ratios, megapixel convention (1 MP = 1024 x 1024 pixels), rounding +and live preview as ComfyUI's built-in Resolution Selector. Only the megapixel +limit differs: the outputs may use the full ComfyUI canvas range. +""" +from __future__ import annotations + +import math + +from nodes import MAX_RESOLUTION + +ASPECT_RATIOS = { + "1:1 (Square)": (1, 1), + "2:3 (Portrait Photo)": (2, 3), + "3:2 (Photo)": (3, 2), + "3:4 (Portrait Standard)": (3, 4), + "4:3 (Standard)": (4, 3), + "9:16 (Portrait Widescreen)": (9, 16), + "16:9 (Widescreen)": (16, 9), + "21:9 (Ultrawide)": (21, 9), +} +# Largest square canvas ComfyUI latents accept, in 1024 x 1024 pixel megapixels. +MAX_MEGAPIXELS = float((MAX_RESOLUTION // 1024) ** 2) + +NODE_ID = "ScaleLockedResolutionSelector" +DISPLAY_NAME = "Resolution Selector (Uncapped)" +DESCRIPTION = ( + "Calculate width and height from an aspect ratio and a megapixel target, like the built-in " + "Resolution Selector but without its 16-megapixel limit. 1 MP = 1024x1024 pixels." +) + + +def compute_resolution(aspect_ratio: str, megapixels: float, multiple: int) -> tuple[int, int]: + if aspect_ratio not in ASPECT_RATIOS: + raise ValueError(f"unknown aspect ratio {aspect_ratio!r}") + megapixels = float(megapixels) + multiple = int(multiple) + if not math.isfinite(megapixels) or megapixels <= 0.0: + raise ValueError("megapixels must be positive") + if multiple < 1: + raise ValueError("multiple must be positive") + w_ratio, h_ratio = ASPECT_RATIOS[aspect_ratio] + scale = math.sqrt(megapixels * 1024 * 1024 / (w_ratio * h_ratio)) + width = max(multiple, round(w_ratio * scale / multiple) * multiple) + height = max(multiple, round(h_ratio * scale / multiple) * multiple) + if max(width, height) > MAX_RESOLUTION: + raise ValueError( + f"{aspect_ratio} at {megapixels:g} MP gives {width}x{height}, beyond ComfyUI's " + f"{MAX_RESOLUTION}-pixel canvas limit per side" + ) + return width, height + + +_TOOLTIPS = { + "aspect_ratio": "The aspect ratio for the output dimensions.", + "megapixels": "Target total megapixels. 1.0 MP = 1024x1024 for square. No 16 MP cap.", + "multiple": "Nearest multiple of the result to set the selected resolution to.", + "width": "Calculated width in pixels multiplied by the selected multiple.", + "height": "Calculated height in pixels multiplied by the selected multiple.", +} + + +def _build_v3_node(): + from comfy_api.latest import io + + if not hasattr(io, "ResolutionPreview"): + return None + + class ResolutionSelectorUncapped(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id=NODE_ID, + display_name=DISPLAY_NAME, + category="utilities", + description=DESCRIPTION, + inputs=[ + io.Combo.Input("aspect_ratio", options=list(ASPECT_RATIOS), + default="1:1 (Square)", tooltip=_TOOLTIPS["aspect_ratio"]), + io.Float.Input("megapixels", default=1.0, min=0.1, max=MAX_MEGAPIXELS, step=0.1, + tooltip=_TOOLTIPS["megapixels"]), + io.ResolutionPreview.Input("preview", tooltip="Live preview of the calculated output resolution."), + io.Int.Input(id="multiple", default=8, min=8, max=128, step=4, + tooltip=_TOOLTIPS["multiple"], advanced=True), + ], + outputs=[ + io.Int.Output("width", tooltip=_TOOLTIPS["width"]), + io.Int.Output("height", tooltip=_TOOLTIPS["height"]), + ], + ) + + @classmethod + def execute(cls, aspect_ratio: str, megapixels: float, multiple: int, preview=None) -> io.NodeOutput: + return io.NodeOutput(*compute_resolution(aspect_ratio, megapixels, multiple)) + + return ResolutionSelectorUncapped + + +class ResolutionSelectorUncappedClassic: + """Fallback for ComfyUI builds without the resolution preview widget.""" + + CATEGORY = "utilities" + DESCRIPTION = DESCRIPTION + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("width", "height") + FUNCTION = "execute" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "aspect_ratio": (list(ASPECT_RATIOS), {"default": "1:1 (Square)", "tooltip": _TOOLTIPS["aspect_ratio"]}), + "megapixels": ("FLOAT", {"default": 1.0, "min": 0.1, "max": MAX_MEGAPIXELS, "step": 0.1, + "tooltip": _TOOLTIPS["megapixels"]}), + "multiple": ("INT", {"default": 8, "min": 8, "max": 128, "step": 4, "tooltip": _TOOLTIPS["multiple"]}), + } + } + + def execute(self, aspect_ratio, megapixels, multiple): + return compute_resolution(aspect_ratio, megapixels, multiple) + + +try: + ResolutionSelectorUncapped = _build_v3_node() or ResolutionSelectorUncappedClassic +except ImportError: + ResolutionSelectorUncapped = ResolutionSelectorUncappedClassic diff --git a/slrd_core.py b/slrd_core.py index 6480ac7..67b6e81 100644 --- a/slrd_core.py +++ b/slrd_core.py @@ -50,6 +50,10 @@ def _interp_mode(src_hw: tuple[int, int], dst_hw: tuple[int, int]) -> str: def resize_4d_tensor(x: torch.Tensor, size_hw: tuple[int, int]) -> torch.Tensor: """Resize BCHW tensor with sane defaults for latent tensors.""" + if x.ndim == 5 and x.shape[2] == 1: + # Native image models may retain a singleton temporal axis. Resize only + # their spatial plane while preserving the native latent layout. + return resize_4d_tensor(x.squeeze(2), size_hw).unsqueeze(2) if tuple(x.shape[-2:]) == tuple(size_hw): return x mode = _interp_mode(tuple(x.shape[-2:]), tuple(size_hw)) @@ -97,7 +101,9 @@ def latent_target_hw_from_megapixels( Compute a target latent spatial size that preserves aspect ratio while targeting a pixel-space megapixel count. Latent tensors are expected to be 8x downsampled in H/W for standard image models used by ComfyUI. """ - _, _, h_lat, w_lat = latent_bchw.shape + if latent_bchw.ndim != 4 and not (latent_bchw.ndim == 5 and latent_bchw.shape[2] == 1): + raise ValueError("Expected an image latent BCHW or BC1HW") + h_lat, w_lat = latent_bchw.shape[-2:] h_px = h_lat * 8 w_px = w_lat * 8 aspect = w_px / max(1.0, float(h_px)) @@ -722,7 +728,13 @@ def init_scale_lock_state( manifold_energy_gain_cap: float = 1.75, manifold_max_shift_px: float = 3.0, manifold_spatial_mask: Optional[torch.Tensor] = None, + full_lock_sigma: Optional[float] = None, ) -> None: + if full_lock_sigma is not None: + full_lock_sigma = float(full_lock_sigma) + if not math.isfinite(full_lock_sigma) or full_lock_sigma < 0.0: + raise ValueError("full_lock_sigma must be a finite non-negative noise level") + guider._slrd_full_lock_sigma = full_lock_sigma guider._slrd_model = model guider._slrd_anchors_x0_cpu = list(anchors_x0_cpu) guider._slrd_planner_sigmas = [float(x) for x in planner_sigmas] if planner_sigmas is not None else None @@ -805,6 +817,84 @@ def resolve_scale_lock_step_index(guider, timestep: torch.Tensor | float | int) return min(step_index, len(guider._slrd_anchors_x0_cpu) - 1) +def scale_lock_full_lock_active(guider, timestep: torch.Tensor | float | int) -> bool: + """Whether this target step reuses the complete planner prediction. + + While active, the target state stays the planner state with the target + noise field. Its own composition cannot enter any frequency band. + """ + threshold = getattr(guider, "_slrd_full_lock_sigma", None) + return threshold is not None and _sigma_scalar(timestep) > threshold + + +def full_lock_prefix_length(sigmas: Sequence[float], threshold: Optional[float]) -> int: + """Number of leading sampler steps whose input sigma exceeds the threshold.""" + if threshold is None: + return 0 + count = 0 + for sigma in list(sigmas)[:-1]: + if float(sigma) <= threshold: + break + count += 1 + return count + + +def parse_stage_megapixels(text: str) -> list[float]: + """Parse a comma- or space-separated list of intermediate stage sizes.""" + values = [] + for part in str(text or "").replace(";", ",").replace(" ", ",").split(","): + part = part.strip() + if not part: + continue + try: + value = float(part) + except ValueError as exc: + raise ValueError(f"stage_megapixels entry {part!r} is not a number") from exc + if not math.isfinite(value) or value <= 0.0: + raise ValueError("stage_megapixels entries must be positive megapixel values") + values.append(value) + return values + + +def plan_stage_megapixels(planner_area_px: float, target_area_px: float, max_stage_upscale: float) -> list[float]: + """Geometric intermediate stage sizes so no stage grows the area by more than the limit.""" + max_stage_upscale = float(max_stage_upscale) + if max_stage_upscale <= 0.0: + return [] + if not math.isfinite(max_stage_upscale) or max_stage_upscale <= 1.0: + raise ValueError("max_stage_upscale must be 0 (single stage) or greater than 1") + ratio = float(target_area_px) / float(planner_area_px) + if ratio <= 1.0: + return [] + stages = max(1, math.ceil(math.log(ratio) / math.log(max_stage_upscale) - 1e-9)) + return [planner_area_px * ratio ** (i / stages) / 1_000_000.0 for i in range(1, stages)] + + +def flow_euler_locked_clean_state( + initial_clean: torch.Tensor, + predictions: Iterable[torch.Tensor], + sigmas: Sequence[float], +) -> torch.Tensor: + """Clean component after flow-matching Euler steps with given predictions. + + Flow Euler keeps ``x_k = sigma_k * noise + (1 - sigma_k) * C_k``. Each + step replaces ``C`` by a convex combination of ``C`` and the prediction, + so the recurrence needs no division by ``1 - sigma``. ``sigmas`` holds one + more entry than ``predictions``; the result is ``C`` at the last sigma. + """ + sigmas = [float(v) for v in sigmas] + predictions = list(predictions) + if len(sigmas) != len(predictions) + 1: + raise ValueError("flow Euler state needs one more sigma than predictions") + clean = initial_clean.to(dtype=torch.float32) + for sigma, sigma_next, prediction in zip(sigmas[:-1], sigmas[1:], predictions): + if not (sigma > sigma_next >= 0.0) or sigma > 1.0: + raise ValueError("flow Euler sigmas must decrease within [0, 1]") + keep = 0.0 if sigma >= 1.0 else sigma_next * (1.0 - sigma) / (sigma * (1.0 - sigma_next)) + clean = torch.lerp(prediction.to(device=clean.device, dtype=torch.float32), clean, keep) + return clean + + def scale_lock_strength_for_step(guider, idx: int) -> float: return scale_lock_strengths_for_step(guider, idx)[0] diff --git a/slrd_runtime.py b/slrd_runtime.py index 838a876..953be12 100644 --- a/slrd_runtime.py +++ b/slrd_runtime.py @@ -9,11 +9,14 @@ from typing import Any, Optional import torch import comfy.model_management +import comfy.model_sampling import comfy.sample import comfy.samplers import comfy.utils import latent_preview +from .krea2_flow import build_patch_coupled_noise + from .slrd_core import ( TrajectoryRecorder, build_nested_noise, @@ -22,10 +25,14 @@ from .slrd_core import ( latent_manifold_compand, latent_target_hw_from_megapixels, resolve_scale_lock_step_index, + resize_4d_tensor, resize_latent_dict, resize_mask, residual_lock_multiband, scale_lock_anchor_for, + scale_lock_full_lock_active, + full_lock_prefix_length, + flow_euler_locked_clean_state, scale_lock_manifold_mask_for, scale_lock_manifold_strength_for_step, scale_lock_mask_for, @@ -104,6 +111,9 @@ class ScaleLockConfig: manifold_energy_gain_cap: float = 1.75 manifold_max_shift_px: float = 3.0 manifold_spatial_mask: Optional[torch.Tensor] = None + # Target steps whose input sigma exceeds this value use the planner + # prediction in every frequency band. None disables the full lock. + full_lock_sigma: Optional[float] = None @dataclass @@ -321,6 +331,13 @@ def apply_scale_lock_to_noise_prediction(guider, base_noise: torch.Tensor, x, ti idx = resolve_scale_lock_step_index(guider, timestep) anchor = scale_lock_anchor_for(guider, idx, base_noise) + if scale_lock_full_lock_active(guider, timestep): + mask = scale_lock_mask_for(guider, base_noise) + if mask is None: + # The anchor may alias stored planner state; never hand it out. + return anchor.clone() + return base_noise + mask * (anchor - base_noise) + corrected = base_noise low_strength, mid_strength = scale_lock_strengths_for_step(guider, idx) if low_strength > 0.0 or mid_strength > 0.0: @@ -412,6 +429,7 @@ def apply_scale_lock_to_guider(guider, runtime: ScaleLockedRuntimeContext, confi manifold_energy_gain_cap=config.manifold_energy_gain_cap, manifold_max_shift_px=config.manifold_max_shift_px, manifold_spatial_mask=manifold_spatial_mask, + full_lock_sigma=config.full_lock_sigma, ) original_predict_noise = getattr(guider, "_slrd_original_predict_noise", guider.predict_noise) @@ -433,6 +451,17 @@ def restore_original_predict_noise(guider) -> None: guider.predict_noise = original_predict_noise +def _planner_callback(model, recorder: TrajectoryRecorder, sigmas: torch.Tensor): + """Record planner anchors and report progress like the other passes.""" + preview = make_preview_callback(model, max(0, len(sigmas) - 1), {}) + + def callback(step, x0, x, total_steps): + recorder.callback(step, x0, x, total_steps) + preview(step, x0, x, total_steps) + + return callback + + def _run_lowres_planner_advanced( *, guider, @@ -475,8 +504,8 @@ def _run_lowres_planner_advanced( sampler, sigmas, denoise_mask=noise_mask, - callback=recorder.callback, - disable_pbar=True, + callback=_planner_callback(model, recorder, sigmas), + disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED, seed=noise_seed(noise), ) samples = samples.to(comfy.model_management.intermediate_device()) @@ -546,8 +575,8 @@ def _run_lowres_planner( sampler_obj, sigmas, denoise_mask=noise_mask, - callback=recorder.callback, - disable_pbar=True, + callback=_planner_callback(model, recorder, sigmas), + disable_pbar=not comfy.utils.PROGRESS_BAR_ENABLED, seed=seed, ) samples = samples.to(comfy.model_management.intermediate_device()) @@ -561,7 +590,13 @@ def _run_lowres_planner( return out, recorder.x0_steps, planner_sigmas, planner_noise -def _build_highres_noise(highres_latent: dict, lowres_noise: torch.Tensor, seed: int, hf_strength: float): +def _build_highres_noise( + highres_latent: dict, + lowres_noise: torch.Tensor, + seed: int, + hf_strength: float, + mode: str = "legacy", +): highres_samples = highres_latent["samples"] target_device = highres_samples.device target_dtype = highres_samples.dtype @@ -573,6 +608,15 @@ def _build_highres_noise(highres_latent: dict, lowres_noise: torch.Tensor, seed: device=target_device, ) + if mode == "patch_coupled": + return build_patch_coupled_noise( + lowres_noise=lowres_noise, + target_shape=tuple(highres_samples.shape), + seed=seed, + ).to(device=target_device, dtype=target_dtype, non_blocking=True) + if mode != "legacy": + raise ValueError(f"unsupported nested-noise mode: {mode}") + return build_nested_noise( lowres_noise=lowres_noise, target_shape=tuple(highres_samples.shape), @@ -654,6 +698,44 @@ def build_runtime_context_from_advanced( ) +def _full_lock_handoff(runtime: ScaleLockedRuntimeContext, config: ScaleLockConfig): + """Target state at the composition-lock release, without model calls. + + Locked steps discard the target prediction, so evaluating the model there + is wasted work. For flow-matching Euler the state after those steps is + ``sigma * noise + (1 - sigma) * clean`` with ``clean`` a fixed convex + combination of the resized planner predictions. Returns + ``(locked_steps, clean)`` in the model's internal latent space, or None + when no step is locked or the masked fallback path must run instead. + """ + sigma_values = [float(v) for v in runtime.sigmas.detach().cpu().tolist()] + locked_steps = full_lock_prefix_length(sigma_values, config.full_lock_sigma) + if locked_steps == 0: + return None + if config.spatial_mask is not None or isinstance(runtime.highres_latent.get("noise_mask"), torch.Tensor): + # Masked locks blend with the target prediction; it must be evaluated. + return None + if config.manifold_enabled: + return None + if len(runtime.anchors_x0) < locked_steps: + raise RuntimeError("planner recorded fewer anchors than locked target steps") + + target = runtime.highres_latent["samples"] + initial = target + if torch.count_nonzero(initial) > 0: + # Mirrors CFGGuider.inner_sample, which leaves empty latents unshifted. + initial = runtime.model.model.process_latent_in(initial) + predictions = ( + resize_4d_tensor( + runtime.anchors_x0[i].to(device=target.device, dtype=torch.float32, non_blocking=True), + tuple(target.shape[-2:]), + ) + for i in range(locked_steps) + ) + clean = flow_euler_locked_clean_state(initial, predictions, sigma_values[: locked_steps + 1]) + return locked_steps, clean + + def sample_with_runtime( *, guider, @@ -661,24 +743,44 @@ def sample_with_runtime( runtime: ScaleLockedRuntimeContext, config: ScaleLockConfig, restore_after: bool = True, + analytic_full_lock: bool = False, ) -> ScaleLockedSampleResult: if runtime.sigmas.numel() == 0: out = clean_latent(runtime.highres_latent) return ScaleLockedSampleResult(output=out, lowres_planner=clean_latent(runtime.lowres_out), denoised_output=out) + sigmas = runtime.sigmas + latent_samples = runtime.highres_latent["samples"] + handoff = _full_lock_handoff(runtime, config) if analytic_full_lock else None + if handoff is not None: + locked_steps, clean = handoff + if locked_steps == len(sigmas) - 1: + # Every step is locked: the result is the final clean state. + out = clean_latent(runtime.highres_latent) + out["samples"] = runtime.model.model.process_latent_out(clean).to( + comfy.model_management.intermediate_device() + ) + return ScaleLockedSampleResult( + output=out, lowres_planner=clean_latent(runtime.lowres_out), denoised_output=out + ) + # The sampler re-applies process_latent_in and the flow noise scaling, + # reproducing sigma * noise + (1 - sigma) * clean at the release step. + sigmas = sigmas[locked_steps:] + latent_samples = runtime.model.model.process_latent_out(clean).to(dtype=latent_samples.dtype) + apply_scale_lock_to_guider(guider, runtime, config) x0_output = {} - callback = make_preview_callback(runtime.model, len(runtime.sigmas) - 1, x0_output) + callback = make_preview_callback(runtime.model, len(sigmas) - 1, x0_output) noise_mask = runtime.highres_latent.get("noise_mask", None) disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED try: samples = guider.sample( runtime.highres_noise, - runtime.highres_latent["samples"], + latent_samples, sampler, - runtime.sigmas, + sigmas, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, @@ -709,6 +811,153 @@ def sample_with_runtime( ) +def _stage_seed(detail_seed: int, stage_index: int) -> int: + return (int(detail_seed) ^ ((stage_index + 1) * 0xD1B54A32D192ED03)) & 0xFFFFFFFFFFFFFFFF + + +def _run_resolution_ladder( + *, + model, + positive, + negative, + cfg, + sampler_name: str, + sigmas: torch.Tensor, + highres_latent: dict, + lowres_latent: dict, + lowres_out: dict, + anchors_x0: list, + planner_sigmas: list, + lowres_noise: torch.Tensor, + detail_seed: int, + nested_noise_strength: float, + noise_mode: str, + config: ScaleLockConfig, + stage_hw: tuple, +) -> ScaleLockedSampleResult: + """Refine through increasing resolutions, each stage finishing its image. + + The first stage starts at the composition-lock release in closed form + from the planner predictions, exactly like the single-stage path. Every + later stage re-noises the finished image of the stage below, resized to + its own size, to the same release noise level with its own patch-coupled + noise, and runs the remaining steps. Each stage therefore repeats the + planner-to-2x handoff at a bounded scale ratio, starting from a sharp + image rather than from mid-noise predictions, whose posterior-mean blur + a late, low-noise start cannot remove. + """ + if config.full_lock_sigma is None: + raise ValueError("resolution stages require the composition lock") + if _normalize_sampler_name(sampler_name) != "euler" or not isinstance( + _model_sampling_obj(model), comfy.model_sampling.CONST + ): + raise ValueError("resolution stages require Euler with flow model sampling") + if ( + isinstance(highres_latent.get("noise_mask"), torch.Tensor) + or config.spatial_mask is not None + or config.manifold_enabled + ): + raise ValueError("resolution stages do not support latent noise masks, spatial lock masks or manifold correction") + + target = highres_latent["samples"] + sigma_values = [float(v) for v in sigmas.detach().cpu().tolist()] + total_steps = len(sigma_values) - 1 + release = full_lock_prefix_length(sigma_values, config.full_lock_sigma) + if release >= total_steps: + raise ValueError("resolution stages need at least one step after the composition lock") + sizes = [tuple(map(int, hw)) for hw in stage_hw] + [tuple(map(int, target.shape[-2:]))] + previous = tuple(map(int, lowres_latent["samples"].shape[-2:])) + for hw in sizes: + if hw[0] < previous[0] or hw[1] < previous[1] or hw == previous or hw[0] % 2 or hw[1] % 2: + raise ValueError("resolution stages must grow monotonically in both dimensions on the 16-pixel grid") + previous = hw + + stage_noise = lowres_noise + sampler = comfy.samplers.sampler_object(sampler_name) + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + x0_output: dict = {} + finished = None + + for index, hw in enumerate(sizes): + final = index == len(sizes) - 1 + stage_latent = highres_latent if final else resize_latent_dict(highres_latent, hw) + stage_samples = stage_latent["samples"] + seed = detail_seed if final else _stage_seed(detail_seed, index) + stage_noise = _build_highres_noise( + highres_latent=stage_latent, + lowres_noise=stage_noise, + seed=seed, + hf_strength=nested_noise_strength, + mode=noise_mode, + ) + + if finished is None: + start_latent = stage_samples + if release > 0: + initial = stage_samples + if torch.count_nonzero(initial) > 0: + initial = model.model.process_latent_in(initial) + predictions = ( + resize_4d_tensor( + anchors_x0[k].to(device=stage_samples.device, dtype=torch.float32, non_blocking=True), hw + ) + for k in range(release) + ) + clean = flow_euler_locked_clean_state(initial, predictions, sigma_values[: release + 1]) + start_latent = model.model.process_latent_out(clean).to(dtype=stage_samples.dtype) + else: + # Output-space latent: CFGGuider.inner_sample applies process_latent_in + # and the flow noise scaling re-noises it to the release sigma. + start_latent = resize_4d_tensor( + finished.to(device=stage_samples.device, dtype=torch.float32), hw + ).to(dtype=stage_samples.dtype) + + stage_runtime = ScaleLockedRuntimeContext( + model=model, + highres_latent=stage_latent, + lowres_latent=lowres_latent, + lowres_out=lowres_out, + sigmas=sigmas, + anchors_x0=anchors_x0, + planner_sigmas=planner_sigmas, + highres_noise=stage_noise, + noise_seed=seed, + ) + guider = create_cfg_guider(model, positive, negative, cfg) + apply_scale_lock_to_guider(guider, stage_runtime, config) + stage_x0: dict = {} + callback = make_preview_callback(model, total_steps - release, stage_x0) + try: + finished = guider.sample( + stage_noise, + start_latent, + sampler, + sigmas[release:], + denoise_mask=None, + callback=callback, + disable_pbar=disable_pbar, + seed=seed, + ) + finally: + restore_original_predict_noise(guider) + if final: + x0_output = stage_x0 + + samples = finished.to(comfy.model_management.intermediate_device()) + out = clean_latent(highres_latent) + out["samples"] = samples + if "x0" in x0_output: + try: + x0_out = model.model.process_latent_out(x0_output["x0"].cpu()) + except Exception: + x0_out = x0_output["x0"].detach().cpu() + denoised = clean_latent(highres_latent) + denoised["samples"] = x0_out + else: + denoised = out + return ScaleLockedSampleResult(output=out, lowres_planner=clean_latent(lowres_out), denoised_output=denoised) + + def run_scale_locked_ksampler( *, model, @@ -727,6 +976,8 @@ def run_scale_locked_ksampler( pin_anchors, sampler_guard, config: ScaleLockConfig, + noise_mode: str = "legacy", + stage_hw: tuple = (), ) -> ScaleLockedSampleResult: if steps < 1: raise ValueError("steps must be >= 1") @@ -762,6 +1013,16 @@ def run_scale_locked_ksampler( if len(anchors_x0) == 0: raise RuntimeError("ScaleLockedResidualKSampler: planner pass did not record any x0 anchors.") + if stage_hw: + return _run_resolution_ladder( + model=model, positive=positive, negative=negative, cfg=cfg, + sampler_name=sampler_name, sigmas=sigmas, highres_latent=highres_latent, + lowres_latent=lowres_latent, lowres_out=lowres_out, anchors_x0=anchors_x0, + planner_sigmas=planner_sigmas, lowres_noise=lowres_noise, detail_seed=detail_seed, + nested_noise_strength=nested_noise_strength, noise_mode=noise_mode, + config=config, stage_hw=tuple(stage_hw), + ) + runtime = ScaleLockedRuntimeContext( model=model, highres_latent=highres_latent, @@ -775,10 +1036,20 @@ def run_scale_locked_ksampler( lowres_noise=lowres_noise, seed=detail_seed, hf_strength=nested_noise_strength, + mode=noise_mode, ), noise_seed=detail_seed, ) guider = create_cfg_guider(model, positive, negative, cfg) sampler = comfy.samplers.sampler_object(sampler_name) - return sample_with_runtime(guider=guider, sampler=sampler, runtime=runtime, config=config) + # The closed-form lock handoff is exact only for the flow-matching Euler update. + analytic_full_lock = ( + config.full_lock_sigma is not None + and _normalize_sampler_name(sampler_name) == "euler" + and isinstance(_model_sampling_obj(model), comfy.model_sampling.CONST) + ) + return sample_with_runtime( + guider=guider, sampler=sampler, runtime=runtime, config=config, + analytic_full_lock=analytic_full_lock, + ) diff --git a/tests/smoke_krea2_native.py b/tests/smoke_krea2_native.py new file mode 100644 index 0000000..0e38d09 --- /dev/null +++ b/tests/smoke_krea2_native.py @@ -0,0 +1,154 @@ +"""CPU integration with real ComfyUI and a tiny random native Krea 2 model. + +This checks execution/layout only. It does not establish real-checkpoint image +quality or GPU speed. Run with ComfyUI's Python dependencies installed. +""" +import argparse +import importlib.util +import json +from pathlib import Path +import sys + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--comfy-root", type=Path, required=True) + parser.add_argument("--spectrum-root", type=Path) + args = parser.parse_args() + root = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(args.comfy_root.resolve())) + sys.argv = [sys.argv[0], "--cpu", "--use-pytorch-cross-attention", "--preview-method", "none"] + import comfy.options + comfy.options.enable_args_parsing() + import torch + import comfy.supported_models + import comfy.model_patcher + import comfy.model_base + torch.set_num_threads(1) + torch.manual_seed(73) + config = comfy.supported_models.Krea2(dict( + image_model="krea2", features=32, tdim=8, txtdim=16, heads=2, kvheads=1, + multiplier=1, layers=2, patch=2, channels=16, txtlayers=2, + txtheads=2, txtkvheads=1, + )) + config.set_inference_dtype(torch.float32, None) + base = comfy.model_base.Krea2(config, device=torch.device("cpu")) + with torch.no_grad(): + for param in base.diffusion_model.parameters(): + param.copy_(torch.randn_like(param) * .02) + model = comfy.model_patcher.ModelPatcher(base, torch.device("cpu"), torch.device("cpu")) + spec = importlib.util.spec_from_file_location("slrd_native_smoke", root / "__init__.py", submodule_search_locations=[str(root)]) + package = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = package + spec.loader.exec_module(package) + node = package.NODE_CLASS_MAPPINGS["Krea2TurboFlowAlignedRegenerate"]() + positive = [[torch.randn(1, 3, 32), {}]] + negative = [[torch.zeros(1, 3, 32), {}]] + cases = [("native", model)] + reports_context = None + if args.spectrum_root: + sys.path.insert(0, str(args.spectrum_root.resolve())) + from comfyui_spectrum_krea2.nodes import SpectrumApplyKrea2 + from comfyui_spectrum_krea2.runtime import REPORTS + reports_context = REPORTS + cases.append(("spectrum", SpectrumApplyKrea2().apply(model, preset="turbo")[0])) + runtime = sys.modules[spec.name + ".slrd_runtime"] + calls = [] + locked_by_steps = {8: 4, 12: 6} + + def run(selected, steps, analytic=True, count_calls=False): + reports = [] + token = reports_context.set(reports) if reports_context is not None else None + original = runtime._full_lock_handoff + if not analytic: + runtime._full_lock_handoff = lambda *args, **kwargs: None + calls.clear() + # Spectrum declines hooked models, so only count calls without it. + hook = (base.diffusion_model.register_forward_pre_hook( + lambda module, args: calls.append(tuple(args[0].shape[-2:]))) if count_calls else None) + try: + with torch.inference_mode(): + # Small latent canvases keep genuine CPU model execution practical. + result = node.sample(selected, positive, negative, + {"samples": torch.zeros(1, 4, 16, 20)}, 42, .01, 0.0, 0.0, steps=steps) + finally: + if hook is not None: + hook.remove() + runtime._full_lock_handoff = original + if reports_context is not None: + reports_context.reset(token) + target_calls = sum(shape == (16, 20) for shape in calls) + return result, reports, target_calls + + for label, selected, steps in [(label,selected,steps) for label,selected in cases for steps in (8,12)]: + locked = locked_by_steps[steps] + result, reports, target_calls = run(selected, steps, count_calls=label == "native") + if label == "native": + # Locked target steps are skipped, not evaluated and discarded. + assert target_calls == steps - locked, (target_calls, steps, locked) + evaluated, _, evaluated_calls = run(selected, steps, analytic=False, count_calls=True) + assert evaluated_calls == steps + for fast, slow in zip(result, evaluated): + difference = float((fast["samples"] - slow["samples"]).abs().max()) + scale = float(slow["samples"].abs().max()) + assert difference <= 1e-4 * max(1.0, scale), (difference, scale) + if label == "spectrum": + assert len(reports) == 2 + assert reports[0]["outer_steps"] == steps + assert reports[1]["outer_steps"] == steps - locked + assert reports[1]["denoiser_calls"] == steps - locked + assert all(r["fallback"] is None for r in reports) + assert reports[0]["forecast"] > 0 + shapes = [list(x["samples"].shape) for x in result] + assert shapes[0] == shapes[2] == [1, 16, 1, 16, 20] + assert all(torch.isfinite(x["samples"]).all() for x in result) + assert shapes[1][2] == 1 + assert not base.diffusion_model._forward_pre_hooks + assert not base.diffusion_model.last._forward_pre_hooks + print(json.dumps({"case": label, "steps_per_pass": steps, "locked_target_steps": locked, + "target_model_calls": target_calls, "output_shapes": shapes, "reports": reports})) + + # Resolution ladder: planner 12x14 -> stage 16x20 -> output 24x28 latent cells. + # Each stage finishes its image and runs every step after the lock. + for label, selected, steps in [(label,selected,steps) for label,selected in cases for steps in (8,12)]: + for stage_options in ({"stage_megapixels": "0.02"}, {"max_stage_upscale": 2.0}): + reports = [] + token = reports_context.set(reports) if reports_context is not None else None + calls.clear() + hook = (base.diffusion_model.register_forward_pre_hook( + lambda module, args: calls.append(tuple(args[0].shape[-2:]))) if label == "native" else None) + try: + with torch.inference_mode(): + result = node.sample(selected, positive, negative, {"samples": torch.zeros(1, 4, 24, 28)}, + 42, .01, 0.0, 0.0, steps=steps, **stage_options) + finally: + if hook is not None: + hook.remove() + if reports_context is not None: + reports_context.reset(token) + free = steps - locked_by_steps[steps] + if label == "native": + counts = {shape: calls.count(shape) for shape in set(calls)} + assert counts == {(12, 14): steps, (16, 20): free, (24, 28): free}, counts + else: + assert [r["outer_steps"] for r in reports] == [steps, free, free], reports + assert all(r["fallback"] is None for r in reports) + shapes = [list(x["samples"].shape) for x in result] + assert shapes[0] == shapes[2] == [1, 16, 1, 24, 28] + assert shapes[1] == [1, 16, 1, 12, 14] + assert all(torch.isfinite(x["samples"]).all() for x in result) + assert not base.diffusion_model._forward_pre_hooks + print(json.dumps({"case": label + "_ladder", "options": stage_options, "steps_per_pass": steps, + "stage_steps": free, "output_shapes": shapes, "reports": reports})) + + try: + node.sample(model, positive, negative, {"samples": torch.zeros(1, 4, 24, 28)}, + 42, .01, 0.0, 0.0, steps=2, composition_lock_sigma=0.0, stage_megapixels="0.02") + except ValueError as exc: + assert "at least one step" in str(exc) + else: + raise AssertionError("too few steps for the stages must be rejected before sampling") + + +if __name__ == "__main__": + main() diff --git a/tests/smoke_resolution_selector.py b/tests/smoke_resolution_selector.py new file mode 100644 index 0000000..c4e8c0d --- /dev/null +++ b/tests/smoke_resolution_selector.py @@ -0,0 +1,71 @@ +"""Compare the uncapped Resolution Selector with ComfyUI's built-in node. + +Run with ComfyUI's Python dependencies installed: + python tests/smoke_resolution_selector.py --comfy-root /path/to/ComfyUI +""" +import argparse +import importlib.util +from pathlib import Path +import sys + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--comfy-root", type=Path, required=True) + args = parser.parse_args() + root = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(args.comfy_root.resolve())) + sys.argv = [sys.argv[0], "--cpu"] + import comfy.options + comfy.options.enable_args_parsing() + from comfy_extras.nodes_resolution import ResolutionSelector, AspectRatio + from comfy_api.internal import _ComfyNodeInternal + + spec = importlib.util.spec_from_file_location("slrd_resolution_smoke", root / "__init__.py", + submodule_search_locations=[str(root)]) + package = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = package + spec.loader.exec_module(package) + node = package.NODE_CLASS_MAPPINGS["ScaleLockedResolutionSelector"] + assert issubclass(node, _ComfyNodeInternal), "expected the live-preview node on this ComfyUI" + + info = node.GET_NODE_INFO_V1() + core_info = ResolutionSelector.GET_NODE_INFO_V1() + assert info["input"]["required"]["megapixels"][1]["max"] == 256.0 + assert core_info["input"]["required"]["megapixels"][1]["max"] == 16.0 + # Same widgets, options and preview wiring as the core node, apart from the cap. + assert list(info["input"]["required"]) == list(core_info["input"]["required"]) + assert info["input"]["required"]["aspect_ratio"][0] == core_info["input"]["required"]["aspect_ratio"][0] + assert info["input"].get("optional", {}).keys() == core_info["input"].get("optional", {}).keys() + for section in ("required", "optional"): + for name, spec_ in core_info["input"].get(section, {}).items(): + if name != "megapixels": + assert info["input"][section][name] == spec_, (name, info["input"][section][name], spec_) + assert info["output"] == core_info["output"] and info["output_name"] == core_info["output_name"] + + compared = 0 + for ratio in AspectRatio: + for megapixels in (0.1, 1.0, 2.5, 8.0, 16.0): + for multiple in (8, 16, 64, 128): + ours = node.EXECUTE_NORMALIZED(aspect_ratio=ratio.value, megapixels=megapixels, multiple=multiple).args + core = ResolutionSelector.EXECUTE_NORMALIZED(aspect_ratio=ratio.value, megapixels=megapixels, + multiple=multiple).args + assert tuple(ours) == tuple(core), (ratio, megapixels, multiple, ours, core) + compared += 1 + above = node.EXECUTE_NORMALIZED(aspect_ratio="1:1 (Square)", megapixels=64.0, multiple=16).args + assert tuple(above) == (8192, 8192), above + try: + node.EXECUTE_NORMALIZED(aspect_ratio="21:9 (Ultrawide)", megapixels=256.0, multiple=16) + except ValueError as exc: + assert "canvas limit" in str(exc) + else: + raise AssertionError("canvases beyond ComfyUI's per-side limit must be rejected") + classic = sys.modules[spec.name + ".resolution"].ResolutionSelectorUncappedClassic + assert classic.INPUT_TYPES()["required"]["megapixels"][1]["max"] == 256.0 + assert classic().execute("16:9 (Widescreen)", 2.5, 16) == tuple( + ResolutionSelector.EXECUTE_NORMALIZED(aspect_ratio="16:9 (Widescreen)", megapixels=2.5, multiple=16).args) + print(f"matched core Resolution Selector on {compared} cases; 64 MP square -> {above[0]}x{above[1]}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_krea2_composition_lock.py b/tests/test_krea2_composition_lock.py new file mode 100644 index 0000000..04759db --- /dev/null +++ b/tests/test_krea2_composition_lock.py @@ -0,0 +1,272 @@ +"""CPU checks for the full planner lock used by the Krea 2 Turbo node. + +ComfyUI is replaced by import stubs. The Euler loop below matches ComfyUI's +``sample_euler`` update for flow models. These checks establish trajectory +algebra and gating only; they do not establish image quality. +""" +from __future__ import annotations + +import importlib.util +import math +from pathlib import Path +import sys +import types + +import pytest +import torch + + +ROOT = Path(__file__).resolve().parents[1] +PACKAGE = "slrd_composition_lock_test" + + +def _load_runtime(): + for name in ("comfy", "comfy.model_management", "comfy.model_sampling", "comfy.sample", "comfy.samplers", "comfy.utils", "latent_preview"): + sys.modules.setdefault(name, types.ModuleType(name)) + comfy = sys.modules["comfy"] + for child in ("model_management", "model_sampling", "sample", "samplers", "utils"): + setattr(comfy, child, sys.modules[f"comfy.{child}"]) + package = types.ModuleType(PACKAGE) + package.__path__ = [str(ROOT)] + sys.modules[PACKAGE] = package + for name in ("krea2_flow", "slrd_core", "slrd_runtime"): + spec = importlib.util.spec_from_file_location(f"{PACKAGE}.{name}", ROOT / f"{name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return sys.modules[f"{PACKAGE}.slrd_runtime"], sys.modules[f"{PACKAGE}.slrd_core"] + + +runtime, core = _load_runtime() + + +def _config(full_lock_sigma): + return runtime.ScaleLockConfig( + lock_strength=0.75, lock_strength_start=1.0, lock_strength_end=0.2, + cutoff=0.33, mid_cutoff=0.6, mid_strength=0.2, schedule="hold_then_drop", + schedule_power=2.0, schedule_hold=0.25, mid_strength_start=1.0, + mid_strength_end=0.0, mid_schedule="linked", mid_schedule_power=2.0, + mid_schedule_hold=0.0, full_lock_sigma=full_lock_sigma, + ) + + +def _flow_sigmas(steps, shift=1.15): + # ComfyUI's flux_time_shift form, as used by Krea 2's model sampling. + t = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64) + sigmas = torch.zeros_like(t) + live = t > 0 + sigmas[live] = math.exp(shift) / (math.exp(shift) + (1.0 / t[live] - 1.0)) + return sigmas + + +def _euler(x, sigmas, predict): + for i in range(len(sigmas) - 1): + denoised = predict(x, sigmas[i]) + x = x + (x - denoised) / sigmas[i] * (sigmas[i + 1] - sigmas[i]) + return x + + +class _Guider: + def __init__(self, model_prediction): + self._model_prediction = model_prediction + self.calls = 0 + + def predict_noise(self, x, timestep, model_options=None, seed=None): + self.calls += 1 + return self._model_prediction(x, timestep) + + +def _locked_guider(model_prediction, anchors, planner_sigmas, highres, full_lock_sigma, spatial_mask=None): + guider = _Guider(model_prediction) + config = _config(full_lock_sigma) + config.spatial_mask = spatial_mask + context = runtime.ScaleLockedRuntimeContext( + model=None, highres_latent={"samples": highres}, lowres_latent=None, + lowres_out=None, sigmas=None, anchors_x0=anchors, + planner_sigmas=planner_sigmas, highres_noise=None, noise_seed=0, + ) + return runtime.apply_scale_lock_to_guider(guider, context, config) + + +def _planner(low_shape, sigmas, seed=5): + generator = torch.Generator().manual_seed(seed) + weight = torch.randn(low_shape[1], low_shape[1], generator=generator, dtype=torch.float64) * 0.3 + bias = torch.randn(low_shape, generator=generator, dtype=torch.float64) + + def model(x, sigma): + return torch.einsum("oc,bc...->bo...", weight, x) * float(sigma) + bias * (1.0 - float(sigma)) + + noise = torch.randn(low_shape, generator=generator, dtype=torch.float64) + anchors = [] + + def record(x, sigma): + out = model(x, sigma) + anchors.append(out.clone()) + return out + + final = _euler(noise, sigmas, record) + return noise, anchors, final + + +@pytest.mark.parametrize("layout", ["bchw", "bc1hw"]) +def test_locked_target_state_is_planner_state_with_target_noise(layout): + sigmas = _flow_sigmas(8) + low_shape, high_shape = (1, 16, 8, 10), (1, 16, 14, 18) + if layout == "bc1hw": + low_shape, high_shape = (1, 16, 1, 8, 10), (1, 16, 1, 14, 18) + _, anchors, _ = _planner(low_shape, sigmas) + planner_sigmas = [float(s) for s in sigmas[:-1]] + release = 0.8 + locked_steps = sum(float(s) > release for s in sigmas[:-1]) + assert locked_steps == 4 + + target_noise = torch.randn(high_shape, generator=torch.Generator().manual_seed(9), dtype=torch.float64) + + def displaced_model(x, sigma): + # A target prediction unrelated to the planner composition. + return torch.roll(x, shifts=3, dims=-1) * 0.5 + 2.0 + + guider = _locked_guider(displaced_model, anchors, planner_sigmas, torch.zeros(high_shape, dtype=torch.float64), release) + x = target_noise + for i in range(locked_steps): + denoised = guider.predict_noise(x, sigmas[i]) + assert torch.equal(denoised, core.resize_4d_tensor(anchors[i], high_shape[-2:])) + x = x + (x - denoised) / sigmas[i] * (sigmas[i + 1] - sigmas[i]) + + # Flow Euler keeps x_k = sigma_k * noise + (1 - sigma_k) * A_k. While locked, + # A_k is the planner's own accumulated prediction, resized to the target. + low_x = torch.zeros(low_shape, dtype=torch.float64) + planner_clean = _euler(low_x, sigmas[: locked_steps + 1], lambda _x, s: anchors[[float(v) for v in sigmas].index(float(s))]) + sigma = float(sigmas[locked_steps]) + expected = sigma * target_noise + core.resize_4d_tensor(planner_clean, high_shape[-2:]) + assert torch.allclose(x, expected, atol=1e-10, rtol=1e-10) + assert guider.calls == locked_steps + + +@pytest.mark.parametrize("steps, release", [(8, 0.8), (12, 0.8), (8, 0.6), (5, 0.0)]) +def test_closed_form_handoff_matches_locked_euler_steps(steps, release): + sigmas = _flow_sigmas(steps) + low_shape, high_shape = (1, 16, 8, 10), (1, 16, 14, 18) + _, anchors, _ = _planner(low_shape, sigmas) + planner_sigmas = [float(s) for s in sigmas[:-1]] + locked = core.full_lock_prefix_length([float(s) for s in sigmas], release) + assert locked == sum(float(s) > release for s in sigmas[:-1]) + + target_noise = torch.randn(high_shape, generator=torch.Generator().manual_seed(4), dtype=torch.float64) + guider = _locked_guider(lambda x, s: x * 0.0 + 5.0, anchors, planner_sigmas, + torch.zeros(high_shape, dtype=torch.float64), release) + x = target_noise + for i in range(locked): + denoised = guider.predict_noise(x, sigmas[i]) + x = x + (x - denoised) / sigmas[i] * (sigmas[i + 1] - sigmas[i]) + + clean = core.flow_euler_locked_clean_state( + torch.zeros(high_shape, dtype=torch.float64), + (core.resize_4d_tensor(anchors[i], high_shape[-2:]) for i in range(locked)), + [float(s) for s in sigmas[: locked + 1]], + ) + sigma = float(sigmas[locked]) + # Same tolerance class as float32 accumulation in the closed form. + assert torch.allclose(sigma * target_noise + (1.0 - sigma) * clean.double(), x, atol=1e-5, rtol=1e-5) + + +def test_closed_form_handoff_rejects_invalid_schedules(): + zero = torch.zeros(1, 16, 4, 4) + with pytest.raises(ValueError, match="one more sigma"): + core.flow_euler_locked_clean_state(zero, [zero], [1.0]) + with pytest.raises(ValueError, match="decrease"): + core.flow_euler_locked_clean_state(zero, [zero], [0.5, 0.7]) + assert core.full_lock_prefix_length([1.0, 0.5, 0.0], None) == 0 + assert core.full_lock_prefix_length([1.0, 0.9, 0.0], 0.0) == 2 + + +def test_release_threshold_is_strict_and_none_disables(): + sigmas = _flow_sigmas(8) + low_shape, high_shape = (1, 16, 8, 10), (1, 16, 16, 20) + _, anchors, _ = _planner(low_shape, sigmas) + planner_sigmas = [float(s) for s in sigmas[:-1]] + + def model(x, sigma): + return torch.full_like(x, 7.0) + + x = torch.zeros(high_shape, dtype=torch.float64) + full = _locked_guider(model, anchors, planner_sigmas, x, 1.0) + # sigma == threshold is not above it, so 1.0 disables the full lock. + assert not torch.equal(full.predict_noise(x, sigmas[0]), core.resize_4d_tensor(anchors[0], high_shape[-2:])) + + zero = _locked_guider(model, anchors, planner_sigmas, x, 0.0) + for i in range(len(planner_sigmas)): + assert torch.equal(zero.predict_noise(x, sigmas[i]), core.resize_4d_tensor(anchors[i], high_shape[-2:])) + + legacy = _locked_guider(model, anchors, planner_sigmas, x, None) + reference = _Guider(model) + core.init_scale_lock_state( + reference, model=None, anchors_x0_cpu=anchors, planner_sigmas=planner_sigmas, + lock_strength=0.75, lock_strength_start=1.0, lock_strength_end=0.2, cutoff=0.33, + mid_cutoff=0.6, mid_strength=0.2, schedule="hold_then_drop", schedule_power=2.0, + schedule_hold=0.25, mid_strength_start=1.0, mid_strength_end=0.0, mid_schedule="linked", + mid_schedule_power=2.0, mid_schedule_hold=0.0, spatial_mask=None, + ) + for i in range(len(planner_sigmas)): + expected = runtime.apply_scale_lock_to_noise_prediction(reference, model(x, sigmas[i]), x, sigmas[i]) + assert torch.equal(legacy.predict_noise(x, sigmas[i]), expected) + + +def test_locked_prediction_does_not_alias_stored_anchor_and_respects_mask(): + sigmas = _flow_sigmas(4) + shape = (1, 16, 8, 10) + _, anchors, _ = _planner(shape, sigmas) + planner_sigmas = [float(s) for s in sigmas[:-1]] + + def model(x, sigma): + return torch.full_like(x, 3.0) + + x = torch.zeros(shape, dtype=torch.float64) + guider = _locked_guider(model, anchors, planner_sigmas, x, 0.5) + out = guider.predict_noise(x, sigmas[0]) + assert torch.equal(out, anchors[0]) and out.data_ptr() != anchors[0].data_ptr() + out.add_(1.0) + assert not torch.equal(out, anchors[0]) + + mask = torch.zeros((1, 1, 8, 10), dtype=torch.float64) + mask[..., :, :5] = 1.0 + masked = _locked_guider(model, anchors, planner_sigmas, x, 0.5, spatial_mask=mask) + out = masked.predict_noise(x, sigmas[0]) + assert torch.allclose(out[..., :5], anchors[0][..., :5], atol=1e-12, rtol=0.0) + assert torch.allclose(out[..., 5:], torch.full_like(out[..., 5:], 3.0), atol=1e-12, rtol=0.0) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -0.1]) +def test_invalid_threshold_rejected(value): + with pytest.raises(ValueError, match="full_lock_sigma"): + core.init_scale_lock_state( + _Guider(None), model=None, anchors_x0_cpu=[], planner_sigmas=None, + lock_strength=1.0, lock_strength_start=1.0, lock_strength_end=1.0, cutoff=0.33, + mid_cutoff=0.6, mid_strength=0.0, schedule="flat", schedule_power=2.0, + schedule_hold=0.0, mid_strength_start=1.0, mid_strength_end=0.0, mid_schedule="linked", + mid_schedule_power=2.0, mid_schedule_hold=0.0, spatial_mask=None, full_lock_sigma=value, + ) + + +def test_stage_planning_is_geometric_and_bounded(): + planner, target = 3.0e6, 14.0e6 + stages = core.plan_stage_megapixels(planner, target, 4.0) + assert len(stages) == 1 + assert stages[0] == pytest.approx((3.0 * 14.0) ** 0.5) + sizes = [planner / 1e6, *stages, target / 1e6] + assert all(b / a <= 4.0 + 1e-9 for a, b in zip(sizes, sizes[1:])) + assert len(core.plan_stage_megapixels(1.0e6, 64.0e6, 4.0)) == 2 + # Exactly the limit needs no intermediate stage, despite rounding in the logarithm. + assert core.plan_stage_megapixels(1048576.0, 4194304.0, 4.0) == [] + assert core.plan_stage_megapixels(1.0e6, 14.0e6, 0.0) == [] + with pytest.raises(ValueError, match="max_stage_upscale"): + core.plan_stage_megapixels(1.0e6, 14.0e6, 1.0) + + +def test_stage_megapixels_parsing(): + assert core.parse_stage_megapixels("") == [] + assert core.parse_stage_megapixels(" 4, 8.5 ;12 ") == [4.0, 8.5, 12.0] + with pytest.raises(ValueError, match="not a number"): + core.parse_stage_megapixels("4, big") + with pytest.raises(ValueError, match="positive"): + core.parse_stage_megapixels("4, -1") diff --git a/tests/test_krea2_flow.py b/tests/test_krea2_flow.py new file mode 100644 index 0000000..251c842 --- /dev/null +++ b/tests/test_krea2_flow.py @@ -0,0 +1,124 @@ +"""CPU numerical tests for Krea 2's phase-coupled noise transport.""" +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch + + +spec = importlib.util.spec_from_file_location( + "krea2_flow", Path(__file__).resolve().parents[1] / "krea2_flow.py" +) +module = importlib.util.module_from_spec(spec) +spec.loader.exec_module(module) +build_patch_coupled_noise = module.build_patch_coupled_noise + + +@pytest.mark.parametrize( + "low, high", + [ + ((1, 16, 8, 10), (1, 16, 16, 20)), + ((2, 16, 8, 10), (2, 16, 14, 24)), + ((1, 16, 8, 10), (1, 16, 8, 14)), + ((1, 16, 10, 6), (1, 16, 16, 10)), + ], +) +def test_phase_projection(low, high): + source = torch.randn(low, generator=torch.Generator().manual_seed(42)) + original = source.clone() + target = build_patch_coupled_noise(source, high, seed=123) + assert target.shape == high + assert torch.equal(source, original) + + sy, sx = low[-2] // 2, low[-1] // 2 + ty, tx = high[-2] // 2, high[-1] // 2 + owner_y = torch.arange(ty) * sy // ty + owner_x = torch.arange(tx) * sx // tx + counts = ( + torch.bincount(owner_y, minlength=sy)[:, None] + * torch.bincount(owner_x, minlength=sx)[None, :] + ) + for py in range(2): + for px in range(2): + phase = target[..., py::2, px::2] + row_sums = torch.zeros((*phase.shape[:2], sy, tx)) + row_sums.index_add_(2, owner_y, phase) + sums = torch.zeros((*phase.shape[:2], sy, sx)) + sums.index_add_(3, owner_x, row_sums) + assert torch.allclose( + sums / counts.sqrt(), source[..., py::2, px::2], + atol=1e-5, rtol=1e-5, + ) + + +def test_determinism_and_marginal_statistics(): + source = torch.randn((3, 16, 60, 60), generator=torch.Generator().manual_seed(888)) + first = build_patch_coupled_noise(source, (3, 16, 120, 120), seed=9) + second = build_patch_coupled_noise(source, (3, 16, 120, 120), seed=9) + third = build_patch_coupled_noise(source, (3, 16, 120, 120), seed=10) + assert torch.equal(first, second) + assert not torch.equal(first, third) + assert abs(float(first.mean())) < 0.025 + assert 0.97 < float(first.std()) < 1.03 + + +def test_dtype_preserved(): + source = torch.randn((1, 16, 8, 10)).to(torch.bfloat16) + out = build_patch_coupled_noise(source, (1, 16, 14, 22), seed=2) + assert out.dtype == torch.bfloat16 + assert bool(torch.isfinite(out).all()) + + +@pytest.mark.parametrize( + "low, target", + [ + (torch.ones((1, 16, 9, 10)), (1, 16, 16, 20)), + (torch.ones((1, 16, 8, 10)), (1, 16, 6, 20)), + (torch.ones((1, 16, 8, 10)), (1, 16, 8, 10)), + (torch.ones((1, 16, 8, 10)), (1, 8, 16, 20)), + ], +) +def test_invalid_geometries_rejected(low, target): + with pytest.raises((ValueError, TypeError)): + build_patch_coupled_noise(low, target, seed=1) + + +def test_nonfinite_source_rejected(): + source = torch.ones((1, 16, 8, 10)) + source[0, 0, 0, 0] = float("nan") + with pytest.raises(ValueError, match="finite"): + build_patch_coupled_noise(source, (1, 16, 16, 20), seed=1) + + +def test_single_frame_noise_preserves_native_layout_and_four_dimensional_coupling(): + source = torch.randn((2, 16, 1, 8, 10), generator=torch.Generator().manual_seed(73)) + target = build_patch_coupled_noise(source, (2, 16, 1, 14, 24), seed=19) + expected = build_patch_coupled_noise(source.squeeze(2), (2, 16, 14, 24), seed=19) + assert target.shape == (2, 16, 1, 14, 24) + assert torch.equal(target.squeeze(2), expected) + + +def test_multiple_temporal_slices_are_not_image_noise(): + with pytest.raises(ValueError, match="single temporal slice"): + build_patch_coupled_noise(torch.zeros(1, 16, 2, 8, 10), (1, 16, 2, 16, 20), seed=1) + + +def test_single_frame_spatial_helpers_preserve_native_latent_layout(): + import sys + core_spec = importlib.util.spec_from_file_location( + "slrd_core_image_test", Path(__file__).resolve().parents[1] / "slrd_core.py" + ) + core = importlib.util.module_from_spec(core_spec) + sys.modules[core_spec.name] = core + core_spec.loader.exec_module(core) + source = torch.randn((2, 16, 1, 8, 10), generator=torch.Generator().manual_seed(23)) + image = source.squeeze(2) + for size in ((4, 6), (14, 24), (8, 10)): + out = core.resize_4d_tensor(source, size) + assert out.shape == (2, 16, 1, *size) + assert torch.equal(out.squeeze(2), core.resize_4d_tensor(image, size)) + assert core.latent_target_hw_from_megapixels(source, 1.0) == core.latent_target_hw_from_megapixels(image, 1.0) + with pytest.raises(ValueError, match="image latent"): + core.latent_target_hw_from_megapixels(source.repeat(1, 1, 2, 1, 1), 1.0)