devatlas-ai-gpu / tests /test_reference_required.py
YashKotak's picture
fix companion reference image resolution
6bb8992
Raw History Blame Contribute Delete
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] = []
@property
def name(self) -> str:
return "recording"
@property
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):
@property
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.
@pytest.mark.parametrize(
"prompt",
[
"wear a saree",
"wear a dress",
"wear a top",
"wear a bikini",
"wear a shirt",
"wear a jacket",
"wear a kimono",
"wear a lehenga",
"put on a dress",
"change into a blue top",
"a red saree",
"dress her in a saree",
],
)
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
@pytest.mark.parametrize(
"prompt",
["wear an entirely unknown garment word", "put on some unheard-of thing"],
)
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