| import importlib.util |
| import sys |
| import types |
| from pathlib import Path |
|
|
|
|
| def _load_encode_module(monkeypatch): |
| captured = {} |
| package = types.ModuleType("prompt_relay_under_test") |
| package.__path__ = [] |
| utils = types.ModuleType("prompt_relay_under_test.utils") |
| patches = types.ModuleType("prompt_relay_under_test.patches") |
|
|
| utils.get_raw_tokenizer = lambda _clip: object() |
| utils.map_token_indices = lambda _tokenizer, _global, _locals: ("full prompt", [(1, 2)]) |
| utils.build_segments = lambda *_args: ["segment"] |
| utils.create_mask_fn = lambda *_args: "mask" |
|
|
| def distribute_segment_lengths(count, latent_frames, parsed_lengths): |
| captured["distribution"] = (count, latent_frames, parsed_lengths) |
| return [latent_frames] |
|
|
| utils.distribute_segment_lengths = distribute_segment_lengths |
| patches.detect_model_type = lambda _model: ("ltx", (1, 1, 1), 8) |
| patches.apply_patches = lambda *_args: None |
|
|
| monkeypatch.setitem(sys.modules, package.__name__, package) |
| monkeypatch.setitem(sys.modules, utils.__name__, utils) |
| monkeypatch.setitem(sys.modules, patches.__name__, patches) |
|
|
| path = Path(__file__).parents[1] / "modules" / "prompt_relay" / "encode.py" |
| spec = importlib.util.spec_from_file_location( |
| "prompt_relay_under_test.encode", |
| path, |
| ) |
| module = importlib.util.module_from_spec(spec) |
| monkeypatch.setitem(sys.modules, spec.name, module) |
| spec.loader.exec_module(module) |
| return module, captured |
|
|
|
|
| def test_encode_relay_preserves_the_first_temporal_latent_frame(monkeypatch): |
| module, captured = _load_encode_module(monkeypatch) |
|
|
| class Model: |
| def clone(self): |
| return self |
|
|
| class Clip: |
| def tokenize(self, text): |
| return text |
|
|
| def encode_from_tokens_scheduled(self, tokens): |
| return tokens |
|
|
| module._encode_relay( |
| Model(), Clip(), 73, 512, 512, "global", "local", "", 0.001, |
| ) |
|
|
| assert captured["distribution"] == (1, 10, None) |
|
|
