Fast_deploy / ComfyUI-MiniMaxH3-Context-Loop /tests /_drift_control_av_unit_test.py

Upload folder using huggingface_hub (part 2)

bc50ea0 verified about 1 month ago

8.5 kB

#!/usr/bin/env python3
"""Focused CPU tests for schedule-matched Drift-Control AV masks."""

import copy
import importlib.util
import pathlib
import sys
import types

import torch


ROOT = pathlib.Path(__file__).resolve().parents[1]
PACKAGE_NAME = "h3_drift_control_test_package"
package = types.ModuleType(PACKAGE_NAME)
package.__path__ = [str(ROOT)]
sys.modules[PACKAGE_NAME] = package


def load_module(name):
    qualified = PACKAGE_NAME + "." + name
    spec = importlib.util.spec_from_file_location(
        qualified, ROOT / (name + ".py"))
    module = importlib.util.module_from_spec(spec)
    sys.modules[qualified] = module
    spec.loader.exec_module(module)
    return module


contracts = load_module("contracts_v05")
drift = load_module("drift_control")


schedule = torch.tensor([1.0, 0.8, 0.5, 0.2, 0.0])
assert abs(drift.next_schedule_sigma(0.8, schedule) - 0.5) < 1e-7
assert abs(drift.next_schedule_sigma(0.65, schedule) - 0.5) < 1e-7
assert abs(drift.next_schedule_sigma(1.1, schedule) - 1.0) < 1e-7
assert drift.next_schedule_sigma(0.0, schedule) == 0.0
assert abs(drift.matched_noise_ratio(0.8, schedule) - 0.625) < 1e-7
assert drift.matched_noise_ratio(0.2, schedule) == 0.0

weights = drift.temporal_prefix_weights(12, 4)
assert weights == (1.0,) * 8 + (0.75, 0.5, 0.25, 0.0)

video_shape = (1, 2, 15, 2, 2)
video_elements = 2 * 15 * 2 * 2
audio_elements = 24
base = torch.ones((1, 1, video_elements + audio_elements))
dynamic, h3_mask = drift.apply_dynamic_prefix_mask(
    base, video_shape, 12, 0.5, 4)
video = dynamic[..., :video_elements].reshape(video_shape)
for index, expected in enumerate(
        (0.5,) * 8 + (0.375, 0.25, 0.125, 0.0)):
    assert torch.allclose(
        video[:, :, index], torch.full_like(video[:, :, index], expected))
assert torch.equal(video[:, :, 12:], torch.ones_like(video[:, :, 12:]))
assert torch.equal(dynamic[..., video_elements:], base[..., video_elements:])
assert tuple(h3_mask.shape) == (1, 1, 15, 2, 2)

# The sampler uses the exact continuous ratio while H3 receives the final
# merged PR's ceil-quantized 1/256 token-grid convention.
third, quantized = drift.apply_dynamic_prefix_mask(
    base, video_shape, 12, 1.0 / 3.0, 4)
assert abs(float(third[..., 0]) - (1.0 / 3.0)) < 1e-7
assert abs(float(quantized[..., 0, 0, 0]) - (86.0 / 256.0)) < 1e-7

state = drift._DriftControlMaskState(video_shape, 12)
updated = state.denoise_mask_function(
    torch.tensor([0.8]), base, {"sigmas": schedule})
assert abs(state.last_ratio - 0.625) < 1e-7
assert state.current_video_mask is not None
assert abs(float(updated[..., 0]) - 0.625) < 1e-7
captured = {}


def executor(*args, **kwargs):
    captured.update(kwargs)
    return "ok"


assert state.apply_model_wrapper(
    executor, "x", denoise_mask=torch.ones((1, 1, 15, 2, 2))) == "ok"
assert captured["denoise_mask"] is state.current_video_mask

