"""Update the prompt tests from the SD1.5/CLIP-77 image contract to Klein's. Every pair below is an exact, unique anchor: the script refuses to run if an anchor is missing or ambiguous, so the migration cannot half-apply. """ from __future__ import annotations import pathlib TESTS = pathlib.Path(__file__).resolve().parent / "tests" PAIRS = { "test_clip_prompt_budget.py": [ ( "def test_build_image_prompt_never_exceeds_77_tokens_for_long_prompt() -> None:", "def test_build_image_prompt_stays_inside_the_model_context_for_long_prompt() -> None:", ), ( ' count = companion_prompt.clip_token_count(built)\n' ' assert count <= 77, f"Token count {count} exceeded limit 77"', " count = companion_prompt.model_token_count(built)\n" " assert count <= companion_prompt.MODEL_MAX_TOKENS, (\n" ' f"Token count {count} exceeded the model context"\n' " )", ), ( "def test_sd15_provider_applies_the_clip_77_token_budget(tmp_path) -> None:\n" ' """Whatever the request, the SD1.5 pipeline is never handed over 77 tokens.\n' "\n" " This is the reported failure: the provider used to compose its own prompt and\n" ' send it straight to `StableDiffusionImg2ImgPipeline`, which then died with\n' ' "CLIP can only handle sequences up to 77 tokens". Every prompt now goes through\n' ' `build_image_prompt` + `validate_prompt_for_inference` first.\n' ' """', "def test_klein_provider_applies_the_model_token_budget(tmp_path) -> None:\n" ' """Whatever the request, the FLUX.2 Klein pipeline is measured first.\n' "\n" " The reported failure this replaces: the old SD1.5 provider composed its own\n" " prompt and sent it straight to `StableDiffusionImg2ImgPipeline`, which died\n" ' with "CLIP can only handle sequences up to 77 tokens". Every prompt now goes\n' " through `build_image_prompt` + `validate_prompt_for_model`, and the budget is\n" " the IMAGE model's own 512-token context - not CLIP's 77.\n" ' """', ), ( ' prompt1 = pipe.calls[0]["prompt"]\n' " count1 = companion_prompt.clip_token_count(prompt1)\n" " assert count1 <= companion_prompt.CLIP_MAX_TOKENS", ' prompt1 = pipe.calls[0]["prompt"]\n' " count1 = companion_prompt.model_token_count(prompt1)\n" " assert count1 <= companion_prompt.MODEL_MAX_TOKENS", ), ( ' prompt2 = pipe.calls[1]["prompt"]\n' " count2 = companion_prompt.clip_token_count(prompt2)\n" " assert count2 <= companion_prompt.CLIP_MAX_TOKENS\n" " # Compaction trims the OPTIONAL decoration, never the user's own clothing or\n" " # action request.", ' prompt2 = pipe.calls[1]["prompt"]\n' " count2 = companion_prompt.model_token_count(prompt2)\n" " assert count2 <= companion_prompt.MODEL_MAX_TOKENS\n" " # Nothing is taken away from the user's own clothing or action request, and\n" " # the long prompt fits the model's context without losing its end.", ), ( "def test_sd15_provider_never_hands_the_pipeline_personality_text(tmp_path) -> None:", "def test_klein_provider_never_hands_the_pipeline_personality_text(tmp_path) -> None:", ), ( ' assert companion_prompt.clip_token_count(call["prompt"]) <= 77', " assert (\n" ' companion_prompt.model_token_count(call["prompt"])\n' " <= companion_prompt.MODEL_MAX_TOKENS\n" " )", ), ( ' assert many.lower().count("same face as reference") == 1', ' assert many.lower().count("primary identity") == 1', ), ( 'def _real_token_count(text: str) -> int:\n' ' """Token count from the snapshot\'s own tokenizer, or -1 when not fetched."""\n' " import os\n" "\n" " try:\n" " from transformers import CLIPTokenizer\n" "\n" " from app.config.settings import get_settings\n" "\n" ' path = os.path.join(get_settings().local_image_model_dir, "tokenizer")\n' ' return len(CLIPTokenizer.from_pretrained(path)(text)["input_ids"])\n' " except Exception:\n" " return -1", 'def _real_token_count(text: str) -> int:\n' ' """Token count from the IMAGE model\'s own tokenizer, or -1 when not fetched."""\n' " import os\n" "\n" " try:\n" " from transformers import AutoTokenizer\n" "\n" " from app.config.settings import get_settings\n" "\n" ' path = os.path.join(get_settings().local_image_model_dir, "tokenizer")\n' " tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True)\n" ' return len(tokenizer(text, add_special_tokens=False)["input_ids"])\n' " except Exception:\n" " return -1", ), ( "def test_reported_prompt_fits_clip_with_the_models_own_tokenizer() -> None:", "def test_reported_prompt_fits_the_models_own_context() -> None:", ), ( ' pytest.skip("local SD snapshot tokenizer not available")\n' ' assert real <= 77, f"{real} tokens would crash generation: {built!r}"', ' pytest.skip("local FLUX.2 Klein snapshot tokenizer not available")\n' " assert real <= companion_prompt.MODEL_MAX_TOKENS, (\n" ' f"{real} tokens would overrun the model context: {built!r}"\n' " )", ), ( ' assert built.lower().count("same face as reference") == 1', ' assert built.lower().count("primary identity") == 1', ), ( ' """The user\'s own words are the priority: a long direction is compacted\n' " against what the wrapper leaves, and a direction that fills CLIP's whole\n" ' context drops the wrapper entirely instead of being shortened."""', ' """The user\'s own words are the priority: the whole request survives, and\n' ' the wrapper text never displaces it."""', ), ( " assert companion_prompt.clip_token_count(built) <= 77\n" " # Every part the user asked for survives.", " assert (\n" " companion_prompt.model_token_count(built)\n" " <= companion_prompt.MODEL_MAX_TOKENS\n" " )\n" " # Every part the user asked for survives.", ), ], "test_companion_avatar_identity.py": [ ( "def test_the_user_prompt_leads_and_fits_inside_clip_budget() -> None:", "def test_the_user_prompt_leads_and_fits_the_model_context() -> None:", ), ( " assert companion_prompt._clip_token_count(prompt) <= companion_prompt.CLIP_MAX_TOKENS", " assert (\n" " companion_prompt.model_token_count(prompt)\n" " <= companion_prompt.MODEL_MAX_TOKENS\n" " )", ), ( "def test_a_prompt_longer_than_clip_never_reaches_the_model() -> None:\n" ' """The reported crash as a regression test: over-long input must still fit."""', "def test_a_prompt_longer_than_the_model_context_never_reaches_the_model() -> None:\n" ' """An over-long input must be trimmed to the model\'s own measured budget."""', ), ( " assert companion_prompt.clip_token_count(built) <= companion_prompt.CLIP_MAX_TOKENS, (\n" ' f"{companion_prompt.clip_token_count(built)} tokens would crash: {built!r}"\n' " )", " counted = companion_prompt.model_token_count(built)\n" " assert counted <= companion_prompt.MODEL_MAX_TOKENS, (\n" ' f"{counted} tokens would overrun the model context: {built!r}"\n' " )", ), ( ' assert built.startswith("In this scene the companion is smiling warmly")\n' ' assert "same face as reference" in built\n' " assert companion_prompt.clip_token_count(built) <= companion_prompt.CLIP_MAX_TOKENS", ' assert built.startswith("In this scene the companion is smiling warmly")\n' ' assert "primary identity" in built\n' " assert (\n" " companion_prompt.model_token_count(built)\n" " <= companion_prompt.MODEL_MAX_TOKENS\n" " )", ), ( ' assert built.startswith("a flying kiss")\n' " assert companion_prompt.clip_token_count(built) <= companion_prompt.CLIP_MAX_TOKENS", ' assert built.startswith("a flying kiss")\n' " assert (\n" " companion_prompt.model_token_count(built)\n" " <= companion_prompt.MODEL_MAX_TOKENS\n" " )", ), ( "def test_the_identity_clause_is_short_and_the_action_stays_first() -> None:\n" ' """The avatar supplies identity; the short rule reinforces it after the action."""\n' " tokenizer = _local_tokenizer()\n" " if tokenizer is None:\n" " return\n" ' full = len(tokenizer(companion_prompt._IDENTITY_RULE)["input_ids"])\n' ' compact = len(tokenizer(companion_prompt._IMAGE_IDENTITY_RULE)["input_ids"])\n' " assert compact < full\n" " # Identity and continuity rules must not truncate even a simple user action.\n" " built = companion_prompt.build_image_prompt(\n" ' "Give me a flying kiss while smiling at the camera.", "Aria", "",\n' " has_avatar=True,\n" " )\n" ' assert len(tokenizer(built)["input_ids"]) < companion_prompt.CLIP_MAX_TOKENS\n' ' assert "flying kiss" in built', "def test_the_action_stays_first_and_the_identity_clause_follows_it() -> None:\n" ' """The instruction leads; the identity clause is an instruction to the EDITOR.\n' "\n" " FLUX.2 Klein sees the reference image itself, so the identity clause does not\n" " have to be squeezed into a 77-token text context to carry a likeness. What it\n" " must do is follow the user's own words and stay inside the model's context.\n" ' """\n' " built = companion_prompt.build_image_prompt(\n" ' "Give me a flying kiss while smiling at the camera.", "Aria", "",\n' " has_avatar=True,\n" " )\n" ' assert built.startswith("Give me a flying kiss while smiling at the camera.")\n' ' assert "primary identity" in built\n' ' assert built.index("primary identity") > built.index("flying kiss")\n' " assert (\n" " companion_prompt.model_token_count(built)\n" " <= companion_prompt.MODEL_MAX_TOKENS\n" " )", ), ], "test_visual_profile.py": [ ( "def test_the_prompt_still_fits_clip_however_much_is_stored() -> None:", "def test_the_prompt_still_fits_the_models_context_however_much_is_stored() -> None:", ), ], } # Image-path budget asserts: every `clip_token_count` in this file measures an # IMAGE prompt, so all of them move to the image model's own counter. BULK = { "test_visual_profile.py": ( " assert companion_prompt.clip_token_count(built) <= companion_prompt.CLIP_MAX_TOKENS", " assert (\n" " companion_prompt.model_token_count(built)\n" " <= companion_prompt.MODEL_MAX_TOKENS\n" " )", ) } def main() -> int: for filename, pairs in PAIRS.items(): path = TESTS / filename text = path.read_text(encoding="utf-8") for old, new in pairs: if text.count(old) != 1: print( "FAIL %s: anchor appears %d times: %r" % (filename, text.count(old), old[:60]) ) return 1 text = text.replace(old, new, 1) path.write_text(text, encoding="utf-8", newline="\n") print("OK %s (%d rewrites)" % (filename, len(pairs))) for filename, (old, new) in BULK.items(): path = TESTS / filename text = path.read_text(encoding="utf-8") count = text.count(old) if count == 0: print("FAIL %s: bulk anchor not found" % filename) return 1 text = text.replace(old, new) path.write_text(text, encoding="utf-8", newline="\n") print("OK %s (bulk %d)" % (filename, count)) return 0 if __name__ == "__main__": raise SystemExit(main())