Spaces:
Running on Zero
Running on Zero
Download tests/test_reference_required.py from YashKotak/devatlas-ai-gpu: direct link, hf CLI and curl.
- Browser
- Download file 13.8 kB
-
https://huggingface.co/spaces/YashKotak/devatlas-ai-gpu/resolve/main/tests/test_reference_required.py
- Command line
-
hf download hf://spaces/YashKotak/devatlas-ai-gpu/tests/test_reference_required.py
-
curl -L -o test_reference_required.py https://huggingface.co/spaces/YashKotak/devatlas-ai-gpu/resolve/main/tests/test_reference_required.py
13.8 kB
| """The reference image is what carries a companion's face into a render. | |
| These tests pin the one rule that must never bend: a job that was given an | |
| avatar must render as an image edit, and a job whose avatar cannot be resolved | |
| must FAIL rather than quietly degrade to text-to-image and return a different | |
| person than the user selected. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import base64 | |
| from pathlib import Path | |
| import pytest | |
| from PIL import Image | |
| from app.config.settings import Settings | |
| from app.services.companion_service import CompanionService | |
| from app.services.providers import ( | |
| ImageProvider, | |
| LocalCompanionChatProvider, | |
| LocalMemoryProvider, | |
| LocalVoiceProvider, | |
| VideoProvider, | |
| ) | |
| class RecordingImageProvider(ImageProvider): | |
| def __init__(self) -> None: | |
| self.calls: list[dict] = [] | |
| def name(self) -> str: | |
| return "recording" | |
| def supports_reference_image(self) -> bool: | |
| return True | |
| async def generate_image( | |
| self, | |
| prompt: str, | |
| character_name: str, | |
| character_details: str, | |
| aspect_ratio: str, | |
| output_path: str, | |
| reference_image_path=None, | |
| ) -> str: | |
| self.calls.append({"reference_image_path": reference_image_path}) | |
| Path(output_path).parent.mkdir(parents=True, exist_ok=True) | |
| Path(output_path).write_bytes(b"generated") | |
| return output_path | |
| class NoopVideoProvider(VideoProvider): | |
| def name(self) -> str: | |
| return "noop" | |
| async def generate_video( | |
| self, image_path: str, prompt: str, duration: float, aspect_ratio: str, output_path: str | |
| ) -> str: | |
| return output_path | |
| def make_service(image: ImageProvider) -> CompanionService: | |
| return CompanionService( | |
| chat=LocalCompanionChatProvider(), | |
| memory=LocalMemoryProvider(), | |
| image=image, | |
| video=NoopVideoProvider(), | |
| voice=LocalVoiceProvider(), | |
| settings=Settings(service_api_key="test-secret-key"), | |
| ) | |
| def params_for(service: CompanionService, output_path: str, **overrides) -> dict: | |
| base = dict( | |
| prompt="bikini", | |
| character_name="Companion", | |
| character_details="", | |
| aspect_ratio="1:1", | |
| output_path=output_path, | |
| ) | |
| base.update(overrides) | |
| return service._image_params(**base) | |
| def test_supplied_avatar_that_cannot_be_materialised_refuses_the_render(tmp_path) -> None: | |
| """The headline rule: never fall back to text-to-image when an avatar exists.""" | |
| image = RecordingImageProvider() | |
| service = make_service(image) | |
| # An avatar WAS supplied, but nothing landed on disk (undecodable payload, or | |
| # the cache was lost between attempts). | |
| params = params_for( | |
| service, str(tmp_path / "out.png"), avatar_expected=True, avatar_path="", avatar_url="" | |
| ) | |
| with pytest.raises(RuntimeError, match="reference image"): | |
| asyncio.run(service._build_image_task(params)) | |
| assert image.calls == [], "a refused render must not reach the provider at all" | |
| def test_resolved_avatar_is_handed_to_the_provider_as_the_reference(tmp_path) -> None: | |
| image = RecordingImageProvider() | |
| service = make_service(image) | |
| avatar = tmp_path / "avatar.png" | |
| Image.new("RGB", (32, 32), (10, 20, 30)).save(avatar) | |
| params = params_for( | |
| service, str(tmp_path / "out.png"), avatar_expected=True, avatar_path=str(avatar) | |
| ) | |
| asyncio.run(service._build_image_task(params)) | |
| assert image.calls, "the render must actually run" | |
| assert image.calls[-1]["reference_image_path"] == str(avatar) | |
| def test_a_companion_with_no_avatar_still_renders_from_text(tmp_path) -> None: | |
| """Unchanged behaviour: no avatar was ever asked for, so there is nothing to guard.""" | |
| image = RecordingImageProvider() | |
| service = make_service(image) | |
| params = params_for(service, str(tmp_path / "out.png")) | |
| asyncio.run(service._build_image_task(params)) | |
| assert image.calls[-1]["reference_image_path"] is None | |
| def test_avatar_intent_is_recorded_even_when_the_download_failed(tmp_path) -> None: | |
| """`avatar_expected` is the REQUEST's intent, not the download's outcome.""" | |
| service = make_service(RecordingImageProvider()) | |
| params = params_for( | |
| service, | |
| str(tmp_path / "out.png"), | |
| avatar_expected=True, | |
| avatar_path="", | |
| avatar_url="https://example.invalid/avatar.jpg", | |
| ) | |
| assert params["avatar_expected"] is True | |
| plain = params_for(service, str(tmp_path / "out.png")) | |
| assert plain["avatar_expected"] is False | |
| def test_reference_survives_the_container_losing_its_avatar_cache(tmp_path, monkeypatch) -> None: | |
| """The payload rebuilds the reference when the cached file is gone. | |
| On a Hugging Face Space the cached avatar lives on the Space's own container | |
| filesystem, so a restart/sleep/rebuild drops it while the job row survives in | |
| the durable store. `avatar_url` alone cannot recover it there: it points at the | |
| Node backend on a different host, and without `SERVER_HOST` it is | |
| `http://localhost:<port>` - which resolved INSIDE the Space addresses the Space | |
| itself. The render must still be an image edit, not a refusal. | |
| """ | |
| monkeypatch.setenv("VIDEO_STORAGE_DIR", str(tmp_path / "storage")) | |
| image = RecordingImageProvider() | |
| service = make_service(image) | |
| avatar = tmp_path / "avatar.png" | |
| Image.new("RGB", (32, 32), (10, 20, 30)).save(avatar) | |
| payload = "data:image/png;base64," + base64.b64encode(avatar.read_bytes()).decode() | |
| # A cached path that no longer exists (the container restarted) plus the | |
| # host-less URL the backend forwards when `SERVER_HOST` is unset. | |
| params = params_for( | |
| service, | |
| str(tmp_path / "out.png"), | |
| avatar_expected=True, | |
| avatar_path=str(tmp_path / "gone.png"), | |
| avatar_image=payload, | |
| avatar_url="http://localhost:4000/uploads/characters/a.png", | |
| ) | |
| asyncio.run(service._build_image_task(params)) | |
| reference = Path(image.calls[-1]["reference_image_path"]) | |
| assert reference.is_file(), "the reference must be rebuilt on disk" | |
| assert reference.read_bytes() == avatar.read_bytes() | |
| def test_cached_reference_is_still_preferred_over_the_payload(tmp_path) -> None: | |
| """The cache is an immutable, content-keyed file: the fast path is unchanged.""" | |
| image = RecordingImageProvider() | |
| service = make_service(image) | |
| cached = tmp_path / "cached.png" | |
| Image.new("RGB", (32, 32), (1, 2, 3)).save(cached) | |
| params = params_for( | |
| service, | |
| str(tmp_path / "out.png"), | |
| avatar_expected=True, | |
| avatar_path=str(cached), | |
| avatar_image="", | |
| avatar_url="", | |
| ) | |
| asyncio.run(service._build_image_task(params)) | |
| assert image.calls[-1]["reference_image_path"] == str(cached) | |
| # --- The clothing edit must cover the WHOLE outfit --------------------------- | |
| # | |
| # `restore_unmasked_pixels` pastes the reference back over every pixel the mask | |
| # leaves black, so a mask that stops short does not merely risk a partial edit - | |
| # it guarantees one. These pin the region geometry that makes a requested outfit | |
| # replacement complete, and that keeps the head out of it. | |
| def _silhouette(width: int, height: int) -> "Image.Image": | |
| """A plain person-shaped silhouette: head-free torso filling the frame.""" | |
| from PIL import Image, ImageDraw | |
| subject = Image.new("L", (width, height), 0) | |
| draw = ImageDraw.Draw(subject) | |
| draw.rectangle((width * 0.2, 0, width * 0.8, height), fill=255) | |
| return subject | |
| def test_a_clothing_edit_reaches_the_bottom_of_the_subject() -> None: | |
| """The band must not stop at a fraction of the frame height. | |
| It used to be clamped to 85% of the image, which froze everything below the | |
| waist no matter what the model painted there. | |
| """ | |
| from PIL import Image | |
| from app.services.providers.image_provider import ( | |
| _CLOTHING_REGION, | |
| build_edit_mask, | |
| ) | |
| width = height = 400 | |
| mask = build_edit_mask( | |
| Image.new("RGB", (width, height), (90, 90, 90)), | |
| _CLOTHING_REGION, | |
| subject=_silhouette(width, height), | |
| ) | |
| box = mask.getbbox() | |
| assert box is not None | |
| # The subject spans the full frame here, so the band has to reach the bottom. | |
| assert box[3] >= int(height * 0.98), f"band stops at y={box[3]} of {height}" | |
| # ...and it must be a real region, not a sliver. (The bottom edge is the | |
| # property that actually changed; the width only has to stay generous.) | |
| assert (box[2] - box[0]) > width * 0.40 | |
| def test_a_clothing_edit_never_reaches_the_head() -> None: | |
| """Whatever the geometry, the top of the frame stays frozen.""" | |
| from PIL import Image | |
| from app.services.providers.image_provider import ( | |
| _CLOTHING_REGION, | |
| build_edit_mask, | |
| ) | |
| width = height = 400 | |
| mask = build_edit_mask( | |
| Image.new("RGB", (width, height), (90, 90, 90)), | |
| _CLOTHING_REGION, | |
| subject=_silhouette(width, height), | |
| ) | |
| for fraction in (0.02, 0.10, 0.20): | |
| row = int(height * fraction) | |
| assert not any(mask.getpixel((x, row)) for x in range(0, width, 4)), ( | |
| f"the head band at y={fraction:.0%} became editable" | |
| ) | |
| def test_the_geometric_fallback_is_body_shaped_not_a_rectangle() -> None: | |
| """The stand-in silhouette must not make the whole frame editable. | |
| A rectangular stand-in made the background beside the person editable too, | |
| which is what left a visible seam around the edit. | |
| """ | |
| from PIL import Image | |
| from app.services.providers.image_provider import ( | |
| _CLOTHING_REGION, | |
| _fallback_subject, | |
| build_edit_mask, | |
| ) | |
| width = height = 400 | |
| subject = _fallback_subject(width, height) | |
| assert subject.getbbox() is not None, "the fallback silhouette must be usable" | |
| mask = build_edit_mask( | |
| Image.new("RGB", (width, height), (90, 90, 90)), | |
| _CLOTHING_REGION, | |
| subject=subject, | |
| ) | |
| box = mask.getbbox() | |
| assert box is not None | |
| # A full-width band means the stand-in was a box again. | |
| assert (box[2] - box[0]) < width * 0.9, "the editable region spans the whole frame" | |
| # ...and it still stops below the head. | |
| assert box[1] > height * 0.2, "the stand-in put the head inside the editable band" | |
| def test_the_clothing_band_reaches_the_bottom_of_a_real_portrait() -> None: | |
| """End-to-end on the geometry: a worn outfit may be replaced top to bottom.""" | |
| from PIL import Image | |
| from app.services.providers.image_provider import ( | |
| _CLOTHING_REGION, | |
| build_edit_mask, | |
| ) | |
| width = height = 400 | |
| image = Image.new("RGB", (width, height), (200, 200, 200)) | |
| # A person occupying the middle of the frame, from the head down. | |
| for y in range(0, height, 2): | |
| for x in range(int(width * 0.25), int(width * 0.75), 2): | |
| image.putpixel((x, y), (30, 40, 60)) | |
| mask = build_edit_mask(image, _CLOTHING_REGION) | |
| box = mask.getbbox() | |
| assert box is not None | |
| assert box[3] >= int(height * 0.90), f"clothing band stops at y={box[3]}" | |
| assert not mask.getpixel((width // 2, int(height * 0.05))), "the head became editable" | |
| # --- Every clothing request must be a CLOTHING edit, never a whole-person one --- | |
| # | |
| # This is the defect that produced "a newly generated person". When a garment was | |
| # not recognised, `edit_region_for_request` fell through to the SUBJECT region, | |
| # which makes the face, body, hair and background editable - so Qwen invents a | |
| # person instead of editing the companion the user chose. | |
| def test_a_clothing_request_never_becomes_a_whole_person_edit(prompt) -> None: | |
| from app.services.providers.image_provider import ( | |
| _CLOTHING_REGION, | |
| _OUTFIT_AND_SCENE_REGION, | |
| _SUBJECT_REGION, | |
| edit_region_for_request, | |
| ) | |
| region = edit_region_for_request(prompt) | |
| assert region in (_CLOTHING_REGION, _OUTFIT_AND_SCENE_REGION), ( | |
| f"{prompt!r} resolved to {region!r}: the whole person became editable, so " | |
| f"the face and body are regenerated instead of the outfit changing" | |
| ) | |
| assert region != _SUBJECT_REGION | |
| def test_an_unknown_garment_still_reads_as_a_clothing_change(prompt) -> None: | |
| """The verb frame must generalise past the hand-kept noun lists.""" | |
| from app.services.providers.image_provider import ( | |
| _CLOTHING_REGION, | |
| edit_region_for_request, | |
| ) | |
| assert edit_region_for_request(prompt) == _CLOTHING_REGION | |
| def test_the_chat_identity_note_is_not_read_as_a_clothing_request() -> None: | |
| """"Keep the same person, same clothing" must never count as a change. | |
| It contains the literal word "clothing", so a naive test reads it as an | |
| outfit request, which both protects nothing and burns prompt budget. | |
| """ | |
| from app.services.companion_prompt import is_outfit_request | |
| note = ( | |
| "Keep the same person, same clothing, same body proportions." | |
| ) | |
| assert is_outfit_request(note) is False | |
| def test_an_identity_note_that_also_names_a_garment_still_changes_it() -> None: | |
| from app.services.companion_prompt import is_outfit_request | |
| assert is_outfit_request("keep the same outfit but wear a saree") is True |