# A sigma-split sampler exposes only its local schedule.  Supplying the full
# unsplit schedule keeps the boundary's next-sigma ratio identical across
# switched model branches instead of incorrectly treating the split as zero.
split_local = torch.tensor([1.0, 0.8])
split_without_override = drift._DriftControlMaskState(video_shape, 12)
split_without_override.denoise_mask_function(
    torch.tensor([0.8]), base, {"sigmas": split_local})
assert split_without_override.last_ratio == 0.0
split_with_override = drift._DriftControlMaskState(
    video_shape, 12, schedule_override=schedule)
split_updated = split_with_override.denoise_mask_function(
    torch.tensor([0.8]), base, {"sigmas": split_local})
assert abs(split_with_override.last_ratio - 0.625) < 1e-7
assert abs(float(split_updated[..., 0]) - 0.625) < 1e-7

try:
    drift._DriftControlMaskState(video_shape, 11)
except ValueError as exc:
    assert "requires 8 matched + 4 taper" in str(exc)
else:
    raise AssertionError("Drift-Control AV accepted a non-v1 prefix layout")


class ModelType:
    name = "FLOW_AV"


class InnerModel:
    model_type = ModelType()


class FakeModel:
    def __init__(self):
        self.model = InnerModel()
        self.model_options = {}
        self.mask_function = None
        self.wrappers = []

    def clone(self):
        cloned = FakeModel()
        cloned.model_options = copy.copy(self.model_options)
        return cloned

    def set_model_denoise_mask_function(self, function):
        self.mask_function = function
        self.model_options["denoise_mask_function"] = function

    def add_wrapper_with_key(self, wrapper_type, key, wrapper):
        self.wrappers.append((wrapper_type, key, wrapper))


comfy = types.ModuleType("comfy")
patcher_extension = types.ModuleType("comfy.patcher_extension")


class WrappersMP:
    APPLY_MODEL = "apply_model"
    DIFFUSION_MODEL = "diffusion_model"
    SAMPLER_SAMPLE = "sampler_sample"


patcher_extension.WrappersMP = WrappersMP
comfy.patcher_extension = patcher_extension
sys.modules["comfy"] = comfy
sys.modules["comfy.patcher_extension"] = patcher_extension

latent = {"samples": [
    torch.zeros(video_shape),
    torch.zeros((1, 3, 2, 24)),
]}
source_model = FakeModel()
patched = drift.install_drift_control_av_model(source_model, latent, 12)
assert patched is not source_model
assert source_model.model_options == {}
assert callable(patched.model_options["denoise_mask_function"])
assert patched.model_options[
    "h3_context_loop_drift_control_av_recipe"] == (
        contracts.DRIFT_CONTROL_AV_RECIPE)
assert len(patched.wrappers) == 2
assert patched.wrappers[0][0:2] == (
    "apply_model", "h3_context_loop_drift_control_av")
assert patched.wrappers[1][0:2] == (
    "sampler_sample", "h3_context_loop_drift_control_split_handoff")


class FlowSampling:
    noise_scale = 1.0


class InnerSamplerModel:
    model_sampling = FlowSampling()

    @staticmethod
    def process_latent_in(value):
        return value


class FakeGuider:
    inner_model = InnerSamplerModel()


handoff_reference = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]])
stage_a = torch.tensor([[[3.0, 5.0, 7.0, 9.0]]])
zero_noise = torch.zeros_like(stage_a)
handoff_state = drift._DriftControlMaskState(
    video_shape,
    12,
    schedule_override=torch.tensor([1.0, 0.8, 0.5, 0.2, 0.0]),
    reference_samples=handoff_reference,
)
handoff_capture = {}


def handoff_executor(
    guider, sigmas, extra_args, callback, noise, latent_image,
    denoise_mask, disable_pbar,
):
    handoff_capture["noise"] = noise
    handoff_capture["latent_image"] = latent_image
    sigma = float(sigmas[0])
    handoff_capture["start"] = (
        sigma * noise + (1.0 - sigma) * latent_image)
    return "split-ok"


