Own native target preparation before the sampler context split
Add an optional preparation node that calls the verified native latent helper and returns both the prepared context and its exact shared target LATENT. Consume raw downscale metadata, preserve mask and foreign packet values, and bind provenance to the original packet, model, format and helper state. Keep existing consumer identity checks and the original seven node interfaces. Reject spatial resizing after an owned grounded wrapper has primed its VAE cache. Disabled preparation returns the original context and target unchanged. All 201 tests and required compile checks pass. Exact native caller/helper fixtures pass on Python 3.11 and 3.13; serving-host acceptance remains pending.
This commit is contained in:
Kevin Gilper
+3
-1
@@ -48,7 +48,7 @@ NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
from .kg_krea_v12 import (
|
||||
KGKrea2PrepareNativeContextV12, KGKrea2NativeAxisV12, KGKrea2CombinedNativeStudiesV12,
|
||||
KGKrea2NativePhaseGuiderV12, KGKrea2PrepareOwnedGroundedContextV12, KGKrea2RemoveOwnedReferenceV12,
|
||||
KGKrea2NativeEnhancerV12,
|
||||
KGKrea2NativeEnhancerV12, KGKrea2PrepareNativeTargetV12,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update({
|
||||
@@ -59,6 +59,7 @@ NODE_CLASS_MAPPINGS.update({
|
||||
"KGKrea2PrepareOwnedGroundedContextV12": KGKrea2PrepareOwnedGroundedContextV12,
|
||||
"KGKrea2RemoveOwnedReferenceV12": KGKrea2RemoveOwnedReferenceV12,
|
||||
"KGKrea2NativeEnhancerV12": KGKrea2NativeEnhancerV12,
|
||||
"KGKrea2PrepareNativeTargetV12": KGKrea2PrepareNativeTargetV12,
|
||||
})
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
"KGKrea2PrepareNativeContextV12": "KG Krea 2 Prepare Native Context V12",
|
||||
@@ -68,4 +69,5 @@ NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
"KGKrea2PrepareOwnedGroundedContextV12": "KG Krea 2 Prepare Owned Grounded Context V12",
|
||||
"KGKrea2RemoveOwnedReferenceV12": "KG Krea 2 Remove Owned Reference V12",
|
||||
"KGKrea2NativeEnhancerV12": "KG Krea 2 Native Enhancer V12",
|
||||
"KGKrea2PrepareNativeTargetV12": "KG Krea 2 Prepare Native Target V12",
|
||||
})
|
||||
|
||||
@@ -24,6 +24,7 @@ or model asset to install for this interface.
|
||||
| `KGKrea2PrepareOwnedGroundedContextV12` | Original MODEL before edit patch, clean source/donor, source/target latents, VAE and exact native patch/preparation arguments | Context plus native MODEL, positive, negative, target latent |
|
||||
| `KGKrea2RemoveOwnedReferenceV12` | Owned grounded context, `native:donor` OFF | Rebuilt context plus its native MODEL, positive, negative, target latent |
|
||||
| `KGKrea2NativeEnhancerV12` | Native MODEL before edit patch, enabled (default OFF), strength0..2 | MODEL plus attachment report; OFF/strength0 returns original MODEL |
|
||||
| `KGKrea2PrepareNativeTargetV12` | Final native context and enabled (default OFF) | Prepared context plus its exact shared native target LATENT |
|
||||
|
||||
Optional native positive and negative conditioning inputs must be connected
|
||||
together. They adopt the existing paired encodes without re-encoding. Otherwise,
|
||||
@@ -222,6 +223,40 @@ are retained. It adds no solver or model forward. Native latent samples and the
|
||||
latent noise mask must match the context by identity. That noise mask is separate
|
||||
from a conditioning-area mask, which is unsupported by this initial consumer.
|
||||
|
||||
The native `SamplerCustomAdvanced` prepares latent channels, dimensions and empty
|
||||
latent resolution before it calls the guider. For the active consumer, place
|
||||
`KGKrea2PrepareNativeTargetV12` after the final native context (including any owned
|
||||
grounded rebuild or reference removal). Enable preparation, connect its context
|
||||
output0 to `KGKrea2CombinedNativeStudiesV12`, and connect its LATENT output1 to
|
||||
`SamplerCustomAdvanced.latent_image`. Both branches then share the exact prepared
|
||||
target; the original source latent, model, conditioning and mask remain bound.
|
||||
The original seven node interfaces are unchanged.
|
||||
|
||||
This node delegates transformation to the installed
|
||||
[`fix_empty_latent_channels`](https://github.com/Comfy-Org/ComfyUI/blob/v0.41.0/comfy/sample.py).
|
||||
It consumes the raw spatial/temporal downscale markers so the sampler cannot
|
||||
resize the already prepared empty latent again. The reviewed Krea2 configuration
|
||||
uses [`Wan21`](https://github.com/Comfy-Org/ComfyUI/blob/v0.41.0/comfy/supported_models.py);
|
||||
that format inherits the identity `LatentFormat.fix_empty_latent` method from
|
||||
[`latent_formats.py`](https://github.com/Comfy-Org/ComfyUI/blob/v0.41.0/comfy/latent_formats.py).
|
||||
Selected helper/utility/format source, actual loaded functions, format class and
|
||||
normalization attributes must match before preparation. The prepared B1 image
|
||||
latent has sixteen channels and one temporal position; the sampler's second
|
||||
native preparation preserves its samples identity. Unsupported formats reject.
|
||||
For an owned grounded context, preparation must retain target spatial dimensions:
|
||||
its native edit wrapper has already primed the VAE cache for that size. A spatial
|
||||
resize at this stage rejects rather than causing a hidden VAE re-encode during
|
||||
sampling. Channel/dimension normalization at the same image size remains allowed.
|
||||
|
||||
The output carries an in-process provenance receipt binding the original packet,
|
||||
prepared packet, exact model, format and helper functions. Context verification
|
||||
checks the receipt before consumption; same-valued replacement tensors, changed
|
||||
masks/packets/model revisions or replaced helpers reject. Ordinary tensor version
|
||||
counters are retained; inference tensors still require a fresh caller revision
|
||||
after in-place writes. No target values are compared as a substitute for identity.
|
||||
Disabled preparation returns the original context and target by identity before
|
||||
host reads or helper calls. Neutral routing can omit or disable this node.
|
||||
|
||||
Initial consumption requires the exact pinned Krea2/CFGGuider/KSAMPLER classes,
|
||||
default Euler, Euler ancestral or ER-SDE options, B1 matching noise/latent shapes and a
|
||||
strictly decreasing bounded accepted sigma schedule. Each evaluation must match
|
||||
|
||||
@@ -2,9 +2,9 @@
|
||||
from .nodes import (KGKrea2PrepareNativeContextV12, KGKrea2NativeAxisV12,
|
||||
KGKrea2CombinedNativeStudiesV12, KGKrea2NativePhaseGuiderV12,
|
||||
KGKrea2PrepareOwnedGroundedContextV12, KGKrea2RemoveOwnedReferenceV12,
|
||||
KGKrea2NativeEnhancerV12)
|
||||
KGKrea2NativeEnhancerV12, KGKrea2PrepareNativeTargetV12)
|
||||
|
||||
__all__ = ['KGKrea2PrepareNativeContextV12', 'KGKrea2NativeAxisV12',
|
||||
'KGKrea2CombinedNativeStudiesV12', 'KGKrea2NativePhaseGuiderV12',
|
||||
'KGKrea2PrepareOwnedGroundedContextV12', 'KGKrea2RemoveOwnedReferenceV12',
|
||||
'KGKrea2NativeEnhancerV12']
|
||||
'KGKrea2NativeEnhancerV12', 'KGKrea2PrepareNativeTargetV12']
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Own the native target preparation boundary without replacing native math."""
|
||||
from dataclasses import dataclass, replace
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
|
||||
from .compatibility import _ast_fingerprint, _lookup, _source_maps, verify_loaded_functions
|
||||
from .native import NativeContext, _signature, context_signature
|
||||
|
||||
PROVENANCE_KEY = 'kg_native_v12_target_preparation'
|
||||
_SAMPLE_PINS = {'fix_empty_latent_channels': '2ea7af39c924e442d1d5204ae82c9ee9eed2c684dc72a486e95df55ecc05332b'}
|
||||
_UTILS_PINS = {
|
||||
'repeat_to_batch_size': 'b4537d24a8494c5259eac72f2398550bfa396a2b48a9e1dac5f7207d5497ffaf',
|
||||
'common_upscale': '32dc30d1611b33ea16c21f6b56c184895dcf4996503b7b7918560031f1a8360b',
|
||||
}
|
||||
_FORMAT_PINS = {'LatentFormat.fix_empty_latent': '048115b7383529fd85a0c72cff703b826a12ceb11cc05040285730ba778713d8'}
|
||||
_FORMAT_VALUES = dict(latent_channels=16, latent_dimensions=3,
|
||||
spacial_downscale_ratio=8, temporal_downscale_ratio=4)
|
||||
|
||||
|
||||
def _verified_functions(module, pins):
|
||||
source = Path(inspect.getfile(module)).read_bytes()
|
||||
_, definitions = _source_maps(source)
|
||||
if any(name not in definitions or _ast_fingerprint(definitions[name]) != expected for name, expected in pins.items()):
|
||||
raise RuntimeError('V12 native target preparation source is unverified')
|
||||
verify_loaded_functions(module, source, tuple(pins))
|
||||
bindings = []
|
||||
for name in pins:
|
||||
function = _lookup(module, name)
|
||||
if function.__globals__ is not vars(module):
|
||||
raise RuntimeError('V12 native target preparation globals differ')
|
||||
bindings.append((module, name, function, function.__code__))
|
||||
return tuple(bindings)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NativeLatentHost:
|
||||
sample_module: object
|
||||
utils_module: object
|
||||
formats_module: object
|
||||
comfy_module: object
|
||||
format: object
|
||||
format_class: object
|
||||
base_format_class: object
|
||||
tensor_type: object
|
||||
bindings: tuple
|
||||
|
||||
def verify(self, model):
|
||||
fmt = self.format
|
||||
if (model.get_model_object('latent_format') is not fmt
|
||||
or type(fmt) is not self.format_class
|
||||
or self.formats_module.Wan21 is not self.format_class
|
||||
or self.formats_module.LatentFormat is not self.base_format_class
|
||||
or self.format_class.__bases__ != (self.base_format_class,)
|
||||
or self.sample_module.__dict__.get('comfy') is not self.comfy_module
|
||||
or self.comfy_module.utils is not self.utils_module
|
||||
or any(name in vars(fmt) for name in (*_FORMAT_VALUES, 'fix_empty_latent'))
|
||||
or any(type(getattr(fmt, name, None)) is not int or getattr(fmt, name) != value
|
||||
for name, value in _FORMAT_VALUES.items())
|
||||
or getattr(fmt.fix_empty_latent, '__func__', None) is not self.base_format_class.fix_empty_latent):
|
||||
raise RuntimeError('V12 native Krea target latent format changed or is unsupported')
|
||||
for module, name, function, code in self.bindings:
|
||||
if _lookup(module, name) is not function or function.__code__ is not code:
|
||||
raise RuntimeError('V12 native target preparation function changed')
|
||||
|
||||
|
||||
def load_latent_host(model):
|
||||
import torch
|
||||
import comfy
|
||||
import comfy.sample as sample
|
||||
import comfy.utils as utils
|
||||
import comfy.latent_formats as formats
|
||||
from .consumer import load_consumer_host
|
||||
_, _, model_base = load_consumer_host()
|
||||
if type(model.model) is not model_base.Krea2:
|
||||
raise RuntimeError('V12 target preparation requires the pinned native Krea2 MODEL')
|
||||
bindings = (_verified_functions(sample, _SAMPLE_PINS)
|
||||
+ _verified_functions(utils, _UTILS_PINS)
|
||||
+ _verified_functions(formats, _FORMAT_PINS))
|
||||
host = NativeLatentHost(sample, utils, formats, comfy, model.get_model_object('latent_format'),
|
||||
formats.Wan21, formats.LatentFormat, torch.Tensor, bindings)
|
||||
host.verify(model)
|
||||
return host
|
||||
|
||||
|
||||
def _model_signature(model):
|
||||
return _signature((model, getattr(model, 'patches_uuid', None), getattr(model, 'model_options', None)))
|
||||
|
||||
|
||||
def _packet_signature(packet):
|
||||
return _signature({key: value for key, value in packet.items() if key != PROVENANCE_KEY})
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NativeTargetPreparation:
|
||||
model: object
|
||||
original: dict
|
||||
canonical: dict
|
||||
host: object
|
||||
original_signature: tuple
|
||||
canonical_signature: tuple
|
||||
model_signature: tuple
|
||||
|
||||
def verify(self, model, target):
|
||||
if (model is not self.model or target is not self.canonical
|
||||
or target.get(PROVENANCE_KEY) is not self
|
||||
or _signature(self.original) != self.original_signature
|
||||
or _packet_signature(target) != self.canonical_signature
|
||||
or _model_signature(model) != self.model_signature):
|
||||
raise RuntimeError('V12 owned native target preparation provenance changed')
|
||||
self.host.verify(model)
|
||||
|
||||
|
||||
def verify_prepared_target(model, target):
|
||||
receipt = target.get(PROVENANCE_KEY)
|
||||
if type(receipt) is not NativeTargetPreparation:
|
||||
raise RuntimeError('V12 native target preparation receipt is invalid')
|
||||
receipt.verify(model, target)
|
||||
|
||||
|
||||
def prepare_native_target(model, target, enabled=False):
|
||||
if type(enabled) is not bool:
|
||||
raise ValueError('V12 native target preparation enabled must be boolean')
|
||||
if not enabled:
|
||||
return target
|
||||
if type(target) is not dict or 'samples' not in target or PROVENANCE_KEY in target:
|
||||
raise ValueError('V12 native target preparation requires the original LATENT packet')
|
||||
for name in ('downscale_ratio_spacial', 'downscale_ratio_temporal'):
|
||||
ratio = target.get(name)
|
||||
if ratio is not None and (type(ratio) not in (int, float) or not 0 < ratio < float('inf')):
|
||||
raise ValueError('V12 native target downscale ratio must be finite and positive')
|
||||
from .consumer import native_model_preparation_lease
|
||||
with native_model_preparation_lease(model):
|
||||
host = load_latent_host(model)
|
||||
original = target['samples']
|
||||
if (type(original) is not host.tensor_type or original.is_nested
|
||||
or original.ndim not in (4, 5) or original.shape[0] != 1):
|
||||
raise RuntimeError('V12 native target preparation requires one ordinary image latent')
|
||||
before, model_before = _signature(target), _model_signature(model)
|
||||
samples = host.sample_module.fix_empty_latent_channels(model, original,
|
||||
target.get('downscale_ratio_spacial'), target.get('downscale_ratio_temporal'))
|
||||
host.verify(model)
|
||||
if before != _signature(target) or model_before != _model_signature(model):
|
||||
raise RuntimeError('V12 native target preparation changed its original packet or MODEL')
|
||||
if (type(samples) is not host.tensor_type or samples.is_nested
|
||||
or samples.ndim != 5 or samples.shape[:3] != (1, 16, 1)):
|
||||
raise RuntimeError('V12 native prepared target is outside the Krea image latent format')
|
||||
canonical = target.copy()
|
||||
canonical['samples'] = samples
|
||||
# These describe the raw packet and have already been consumed by the
|
||||
# native helper. Keeping them would apply the same resize again.
|
||||
canonical.pop('downscale_ratio_spacial', None)
|
||||
canonical.pop('downscale_ratio_temporal', None)
|
||||
receipt = NativeTargetPreparation(model, target, canonical, host, before,
|
||||
_packet_signature(canonical), model_before)
|
||||
canonical[PROVENANCE_KEY] = receipt
|
||||
receipt.verify(model, canonical)
|
||||
return canonical
|
||||
|
||||
|
||||
def prepare_context_target(context, enabled=False):
|
||||
if not isinstance(context, NativeContext) or type(enabled) is not bool:
|
||||
raise ValueError('V12 target preparation requires a native context and boolean enabled')
|
||||
target = context.native.get('target_latent')
|
||||
if not enabled:
|
||||
return context, target
|
||||
context.verify()
|
||||
canonical = prepare_native_target(context.native['model'], target, enabled=True)
|
||||
if (context.owned_recipe is not None
|
||||
and target['samples'].shape[-2:] != canonical['samples'].shape[-2:]):
|
||||
# The original edit wrapper has already primed its VAE cache for these
|
||||
# target dimensions. Changing them here would trigger hidden re-encoding
|
||||
# during sampling; a broader owned rebuild is outside this preparation.
|
||||
raise RuntimeError('V12 owned grounded target preparation cannot change spatial dimensions')
|
||||
context.verify()
|
||||
bundle = {**context.native, 'target_latent': canonical}
|
||||
prepared = replace(context, native=bundle,
|
||||
signature=context_signature(context.clip, bundle, context.sources))
|
||||
prepared.verify()
|
||||
return prepared, canonical
|
||||
@@ -100,6 +100,10 @@ class NativeContext:
|
||||
if self.owned_recipe is not None:
|
||||
self.owned_recipe.verify()
|
||||
model = self.native['model']
|
||||
target = self.native.get('target_latent')
|
||||
if isinstance(target, dict) and 'kg_native_v12_target_preparation' in target:
|
||||
from .latents import verify_prepared_target
|
||||
verify_prepared_target(model, target)
|
||||
if getattr(model, 'patches_uuid', None) != self.model_revision:
|
||||
raise RuntimeError('V12 native model patch revision changed after preparation')
|
||||
if context_signature(self.clip, self.native, self.sources) != self.signature:
|
||||
|
||||
@@ -204,3 +204,22 @@ class KGKrea2NativeEnhancerV12:
|
||||
from .enhancer import apply_enhancer
|
||||
return apply_enhancer(kwargs['Native MODEL before edit patch'],
|
||||
enabled=kwargs['Enhancer enabled'], strength=kwargs['Enhancer strength'])
|
||||
|
||||
|
||||
class KGKrea2PrepareNativeTargetV12:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {
|
||||
'Native context': (CONTEXT_TYPE,),
|
||||
'Native target preparation enabled': ('BOOLEAN', {'default': False}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = (CONTEXT_TYPE, 'LATENT')
|
||||
RETURN_NAMES = ('native_context', 'native_target_latent')
|
||||
FUNCTION = 'execute'
|
||||
CATEGORY = 'advanced/conditioning/native studies'
|
||||
|
||||
def execute(self, **kwargs):
|
||||
from .latents import prepare_context_target
|
||||
return prepare_context_target(kwargs['Native context'],
|
||||
enabled=kwargs['Native target preparation enabled'])
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Owned target provenance and caller boundary; native source integration is separate."""
|
||||
import importlib.util
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace as NS
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
spec = importlib.util.spec_from_file_location('target_preparation_v12_tests', ROOT / 'kg_krea_v12/__init__.py',
|
||||
submodule_search_locations=[str(ROOT / 'kg_krea_v12')])
|
||||
package = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = package
|
||||
spec.loader.exec_module(package)
|
||||
from target_preparation_v12_tests import latents, native, nodes, consumer
|
||||
|
||||
|
||||
class Tensor:
|
||||
is_nested = False
|
||||
def __init__(self, shape, value=0):
|
||||
self.shape, self.value, self._version = tuple(shape), value, 0
|
||||
self.ndim = len(shape)
|
||||
|
||||
|
||||
class LatentFormat:
|
||||
latent_channels, latent_dimensions = 16, 3
|
||||
spacial_downscale_ratio, temporal_downscale_ratio = 8, 4
|
||||
def fix_empty_latent(self, value):
|
||||
return value
|
||||
|
||||
|
||||
class Wan21(LatentFormat):
|
||||
pass
|
||||
|
||||
|
||||
class TargetPreparationTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.calls = []
|
||||
fmt = Wan21()
|
||||
utils, formats, comfy, sample = (ModuleType(name) for name in ('fixture_utils', 'fixture_formats', 'fixture_comfy', 'fixture_sample'))
|
||||
comfy.utils = utils
|
||||
formats.Wan21, formats.LatentFormat = Wan21, LatentFormat
|
||||
sample.comfy = comfy
|
||||
def helper(model, value, spatial=None, temporal=None):
|
||||
self.calls.append((model, value, spatial, temporal))
|
||||
shape = list(value.shape)
|
||||
if value.value == 0:
|
||||
shape[1] = 16
|
||||
if spatial is not None and spatial != 8:
|
||||
shape[-2:] = [round(length * spatial / 8) for length in shape[-2:]]
|
||||
if len(shape) == 4:
|
||||
shape.insert(2, 1)
|
||||
if tuple(shape) != value.shape:
|
||||
value = Tensor(shape, value.value)
|
||||
return fmt.fix_empty_latent(value)
|
||||
sample.fix_empty_latent_channels = helper
|
||||
self.host = latents.NativeLatentHost(sample, utils, formats, comfy, fmt, Wan21, LatentFormat, Tensor,
|
||||
((sample, 'fix_empty_latent_channels', helper, helper.__code__),))
|
||||
self.model = NS(model=object(), model_options={}, patches_uuid='owned-model', get_model_object=lambda name: fmt)
|
||||
self.raw = Tensor((1, 4, 3, 5))
|
||||
self.mask, self.foreign, self.source = object(), object(), object()
|
||||
self.target = dict(samples=self.raw, noise_mask=self.mask, foreign=self.foreign, downscale_ratio_spacial=16)
|
||||
self.bundle = dict(model=self.model, positive=[], negative=[], source_latent=self.source,
|
||||
target_latent=self.target, noise_mask=self.mask)
|
||||
self.context = native.prepare_context(NS(), self.model, 'owned prompt', revision='owned-target-v1', native_bundle=self.bundle)
|
||||
host_patch = patch.object(latents, 'load_latent_host', return_value=self.host)
|
||||
host_patch.start()
|
||||
self.addCleanup(host_patch.stop)
|
||||
self.addCleanup(lambda: self.assertFalse(consumer._busy))
|
||||
|
||||
def prepare(self):
|
||||
return latents.prepare_context_target(self.context, enabled=True)
|
||||
|
||||
def test_disabled_is_exact_original_context_and_target_without_host_reads(self):
|
||||
with patch.object(latents, 'load_latent_host', side_effect=AssertionError('OFF host read')):
|
||||
context, target = latents.prepare_context_target(self.context)
|
||||
self.assertIs(context, self.context)
|
||||
self.assertIs(target, self.target)
|
||||
self.assertFalse(self.calls)
|
||||
|
||||
def test_shared_output_consumes_scale_and_second_native_call_is_identity(self):
|
||||
context, target = self.prepare()
|
||||
self.assertIs(context.native['target_latent'], target)
|
||||
self.assertEqual(target['samples'].shape, (1, 16, 1, 6, 10))
|
||||
self.assertEqual(len(self.calls), 1)
|
||||
self.assertNotIn('downscale_ratio_spacial', target)
|
||||
self.assertNotIn('downscale_ratio_temporal', target)
|
||||
again = self.host.sample_module.fix_empty_latent_channels(self.model, target['samples'], target.get('downscale_ratio_spacial'), target.get('downscale_ratio_temporal'))
|
||||
self.assertIs(again, target['samples'])
|
||||
self.assertEqual(self.target['samples'].shape, (1, 4, 3, 5))
|
||||
self.assertEqual(self.target['downscale_ratio_spacial'], 16)
|
||||
context.verify()
|
||||
|
||||
def test_foreign_packet_values_mask_source_and_conditioning_retain_identity(self):
|
||||
context, target = self.prepare()
|
||||
self.assertIs(target['noise_mask'], self.mask)
|
||||
self.assertIs(target['foreign'], self.foreign)
|
||||
for key in ('model', 'positive', 'negative', 'source_latent', 'noise_mask'):
|
||||
self.assertIs(context.native[key], self.bundle[key])
|
||||
self.context.verify()
|
||||
|
||||
def test_nonempty_native_values_keep_channels_resolution_and_mask(self):
|
||||
self.raw.shape, self.raw.ndim, self.raw.value = (1, 16, 3, 5), 4, 7
|
||||
self.context = native.prepare_context(NS(), self.model, 'owned prompt', revision='nonempty', native_bundle=self.bundle)
|
||||
context, target = self.prepare()
|
||||
self.assertEqual(target['samples'].shape, (1, 16, 1, 3, 5))
|
||||
self.assertEqual(target['samples'].value, 7)
|
||||
self.assertIs(target['noise_mask'], self.mask)
|
||||
context.verify()
|
||||
|
||||
def test_different_model_packet_or_same_value_tensor_cannot_use_receipt(self):
|
||||
_, target = self.prepare()
|
||||
receipt = target[latents.PROVENANCE_KEY]
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
receipt.verify(NS(**vars(self.model)), target)
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
receipt.verify(self.model, target.copy())
|
||||
target['samples'] = Tensor(target['samples'].shape)
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
receipt.verify(self.model, target)
|
||||
|
||||
def test_original_and_canonical_mask_mutations_reject(self):
|
||||
context, target = self.prepare()
|
||||
self.target['noise_mask'] = object()
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
context.verify()
|
||||
self.target['noise_mask'] = self.mask
|
||||
target['noise_mask'] = object()
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
context.verify()
|
||||
|
||||
def test_raw_tensor_version_and_model_revision_mutations_reject(self):
|
||||
context, _ = self.prepare()
|
||||
self.raw._version += 1
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
context.verify()
|
||||
self.raw._version -= 1
|
||||
self.model.patches_uuid = 'foreign-revision'
|
||||
with self.assertRaisesRegex(RuntimeError, 'provenance changed'):
|
||||
context.verify()
|
||||
|
||||
def test_format_or_helper_replacement_after_preparation_rejects(self):
|
||||
context, _ = self.prepare()
|
||||
self.host.format.latent_channels = 4
|
||||
with self.assertRaisesRegex(RuntimeError, 'format changed'):
|
||||
context.verify()
|
||||
del self.host.format.latent_channels
|
||||
self.host.sample_module.fix_empty_latent_channels = lambda *args: self.raw
|
||||
with self.assertRaisesRegex(RuntimeError, 'function changed'):
|
||||
context.verify()
|
||||
|
||||
def test_unsupported_batch_nested_duplicate_or_ratio_rejects(self):
|
||||
for change in ({'samples': Tensor((2, 4, 3, 5))}, {'downscale_ratio_spacial': float('nan')},
|
||||
{latents.PROVENANCE_KEY: object()}):
|
||||
with self.subTest(change=change), self.assertRaises((RuntimeError, ValueError)):
|
||||
latents.prepare_native_target(self.model, {**self.target, **change}, True)
|
||||
self.raw.is_nested = True
|
||||
with self.assertRaisesRegex(RuntimeError, 'ordinary image latent'):
|
||||
self.prepare()
|
||||
|
||||
def test_additive_node_exports_both_shared_ports(self):
|
||||
cls = nodes.KGKrea2PrepareNativeTargetV12
|
||||
self.assertEqual(cls.RETURN_TYPES, ('KG_KREA_NATIVE_CONTEXT_V12', 'LATENT'))
|
||||
context, target = cls().execute(**{'Native context': self.context, 'Native target preparation enabled': True})
|
||||
self.assertIs(context.native['target_latent'], target)
|
||||
self.assertEqual(cls.INPUT_TYPES()['required']['Native target preparation enabled'][1]['default'], False)
|
||||
|
||||
def test_preparation_cannot_overlap_an_owned_model_operation(self):
|
||||
with consumer.native_model_preparation_lease(self.model):
|
||||
with self.assertRaisesRegex(RuntimeError, 'busy'):
|
||||
self.prepare()
|
||||
self.assertFalse(self.calls)
|
||||
|
||||
def test_grounded_spatial_change_rejects_without_changing_original_handles(self):
|
||||
context = replace(self.context, owned_recipe=NS(verify=lambda: None))
|
||||
with self.assertRaisesRegex(RuntimeError, 'grounded target preparation cannot change spatial'):
|
||||
latents.prepare_context_target(context, enabled=True)
|
||||
self.assertIs(context.native['target_latent'], self.target)
|
||||
self.assertIs(self.target['samples'], self.raw)
|
||||
self.assertIs(self.target['noise_mask'], self.mask)
|
||||
self.assertIs(context.native['source_latent'], self.source)
|
||||
self.assertEqual(self.raw.shape, (1, 4, 3, 5))
|
||||
context.verify()
|
||||
|
||||
def test_grounded_channel_dimension_preparation_with_same_spatial_shape_is_allowed(self):
|
||||
self.target['downscale_ratio_spacial'] = 8
|
||||
context = native.prepare_context(NS(), self.model, 'owned prompt', revision='grounded-no-resize', native_bundle=self.bundle)
|
||||
context = replace(context, owned_recipe=NS(verify=lambda: None))
|
||||
prepared, target = latents.prepare_context_target(context, enabled=True)
|
||||
self.assertEqual(target['samples'].shape, (1, 16, 1, 3, 5))
|
||||
self.assertIs(prepared.owned_recipe, context.owned_recipe)
|
||||
self.assertIs(target['noise_mask'], self.mask)
|
||||
prepared.verify()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user