diff --git a/__init__.py b/__init__.py index a677ee5..6c6ac7b 100644 --- a/__init__.py +++ b/__init__.py @@ -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", }) diff --git a/docs/krea-v12-native-studies.md b/docs/krea-v12-native-studies.md index 63aff95..d36a190 100644 --- a/docs/krea-v12-native-studies.md +++ b/docs/krea-v12-native-studies.md @@ -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 diff --git a/kg_krea_v12/__init__.py b/kg_krea_v12/__init__.py index 2747e59..48fd2f9 100644 --- a/kg_krea_v12/__init__.py +++ b/kg_krea_v12/__init__.py @@ -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'] diff --git a/kg_krea_v12/latents.py b/kg_krea_v12/latents.py new file mode 100644 index 0000000..d24c284 --- /dev/null +++ b/kg_krea_v12/latents.py @@ -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 diff --git a/kg_krea_v12/native.py b/kg_krea_v12/native.py index 53e28a6..077651a 100644 --- a/kg_krea_v12/native.py +++ b/kg_krea_v12/native.py @@ -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: diff --git a/kg_krea_v12/nodes.py b/kg_krea_v12/nodes.py index fd62de3..3cd9e40 100644 --- a/kg_krea_v12/nodes.py +++ b/kg_krea_v12/nodes.py @@ -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']) diff --git a/tests/test_krea_v12_target_preparation.py b/tests/test_krea_v12_target_preparation.py new file mode 100644 index 0000000..d5e6fe9 --- /dev/null +++ b/tests/test_krea_v12_target_preparation.py @@ -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()