assert handoff_state.sampler_sample_wrapper(
    handoff_executor,
    FakeGuider(),
    torch.tensor([0.5, 0.2, 0.0]),
    {},
    None,
    zero_noise,
    stage_a,
    torch.ones_like(stage_a),
    True,
) == "split-ok"
assert torch.equal(handoff_capture["latent_image"], handoff_reference)
# Stock DisableNoise + stage_a would start at (1-sigma)*stage_a.  The corrected
# call reconstructs that exact state while restoring the original reference.
assert torch.allclose(handoff_capture["start"], 0.5 * stage_a)

stage_one_capture = {}


def stage_one_executor(*args):
    stage_one_capture["noise"] = args[4]
    stage_one_capture["latent_image"] = args[5]
    return "stage-one-ok"


assert handoff_state.sampler_sample_wrapper(
    stage_one_executor,
    FakeGuider(),
    torch.tensor([1.0, 0.8, 0.5]),
    {},
    None,
    zero_noise,
    stage_a,
    torch.ones_like(stage_a),
    True,
) == "stage-one-ok"
assert stage_one_capture["noise"] is zero_noise
assert stage_one_capture["latent_image"] is stage_a

marked = drift.mark_drift_control_latent(latent, 12)
assert marked is not latent
assert drift.drift_control_latent_prefix_steps(marked) == 12
try:
    drift.drift_control_latent_prefix_steps(latent)
except ValueError as exc:
    assert "Chain Context's latent" in str(exc)
else:
    raise AssertionError("The inline patch accepted an unmarked stock latent")

conflicting = FakeModel()
conflicting.model_options["denoise_mask_function"] = lambda *args: args[1]
try:
    drift.install_drift_control_av_model(conflicting, latent, 12)
except ValueError as exc:
    assert "another dynamic denoise-mask" in str(exc)
else:
    raise AssertionError("A conflicting dynamic mask patch was accepted")

