import pytest import torch from fastfill.v2.batch import TinyTokenizer, collate_samples from fastfill.v2.model import ModelConfig, build_model, load_model, model_inputs from fastfill.v2.tests.test_batch import sample def config(): return ModelConfig(backbone="tiny", decoder_dim=32, decoder_heads=4, decoder_layers=2, tiny_hidden_size=32, max_objects=16, tiny_layers=1, dropout=0.0) @pytest.mark.parametrize("kwargs", [ {"decoder_dim": True}, {"decoder_layers": 1.5}, {"decoder_ffn_multiplier": 0}, {"dropout": float("nan")}, {"lora_dropout": float("inf")}, {"lora_rank": -1}, {"train_backbone": "false", "lora_rank": 0}, {"local_files_only": "false"}, {"size_reference": (True, 1., 1.)}, {"size_reference": (1e38, 1., 1.)}, {"size_reference": (1e-42, 1., 1.)}, {"backbone": ""}, {"lora_target_modules": [None]}, ]) def test_model_config_rejects_ambiguous_types_and_nonfinite_geometry(kwargs): with pytest.raises(ValueError): ModelConfig(**kwargs) def test_positive_size_bounds_survive_fp16_autocast(): cfg = ModelConfig(backbone="tiny", decoder_dim=32, decoder_heads=4, decoder_layers=1, tiny_hidden_size=32, size_log_limit=30.) model = build_model(cfg) with torch.no_grad(): model.size_head.weight.zero_() model.size_head.bias.copy_(torch.tensor([30., -30., 0.])) batch = collate_samples([sample(1)], TinyTokenizer()) with torch.autocast("cpu", dtype=torch.float16): sizes = model(**model_inputs(batch))["size"] assert sizes.dtype == torch.float32 assert torch.isfinite(sizes).all() and (sizes > 0).all() sizes.log().sum().backward() assert torch.isfinite(model.size_head.bias.grad).all() def test_exact_slots_positive_size_and_gradients_to_condition_backbone(): model = build_model(config()) b = collate_samples([sample(2), sample(1)], TinyTokenizer()) out = model(**model_inputs(b)) assert out["position_normalized"].shape == (2, 2, 3) assert out["yaw_logits"].shape == (2, 2, 12) assert out["slot_mask"].tolist() == [[True, True], [True, False]] assert (out["size"][out["slot_mask"]] > 0).all() assert torch.equal(out["position_normalized"][1, 1], torch.zeros(3)) loss = sum(out[k].square().sum() for k in ("position_normalized", "size", "yaw_logits", "yaw_residuals")) loss.backward() assert model.backbone.embedding.weight.grad.abs().sum() > 0 assert model.position_head.weight.grad.abs().sum() > 0 def test_batch_padding_does_not_change_valid_object_predictions(): torch.manual_seed(2) model = build_model(config()).eval() one = collate_samples([sample(1)], TinyTokenizer()) padded = collate_samples([sample(1), sample(3)], TinyTokenizer()) with torch.no_grad(): a = model(**model_inputs(one)) b = model(**model_inputs(padded)) for key in ("position_normalized", "size", "yaw_logits", "yaw_residuals"): assert torch.allclose(a[key][0, 0], b[key][0, 0], atol=2e-6), key def test_full_condition_after_object_changes_predictions(): model = build_model(config()).eval() a = sample(1) b = sample(1) b["condition"]["constraints"] = [{"type": "faces_direction", "object_id": "chair_0", "direction_xy": [0, -1]}] with torch.no_grad(): x = model(**model_inputs(collate_samples([a], TinyTokenizer()))) y = model(**model_inputs(collate_samples([b], TinyTokenizer()))) assert not torch.allclose(x["position_normalized"], y["position_normalized"]) def test_fixed_size_and_ground_height_override_head(): s = sample(1) s["condition"]["objects"][0].update(fixed_size_local_m=[1.1, .7, 1.2], support_parent="floor") s["target"]["objects"][0]["target_size_local_m"] = [1.1, .7, 1.2] model = build_model(config()) b = collate_samples([s], TinyTokenizer()) out = model(**model_inputs(b)) assert torch.equal(out["size"][0, 0], torch.tensor([1.1, .7, 1.2])) assert out["position_normalized"][0, 0, 2] == 0 def test_checkpoint_round_trip(tmp_path): model = build_model(config()).eval() b = collate_samples([sample(1)], TinyTokenizer()) model.save_pretrained(tmp_path) loaded = load_model(tmp_path).eval() with torch.no_grad(): assert torch.equal(model(**model_inputs(b))["size"], loaded(**model_inputs(b))["size"]) def test_empty_scene_and_mixed_empty_batch_do_not_create_nan(): model = build_model(config()) empty = collate_samples([sample(0)], TinyTokenizer()) out = model(**model_inputs(empty)) assert out["size"].shape == (1, 0, 3) mixed = collate_samples([sample(0), sample(2)], TinyTokenizer()) out = model(**model_inputs(mixed)) assert torch.isfinite(out["size"]).all() assert torch.equal(out["size"][0], torch.zeros(2, 3)) def test_bidirectional_object_interaction_changes_first_slot(): torch.manual_seed(7) model = build_model(config()).eval() a = sample(2) b = sample(2) b["condition"]["objects"][1]["description"] = "large reading chair with an attached lamp" with torch.no_grad(): x = model(**model_inputs(collate_samples([a], TinyTokenizer()))) y = model(**model_inputs(collate_samples([b], TinyTokenizer()))) assert not torch.allclose(x["size"][0, 0], y["size"][0, 0]) @pytest.mark.parametrize("training", ["lora", "full", "frozen"]) def test_local_random_qwen_api_gradients_and_checkpoint(tmp_path, training): transformers = pytest.importorskip("transformers") if training == "lora": pytest.importorskip("peft") source = tmp_path / "qwen" transformers.Qwen2Model(transformers.Qwen2Config(vocab_size=258, hidden_size=32, intermediate_size=64, num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=4096)).save_pretrained(source) cfg = ModelConfig(backbone=str(source), decoder_dim=32, decoder_heads=4, decoder_layers=1, lora_rank=2 if training == "lora" else 0, lora_alpha=4, train_backbone=training == "full", local_files_only=True) model = build_model(cfg) b = collate_samples([sample(1)], TinyTokenizer()) model(**model_inputs(b))["size"].sum().backward() gradients = [p.grad for p in model.backbone.parameters() if p.grad is not None] assert bool(gradients) == (training != "frozen") if gradients: assert sum(float(grad.abs().sum()) for grad in gradients) > 0 checkpoint = tmp_path / "checkpoint" model.eval().save_pretrained(checkpoint) loaded = load_model(checkpoint, local_files_only=True).eval() with torch.no_grad(): assert torch.equal(model(**model_inputs(b))["size"], loaded(**model_inputs(b))["size"])