Spaces:
Running on Zero
Running on Zero
| """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()) | |