print(
    "Drift-Control AV: next-sigma ratios, 8+4 taper, packed sampler mask, "
    "H3 timestep mask, clean-reference split handoff, MODEL patch, and "
    "conflict guard passed")

	#!/usr/bin/env python3
	"""Focused CPU tests for schedule-matched Drift-Control AV masks."""

	import copy
	import importlib.util
	import pathlib
	import sys
	import types

	import torch


	ROOT = pathlib.Path(__file__).resolve().parents[1]
	PACKAGE_NAME = "h3_drift_control_test_package"
	package = types.ModuleType(PACKAGE_NAME)
	package.__path__ = [str(ROOT)]
	sys.modules[PACKAGE_NAME] = package


	def load_module(name):
	qualified = PACKAGE_NAME + "." + name
	spec = importlib.util.spec_from_file_location(
	qualified, ROOT / (name + ".py"))
	module = importlib.util.module_from_spec(spec)
	sys.modules[qualified] = module
	spec.loader.exec_module(module)
	return module


	contracts = load_module("contracts_v05")
	drift = load_module("drift_control")


	schedule = torch.tensor([1.0, 0.8, 0.5, 0.2, 0.0])
	assert abs(drift.next_schedule_sigma(0.8, schedule) - 0.5) < 1e-7
	assert abs(drift.next_schedule_sigma(0.65, schedule) - 0.5) < 1e-7
	assert abs(drift.next_schedule_sigma(1.1, schedule) - 1.0) < 1e-7
	assert drift.next_schedule_sigma(0.0, schedule) == 0.0
	assert abs(drift.matched_noise_ratio(0.8, schedule) - 0.625) < 1e-7
	assert drift.matched_noise_ratio(0.2, schedule) == 0.0

	weights = drift.temporal_prefix_weights(12, 4)
	assert weights == (1.0,) * 8 + (0.75, 0.5, 0.25, 0.0)

	video_shape = (1, 2, 15, 2, 2)
	video_elements = 2 * 15 * 2 * 2
	audio_elements = 24
	base = torch.ones((1, 1, video_elements + audio_elements))
	dynamic, h3_mask = drift.apply_dynamic_prefix_mask(
	base, video_shape, 12, 0.5, 4)
	video = dynamic[..., :video_elements].reshape(video_shape)
	for index, expected in enumerate(
	(0.5,) * 8 + (0.375, 0.25, 0.125, 0.0)):
	assert torch.allclose(
	video[:, :, index], torch.full_like(video[:, :, index], expected))
	assert torch.equal(video[:, :, 12:], torch.ones_like(video[:, :, 12:]))
	assert torch.equal(dynamic[..., video_elements:], base[..., video_elements:])
	assert tuple(h3_mask.shape) == (1, 1, 15, 2, 2)

	# The sampler uses the exact continuous ratio while H3 receives the final
	# merged PR's ceil-quantized 1/256 token-grid convention.
	third, quantized = drift.apply_dynamic_prefix_mask(
	base, video_shape, 12, 1.0 / 3.0, 4)
	assert abs(float(third[..., 0]) - (1.0 / 3.0)) < 1e-7
	assert abs(float(quantized[..., 0, 0, 0]) - (86.0 / 256.0)) < 1e-7

	state = drift._DriftControlMaskState(video_shape, 12)
	updated = state.denoise_mask_function(
	torch.tensor([0.8]), base, {"sigmas": schedule})
	assert abs(state.last_ratio - 0.625) < 1e-7
	assert state.current_video_mask is not None
	assert abs(float(updated[..., 0]) - 0.625) < 1e-7
	captured = {}


	def executor(args, *kwargs):
	captured.update(kwargs)
	return "ok"


	assert state.apply_model_wrapper(
	executor, "x", denoise_mask=torch.ones((1, 1, 15, 2, 2))) == "ok"
	assert captured["denoise_mask"] is state.current_video_mask

	# A sigma-split sampler exposes only its local schedule. Supplying the full
	# unsplit schedule keeps the boundary's next-sigma ratio identical across
	# switched model branches instead of incorrectly treating the split as zero.
	split_local = torch.tensor([1.0, 0.8])
	split_without_override = drift._DriftControlMaskState(video_shape, 12)
	split_without_override.denoise_mask_function(
	torch.tensor([0.8]), base, {"sigmas": split_local})
	assert split_without_override.last_ratio == 0.0
	split_with_override = drift._DriftControlMaskState(
	video_shape, 12, schedule_override=schedule)
	split_updated = split_with_override.denoise_mask_function(
	torch.tensor([0.8]), base, {"sigmas": split_local})
	assert abs(split_with_override.last_ratio - 0.625) < 1e-7
	assert abs(float(split_updated[..., 0]) - 0.625) < 1e-7

	try:
	drift._DriftControlMaskState(video_shape, 11)
	except ValueError as exc:
	assert "requires 8 matched + 4 taper" in str(exc)
	else:
	raise AssertionError("Drift-Control AV accepted a non-v1 prefix layout")


	class ModelType:
	name = "FLOW_AV"


	class InnerModel:
	model_type = ModelType()


	class FakeModel:
	def __init__(self):
	self.model = InnerModel()
	self.model_options = {}
	self.mask_function = None
	self.wrappers = []

	def clone(self):
	cloned = FakeModel()
	cloned.model_options = copy.copy(self.model_options)
	return cloned

	def set_model_denoise_mask_function(self, function):
	self.mask_function = function
	self.model_options["denoise_mask_function"] = function

	def add_wrapper_with_key(self, wrapper_type, key, wrapper):
	self.wrappers.append((wrapper_type, key, wrapper))


	comfy = types.ModuleType("comfy")
	patcher_extension = types.ModuleType("comfy.patcher_extension")


	class WrappersMP:
	APPLY_MODEL = "apply_model"
	DIFFUSION_MODEL = "diffusion_model"
	SAMPLER_SAMPLE = "sampler_sample"


	patcher_extension.WrappersMP = WrappersMP
	comfy.patcher_extension = patcher_extension
	sys.modules["comfy"] = comfy
	sys.modules["comfy.patcher_extension"] = patcher_extension

	latent = {"samples": [
	torch.zeros(video_shape),
	torch.zeros((1, 3, 2, 24)),
	]}
	source_model = FakeModel()
	patched = drift.install_drift_control_av_model(source_model, latent, 12)
	assert patched is not source_model
	assert source_model.model_options == {}
	assert callable(patched.model_options["denoise_mask_function"])
	assert patched.model_options[
	"h3_context_loop_drift_control_av_recipe"] == (
	contracts.DRIFT_CONTROL_AV_RECIPE)
	assert len(patched.wrappers) == 2
	assert patched.wrappers[0][0:2] == (
	"apply_model", "h3_context_loop_drift_control_av")
	assert patched.wrappers[1][0:2] == (
	"sampler_sample", "h3_context_loop_drift_control_split_handoff")


	class FlowSampling:
	noise_scale = 1.0


	class InnerSamplerModel:
	model_sampling = FlowSampling()

	@staticmethod
	def process_latent_in(value):
	return value


	class FakeGuider:
	inner_model = InnerSamplerModel()


	handoff_reference = torch.tensor([[[1.0, 2.0, 3.0, 4.0]]])
	stage_a = torch.tensor([[[3.0, 5.0, 7.0, 9.0]]])
	zero_noise = torch.zeros_like(stage_a)
	handoff_state = drift._DriftControlMaskState(
	video_shape,
	12,
	schedule_override=torch.tensor([1.0, 0.8, 0.5, 0.2, 0.0]),
	reference_samples=handoff_reference,
	)
	handoff_capture = {}


	def handoff_executor(
	guider, sigmas, extra_args, callback, noise, latent_image,
	denoise_mask, disable_pbar,
	):
	handoff_capture["noise"] = noise
	handoff_capture["latent_image"] = latent_image
	sigma = float(sigmas[0])
	handoff_capture["start"] = (
	sigma * noise + (1.0 - sigma) * latent_image)
	return "split-ok"


	assert handoff_state.sampler_sample_wrapper(
	handoff_executor,
	FakeGuider(),
	torch.tensor([0.5, 0.2, 0.0]),
	{},
	None,
	zero_noise,
	stage_a,
	torch.ones_like(stage_a),
	True,
	) == "split-ok"
	assert torch.equal(handoff_capture["latent_image"], handoff_reference)
	# Stock DisableNoise + stage_a would start at (1-sigma)*stage_a. The corrected
	# call reconstructs that exact state while restoring the original reference.
	assert torch.allclose(handoff_capture["start"], 0.5 * stage_a)

	stage_one_capture = {}


	def stage_one_executor(*args):
	stage_one_capture["noise"] = args[4]
	stage_one_capture["latent_image"] = args[5]
	return "stage-one-ok"


	assert handoff_state.sampler_sample_wrapper(
	stage_one_executor,
	FakeGuider(),
	torch.tensor([1.0, 0.8, 0.5]),
	{},
	None,
	zero_noise,
	stage_a,
	torch.ones_like(stage_a),
	True,
	) == "stage-one-ok"
	assert stage_one_capture["noise"] is zero_noise
	assert stage_one_capture["latent_image"] is stage_a

	marked = drift.mark_drift_control_latent(latent, 12)
	assert marked is not latent
	assert drift.drift_control_latent_prefix_steps(marked) == 12
	try:
	drift.drift_control_latent_prefix_steps(latent)
	except ValueError as exc:
	assert "Chain Context's latent" in str(exc)
	else:
	raise AssertionError("The inline patch accepted an unmarked stock latent")

	conflicting = FakeModel()
	conflicting.model_options["denoise_mask_function"] = lambda *args: args[1]
	try:
	drift.install_drift_control_av_model(conflicting, latent, 12)
	except ValueError as exc:
	assert "another dynamic denoise-mask" in str(exc)
	else:
	raise AssertionError("A conflicting dynamic mask patch was accepted")

	print(
	"Drift-Control AV: next-sigma ratios, 8+4 taper, packed sampler mask, "
	"H3 timestep mask, clean-reference split handoff, MODEL patch, and "
	"conflict guard passed")