543 lines
21 KiB
Python
543 lines
21 KiB
Python
import json
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
from jsonschema import Draft202012Validator
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
import nodes as dg # noqa: E402
|
|
import scripts.dg_to_ltx_prompt_injector as prompt_injector # noqa: E402
|
|
|
|
|
|
def _packet(prompt: str, metadata: dict | None = None, segments: list[dict] | None = None) -> dict:
|
|
return {
|
|
"ltx_prompt": prompt,
|
|
"ideogram_prompt": "",
|
|
"negative_prompt": "blurry",
|
|
"scene_segments": segments if segments is not None else [{"index": 0, "duration_seconds": 2, "prompt": prompt}],
|
|
"metadata": metadata or {},
|
|
}
|
|
|
|
|
|
def _split(packet_or_text, context=None):
|
|
text = packet_or_text if isinstance(packet_or_text, str) else json.dumps(packet_or_text)
|
|
return dg.DiffusionGemmaJSONSplitter().split(text, gemma_context=context)
|
|
|
|
|
|
def _metadata(result):
|
|
return json.loads(result[4])
|
|
|
|
|
|
def _media_context(user_prompt: str, media: dict, visual_description: str = ""):
|
|
metadata = dict(media)
|
|
if visual_description:
|
|
metadata["visual_description"] = visual_description
|
|
return dg.GemmaContext(
|
|
user_prompt=user_prompt,
|
|
images=object(),
|
|
source=str(metadata.get("source", "image")),
|
|
media_metadata=metadata,
|
|
visual_description=visual_description,
|
|
warnings=[],
|
|
)
|
|
|
|
|
|
def _assert_ready(result):
|
|
assert result[8] is True, _metadata(result)
|
|
assert result[12] is True, _metadata(result)
|
|
assert result[0], "expected executable LTX prompt"
|
|
|
|
|
|
def _assert_blocked(result, expected_reason: str):
|
|
metadata = _metadata(result)
|
|
assert result[8] is False, metadata
|
|
assert result[12] is False, metadata
|
|
assert result[0] == "", metadata
|
|
assert expected_reason in metadata.get("blocked_reasons", []), metadata
|
|
|
|
|
|
def test_readiness_gates():
|
|
ready = _split(_packet("A red kite rises over a windy beach at sunset."))
|
|
_assert_ready(ready)
|
|
|
|
segmented = _split(
|
|
_packet(
|
|
"Segmented shot.",
|
|
segments=[
|
|
{"index": 0, "duration_seconds": 2, "prompt": "First beat."},
|
|
{"index": 1, "duration_seconds": 3, "prompt": "Second beat."},
|
|
],
|
|
)
|
|
)
|
|
assert segmented[7] == "2, 3", segmented[7]
|
|
|
|
salvaged = _split("ltx_prompt: A clean plain text prompt that was not valid JSON.")
|
|
_assert_blocked(salvaged, "salvaged_output")
|
|
|
|
fallback = _split("", context=dg.GemmaContext(user_prompt="make a cinematic shot", media_metadata={}))
|
|
_assert_blocked(fallback, "template_fallback")
|
|
|
|
metadata_only_context = _media_context(
|
|
"Use this image as an LTX prompt.",
|
|
{
|
|
"source": "image",
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": False,
|
|
},
|
|
)
|
|
metadata_only_claim = _split(_packet("A sunlit daylight courtyard with people in the background."), metadata_only_context)
|
|
_assert_blocked(metadata_only_claim, "metadata_only_media_reference_request")
|
|
|
|
night_context = _media_context(
|
|
"Use the supplied visual description.",
|
|
{
|
|
"source": "image",
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": False,
|
|
},
|
|
"Nighttime locked-off static camera view under neon streetlights.",
|
|
)
|
|
model_media = {
|
|
"media": {
|
|
"source": "image",
|
|
"pixels_sent_to_backend": True,
|
|
"visual_description": "daylight exterior",
|
|
}
|
|
}
|
|
night_ready = _split(_packet("Nighttime neon street scene from a locked-off static camera.", model_media), night_context)
|
|
_assert_ready(night_ready)
|
|
night_metadata = _metadata(night_ready)
|
|
assert "Nighttime" in night_metadata["media"]["visual_description"], night_metadata
|
|
assert "model_reported_media" in night_metadata, night_metadata
|
|
|
|
daylight_contradiction = _split(_packet("A bright daylight street with a blue sky.", model_media), night_context)
|
|
_assert_blocked(daylight_contradiction, "visual_description_contradiction_daylight_claim")
|
|
|
|
camera_contradiction = _split(_packet("The camera slowly zooms in on the neon street.", model_media), night_context)
|
|
_assert_blocked(camera_contradiction, "visual_description_contradiction_camera_motion")
|
|
|
|
|
|
def test_prompt_packet_schema():
|
|
schema = json.loads((ROOT / "schemas" / "prompt_packet.schema.json").read_text(encoding="utf-8"))
|
|
validator = Draft202012Validator(schema)
|
|
valid_packet = _packet("A clean prompt.", {"ready_for_generation": True, "blocked_reasons": []})
|
|
validator.validate(valid_packet)
|
|
invalid_packet = {"ltx_prompt": "missing required fields"}
|
|
errors = sorted(validator.iter_errors(invalid_packet), key=lambda item: item.path)
|
|
assert errors, "schema should reject incomplete prompt packets"
|
|
|
|
|
|
def test_ltx_camera_choreography_contract():
|
|
prompt = dg._build_model_prompt(
|
|
"Recreate the source video as an LTX prompt.",
|
|
dg.DEFAULT_MASTER_PROMPT,
|
|
"ltx",
|
|
{
|
|
"source": "video",
|
|
"duration_seconds": 10,
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": True,
|
|
"visual_grounding_mode": "pixels",
|
|
},
|
|
audio_mode="visual_only",
|
|
creativity_mode="faithful",
|
|
creative_strength=0.2,
|
|
thinking_mode="on",
|
|
)
|
|
lowered = prompt.lower()
|
|
for phrase in (
|
|
"camera choreography",
|
|
"camera-choreography extraction",
|
|
"first-frame framing",
|
|
"camera-subject relationship",
|
|
"zoom/dolly/truck/boom/pan/tilt/roll/orbit/rotation",
|
|
"parallax",
|
|
"180 degrees",
|
|
"do not invent one",
|
|
):
|
|
assert phrase in lowered, phrase
|
|
|
|
|
|
def test_image_identity_video_control_contract():
|
|
media = {
|
|
"source": "image+video",
|
|
"duration_seconds": 3,
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": True,
|
|
"visual_grounding_mode": "pixels",
|
|
"media_synthesis_mode": "image_identity_video_control",
|
|
"image_role": "identity_reference",
|
|
"video_role": "control_structure_pose_depth_canny_composition_motion_camera",
|
|
"reference_image_count": 1,
|
|
"video_sampled_frame_count": 3,
|
|
"sampled_frame_count": 4,
|
|
}
|
|
prompt = dg._build_model_prompt(
|
|
"Replace the woman in the video with the girl in the image.",
|
|
dg.DEFAULT_MASTER_PROMPT,
|
|
"ltx",
|
|
media,
|
|
audio_mode="visual_only",
|
|
creativity_mode="faithful",
|
|
creative_strength=0.2,
|
|
thinking_mode="on",
|
|
)
|
|
lowered = prompt.lower()
|
|
for phrase in (
|
|
"media_synthesis_mode",
|
|
"image_identity_video_control",
|
|
"reference-image subject performing the video action/control structure",
|
|
"image controls identity and appearance",
|
|
"wardrobe, accessories, styling",
|
|
"video controls pose, depth, canny/edge layout",
|
|
"do not copy the control-video subject's clothing",
|
|
"replace the video subject appearance with the reference-image appearance",
|
|
"do not let the person in the control video override",
|
|
):
|
|
assert phrase in lowered, phrase
|
|
|
|
intro = dg._sampled_video_frame_sequence_intro(
|
|
{
|
|
**media,
|
|
"sampled_indices": [0, 12, 24],
|
|
"source_fps": 24,
|
|
},
|
|
4,
|
|
)
|
|
intro_lower = intro.lower()
|
|
assert "optional still reference" not in intro_lower, intro
|
|
assert "first attached image is the identity reference" in intro_lower, intro
|
|
assert "remaining attached images are ordered sampled video control frames" in intro_lower, intro
|
|
assert "video frame times: 1=0.00s, 2=0.50s, 3=1.00s" in intro_lower, intro
|
|
assert "wardrobe, accessories, styling" in intro_lower, intro
|
|
assert "do not copy the control-video subject's clothing" in intro_lower, intro
|
|
|
|
synthesis_context = _media_context(
|
|
"Replace the woman in the video with the girl in the image.",
|
|
media,
|
|
)
|
|
ready = _split(
|
|
_packet("The girl from the reference image follows the video pose, room blocking, and camera motion."),
|
|
synthesis_context,
|
|
)
|
|
_assert_ready(ready)
|
|
ready_metadata = _metadata(ready)
|
|
assert ready_metadata["media_synthesis_mode"] == "image_identity_video_control", ready_metadata
|
|
assert ready_metadata["image_role"] == "identity_reference", ready_metadata
|
|
assert ready_metadata["video_sampled_frame_count"] == 3, ready_metadata
|
|
|
|
missing_image = _media_context(
|
|
"Replace the woman in the video with the girl in the image.",
|
|
{
|
|
"source": "video",
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": True,
|
|
"media_synthesis_mode": "image_identity_video_control",
|
|
"reference_image_count": 0,
|
|
"video_sampled_frame_count": 3,
|
|
"sampled_frame_count": 3,
|
|
},
|
|
)
|
|
_assert_blocked(
|
|
_split(_packet("The reference-image subject follows the control-video pose and camera motion."), missing_image),
|
|
"synthesis_missing_identity_image",
|
|
)
|
|
|
|
unattached_image = _media_context(
|
|
"Replace the woman in the video with the girl in the image.",
|
|
{
|
|
"source": "video",
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": True,
|
|
"media_synthesis_mode": "image_identity_video_control",
|
|
"reference_image_count": 1,
|
|
"reference_image_backend_attached": False,
|
|
"video_sampled_frame_count": 3,
|
|
"sampled_frame_count": 3,
|
|
},
|
|
)
|
|
_assert_blocked(
|
|
_split(_packet("The reference-image subject follows the control-video pose and camera motion."), unattached_image),
|
|
"synthesis_identity_image_not_attached",
|
|
)
|
|
|
|
missing_video = _media_context(
|
|
"Replace the woman in the video with the girl in the image.",
|
|
{
|
|
"source": "image",
|
|
"pixel_tensor_present": True,
|
|
"pixels_sent_to_backend": True,
|
|
"media_synthesis_mode": "image_identity_video_control",
|
|
"reference_image_count": 1,
|
|
"video_sampled_frame_count": 0,
|
|
"sampled_frame_count": 1,
|
|
},
|
|
)
|
|
_assert_blocked(
|
|
_split(_packet("The reference-image subject follows the control-video pose and camera motion."), missing_video),
|
|
"synthesis_missing_control_video",
|
|
)
|
|
|
|
|
|
def test_media_sampler_forces_identity_image_in_synthesis_mode():
|
|
import torch
|
|
|
|
class FakeVideo:
|
|
def __init__(self, images, frame_rate=3.0):
|
|
self.images = images
|
|
self.frame_rate = frame_rate
|
|
|
|
def get_duration(self):
|
|
return 1.0
|
|
|
|
def get_frame_count(self):
|
|
return int(self.images.shape[0])
|
|
|
|
def get_dimensions(self):
|
|
return (int(self.images.shape[2]), int(self.images.shape[1]))
|
|
|
|
def get_components(self):
|
|
return SimpleNamespace(images=self.images, frame_rate=self.frame_rate)
|
|
|
|
image = torch.ones((1, 2, 2, 3))
|
|
media_context, metadata_json = dg.DiffusionGemmaMediaSampler().sample(
|
|
3.0,
|
|
10.0,
|
|
10,
|
|
"trim",
|
|
False,
|
|
image=image,
|
|
video=FakeVideo(torch.zeros((3, 2, 2, 3))),
|
|
media_synthesis_mode="image_identity_video_control",
|
|
)
|
|
metadata = json.loads(metadata_json)
|
|
assert media_context.source == "image+video", metadata
|
|
assert metadata["media_synthesis_mode"] == "image_identity_video_control", metadata
|
|
assert metadata["reference_image_count"] == 1, metadata
|
|
assert metadata["reference_image_backend_attached"] is True, metadata
|
|
assert metadata["reference_image_resized_for_video_control"] is False, metadata
|
|
assert metadata["video_sampled_frame_count"] == 3, metadata
|
|
assert metadata["sampled_frame_count"] == 4, metadata
|
|
assert tuple(media_context.images.shape) == (4, 2, 2, 3)
|
|
|
|
mismatched_image = torch.ones((1, 6, 8, 3))
|
|
mismatched_context, mismatched_metadata_json = dg.DiffusionGemmaMediaSampler().sample(
|
|
3.0,
|
|
10.0,
|
|
10,
|
|
"trim",
|
|
False,
|
|
image=mismatched_image,
|
|
video=FakeVideo(torch.zeros((3, 2, 4, 3))),
|
|
media_synthesis_mode="image_identity_video_control",
|
|
)
|
|
mismatched_metadata = json.loads(mismatched_metadata_json)
|
|
assert mismatched_context.source == "image+video", mismatched_metadata
|
|
assert mismatched_metadata["reference_image_backend_attached"] is True, mismatched_metadata
|
|
assert mismatched_metadata["reference_image_resized_for_video_control"] is True, mismatched_metadata
|
|
assert mismatched_metadata["reference_image_original_width"] == 8, mismatched_metadata
|
|
assert mismatched_metadata["reference_image_original_height"] == 6, mismatched_metadata
|
|
assert mismatched_metadata["reference_image_backend_width"] == 4, mismatched_metadata
|
|
assert mismatched_metadata["reference_image_backend_height"] == 2, mismatched_metadata
|
|
assert mismatched_metadata["sampled_frame_count"] == 4, mismatched_metadata
|
|
assert tuple(mismatched_context.images.shape) == (4, 2, 4, 3)
|
|
assert bool(torch.all(mismatched_context.images[0] > 0.99)), mismatched_context.images[0]
|
|
|
|
|
|
def test_transformers_media_message_labels_identity_and_control_frames():
|
|
import torch
|
|
|
|
media_context = dg.MediaContext(
|
|
images=torch.zeros((4, 2, 2, 3)),
|
|
source="image+video",
|
|
metadata={
|
|
"source": "image+video",
|
|
"media_synthesis_mode": "image_identity_video_control",
|
|
"reference_image_count": 1,
|
|
"video_sampled_frame_count": 3,
|
|
"sampled_frame_count": 4,
|
|
"sampled_indices": [0, 12, 24],
|
|
"source_fps": 24,
|
|
"transformers_video_transport": "sampled_frame_images",
|
|
},
|
|
)
|
|
messages, processor_kwargs = dg._transformers_messages_and_processor_kwargs("Final prompt contract.", media_context)
|
|
assert processor_kwargs == {}, processor_kwargs
|
|
content = messages[0]["content"]
|
|
texts = "\n".join(str(item["text"]) for item in content if item.get("type") == "text").lower()
|
|
assert "identity reference image" in texts, texts
|
|
assert "video control frames" in texts, texts
|
|
assert "do not copy the control-video subject's clothing" in texts, texts
|
|
assert "reference wardrobe and accessories" in texts, texts
|
|
assert "not part of the video timeline" in texts, texts
|
|
assert sum(1 for item in content if item.get("type") == "image") == 4, content
|
|
|
|
|
|
def _minimal_ltx_workflow() -> dict:
|
|
def subgraph(subgraph_id: str, name: str, nodes: list[dict]) -> dict:
|
|
return {"id": subgraph_id, "name": name, "nodes": nodes, "links": []}
|
|
|
|
return {
|
|
"id": "smoke",
|
|
"version": 0.4,
|
|
"nodes": [
|
|
{
|
|
"id": 5496,
|
|
"type": "VHS_LoadVideo",
|
|
"title": "VHS_LoadVideo",
|
|
"widgets_values": {"video": "old_video.mp4", "force_rate": 24},
|
|
}
|
|
],
|
|
"links": [],
|
|
"definitions": {
|
|
"subgraphs": [
|
|
subgraph(
|
|
"fe45e814-d005-4dd9-97dc-e31c4c6e9773",
|
|
"I2V Image",
|
|
[
|
|
{
|
|
"id": 5038,
|
|
"type": "LoadImage",
|
|
"title": "Reference Object",
|
|
"widgets_values": ["old_reference.png", "image"],
|
|
}
|
|
],
|
|
),
|
|
subgraph(
|
|
"ba181d97-b80f-49a8-ae30-e552bd7c8f49",
|
|
"Mask",
|
|
[
|
|
{
|
|
"id": 5433,
|
|
"type": "SAM3Segment",
|
|
"title": "SAM3 Segmentation (RMBG) - 1 Pass",
|
|
"widgets_values": ["old target", "Merged", 0.5],
|
|
}
|
|
],
|
|
),
|
|
subgraph(
|
|
"f935d44a-3a0d-41cf-b9bc-b2261f8ba60f",
|
|
"New Subgraph",
|
|
[
|
|
{
|
|
"id": 5676,
|
|
"type": "SAM3Segment",
|
|
"title": "SAM3 Segmentation (RMBG) - 2 Pass",
|
|
"widgets_values": ["", "Merged", 0.5],
|
|
}
|
|
],
|
|
),
|
|
subgraph(
|
|
"a0619931-10e4-48af-9d29-114a59e7c821",
|
|
"New Subgraph",
|
|
[
|
|
{
|
|
"id": 5745,
|
|
"type": "SAM3Segment",
|
|
"title": "SAM3 Segmentation (RMBG) - 2 Pass",
|
|
"widgets_values": ["", "Merged", 0.5],
|
|
}
|
|
],
|
|
),
|
|
subgraph(
|
|
"483656b1-20e1-4b50-a7e2-3b6126ad4bb7",
|
|
"Prompt",
|
|
[
|
|
{
|
|
"id": 5630,
|
|
"type": "CLIPTextEncode",
|
|
"title": "Manual Prompt",
|
|
"widgets_values": ["old positive"],
|
|
},
|
|
{
|
|
"id": 5626,
|
|
"type": "CLIPTextEncode",
|
|
"title": "CLIPTextEncode",
|
|
"widgets_values": [""],
|
|
},
|
|
],
|
|
),
|
|
]
|
|
},
|
|
}
|
|
|
|
|
|
def test_dg_to_ltx_prompt_injector():
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tmp = Path(tmpdir)
|
|
workflow_path = tmp / "workflow.json"
|
|
dg_path = tmp / "dg_prompt.json"
|
|
output_path = tmp / "patched.json"
|
|
manifest_path = tmp / "manifest.json"
|
|
workflow_path.write_text(json.dumps(_minimal_ltx_workflow()), encoding="utf-8")
|
|
dg_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"sam_target": "scratched panel",
|
|
"positive_prompt": "Repair only the small scratch on the panel.",
|
|
"negative_prompt": "full scene reroll, new object, identity change",
|
|
"preservation_constraints": ["preserve camera", "preserve lighting"],
|
|
"do_not_change": ["background", "panel shape"],
|
|
"repair_scope": "localized_masked_inpaint",
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
manifest = prompt_injector.build_patched_workflow(
|
|
workflow_path,
|
|
dg_path,
|
|
output_path,
|
|
manifest_path,
|
|
video_path="input/source.mp4",
|
|
reference_image_path="input/reference.png",
|
|
)
|
|
patched = json.loads(output_path.read_text(encoding="utf-8"))
|
|
assert manifest["only_prompt_path_widgets_changed"] is True
|
|
assert manifest["ltx_graph_topology_unchanged"] is True
|
|
assert manifest["sam3_remains_in_workflow"] is True
|
|
assert manifest["sam3_node_count"] == 3
|
|
assert manifest["changed_widget_count"] == 7
|
|
assert manifest["generation_run"] is False
|
|
assert manifest["semantic_repair_success_claimed"] is False
|
|
assert not manifest["forbidden_changes"]
|
|
changed_ids = {item["field_id"] for item in manifest["changes"]}
|
|
assert changed_ids == {
|
|
"video_input",
|
|
"reference_image_input",
|
|
"sam3_target_primary",
|
|
"sam3_target_pass2_a",
|
|
"sam3_target_pass2_b",
|
|
"ltx_positive_prompt",
|
|
"ltx_negative_prompt",
|
|
}
|
|
assert patched["nodes"][0]["widgets_values"]["video"] == "input/source.mp4"
|
|
prompt_sg = next(item for item in patched["definitions"]["subgraphs"] if item["id"] == "483656b1-20e1-4b50-a7e2-3b6126ad4bb7")
|
|
positive = next(item for item in prompt_sg["nodes"] if item["id"] == 5630)["widgets_values"][0]
|
|
assert "Repair only the small scratch" in positive
|
|
assert "Preservation constraints" in positive
|
|
assert "Do not change" in positive
|
|
assert next(item for item in prompt_sg["nodes"] if item["id"] == 5626)["widgets_values"][0] == "full scene reroll, new object, identity change"
|
|
inventory = prompt_injector.inspect_workflow(output_path)["field_inventory"]
|
|
assert len(inventory) == 7
|
|
|
|
|
|
def main():
|
|
test_readiness_gates()
|
|
test_prompt_packet_schema()
|
|
test_ltx_camera_choreography_contract()
|
|
test_image_identity_video_control_contract()
|
|
test_media_sampler_forces_identity_image_in_synthesis_mode()
|
|
test_transformers_media_message_labels_identity_and_control_frames()
|
|
test_dg_to_ltx_prompt_injector()
|
|
print("Prompt Builder readiness/schema smoke passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|