from pathlib import Path from tempfile import TemporaryDirectory from unittest import TestCase from catbox.model_backend import CatboxModelBackend, FakeModelRunner from catbox.sd_turbo_runner import SdTurboImageToImageModelRunner class FailingModelRunner(FakeModelRunner): def generate(self, outcome, seed, config=None): raise RuntimeError("CUDA ran out of memory") class ModelBackendTests(TestCase): def test_readiness_reports_starting_until_runner_is_ready(self): with TemporaryDirectory() as output_dir: backend = CatboxModelBackend( model_runner=FakeModelRunner(output_dir=output_dir, ready=False) ) response = backend.readiness() self.assertEqual(response["status"], "starting") self.assertEqual(response["modelBackend"], "starting") def test_observation_returns_generated_outcome_contract(self): with TemporaryDirectory() as output_dir: backend = CatboxModelBackend( model_runner=FakeModelRunner(output_dir=output_dir), seed_source=lambda: 41100, clock=lambda: 12.5, ) response = backend.observe() self.assertEqual(response["status"], "generated") self.assertIn(response["outcome"], {"living", "dead"}) self.assertTrue(response["imageRef"].endswith(".png")) self.assertTrue(Path(response["imageRef"]).exists()) self.assertEqual(response["metadata"]["seed"], 41100) self.assertIs(response["metadata"]["ephemeral"], True) self.assertTrue(response["revealNote"]) def test_generation_failure_is_structured_without_fake_outcome(self): with TemporaryDirectory() as output_dir: backend = CatboxModelBackend( model_runner=FailingModelRunner(output_dir=output_dir), seed_source=lambda: 41100, outcome_source=lambda: "living", ) response = backend.observe() self.assertEqual(response["status"], "generation_failed") self.assertEqual(response["error"]["type"], "RuntimeError") self.assertIn("CUDA", response["error"]["message"]) self.assertNotIn("imageRef", response) self.assertEqual(response["metadata"]["outcome"], "living") self.assertIs(response["metadata"]["ephemeral"], True) def test_invalid_outcome_selection_fails_before_generation(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "ghost", ) response = backend.observe() self.assertEqual(response["status"], "generation_failed") self.assertEqual(response["error"]["type"], "InvalidOutcome") self.assertNotIn("imageRef", response) def test_observation_invokes_runner_once_for_selected_outcome(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "dead", ) response = backend.observe() self.assertEqual(response["status"], "generated") self.assertEqual(runner.generations, [{"outcome": "dead", "seed": 41100}]) def test_dev_controls_can_force_living_outcome(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "dead", ) response = backend.observe_with_dev_controls({"outcome": "living"}) self.assertEqual(response["status"], "generated") self.assertEqual(response["outcome"], "living") self.assertEqual(runner.generations, [{"outcome": "living", "seed": 41100}]) def test_dev_controls_can_force_dead_outcome(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "living", ) response = backend.observe_with_dev_controls({"outcome": "dead"}) self.assertEqual(response["status"], "generated") self.assertEqual(response["outcome"], "dead") self.assertEqual(runner.generations, [{"outcome": "dead", "seed": 41100}]) def test_dev_controls_can_override_seed_for_reproducible_observation(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "living", ) response = backend.observe_with_dev_controls({"seed": 90210}) self.assertEqual(response["status"], "generated") self.assertEqual(response["metadata"]["seed"], 90210) self.assertEqual(runner.generations, [{"outcome": "living", "seed": 90210}]) def test_dev_controls_pass_config_overrides_to_generation_metadata_and_runner(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "dead", ) response = backend.observe_with_dev_controls( {"config": {"steps": 6, "strength": 0.55}} ) self.assertEqual(response["status"], "generated") self.assertEqual( response["metadata"]["configOverrides"], {"steps": 6, "strength": 0.55}, ) self.assertEqual( runner.generations, [ { "outcome": "dead", "seed": 41100, "config": {"steps": 6, "strength": 0.55}, } ], ) def test_invalid_dev_controls_outcome_fails_clearly_without_generation(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, ) response = backend.observe_with_dev_controls({"outcome": "ghost"}) self.assertEqual(response["status"], "generation_failed") self.assertEqual(response["error"]["type"], "InvalidDevControlsOverride") self.assertEqual(response["error"]["field"], "outcome") self.assertIn("ghost", response["error"]["message"]) self.assertNotIn("imageRef", response) self.assertEqual(runner.generations, []) def test_invalid_dev_controls_seed_fails_clearly_without_generation(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, outcome_source=lambda: "living", ) response = backend.observe_with_dev_controls({"seed": "not-a-seed"}) self.assertEqual(response["status"], "generation_failed") self.assertEqual(response["error"]["type"], "InvalidDevControlsOverride") self.assertEqual(response["error"]["field"], "seed") self.assertIn("integer", response["error"]["message"]) self.assertNotIn("imageRef", response) self.assertEqual(runner.generations, []) def test_invalid_dev_controls_config_fails_clearly_without_generation(self): with TemporaryDirectory() as output_dir: runner = FakeModelRunner(output_dir=output_dir) backend = CatboxModelBackend( model_runner=runner, outcome_source=lambda: "living", ) response = backend.observe_with_dev_controls({"config": ["steps", 6]}) self.assertEqual(response["status"], "generation_failed") self.assertEqual(response["error"]["type"], "InvalidDevControlsOverride") self.assertEqual(response["error"]["field"], "config") self.assertIn("object", response["error"]["message"]) self.assertNotIn("imageRef", response) self.assertEqual(runner.generations, []) class StubGeneratedImage: def __init__(self) -> None: self.saved_path = None def save(self, path): self.saved_path = path Path(path).write_bytes(b"generated image") class StubPipeline: def __init__(self) -> None: self.calls = [] self.generated_image = StubGeneratedImage() def to(self, device): self.device = device return self def enable_attention_slicing(self): self.attention_slicing_enabled = True def __call__(self, **kwargs): self.calls.append(kwargs) return type("PipelineResult", (), {"images": [self.generated_image]})() class StubTraceLatents: def detach(self): return self def __truediv__(self, value): return self class StubTraceVae: config = type("VaeConfig", (), {"scaling_factor": 1.0})() def decode(self, latents): return type("Decoded", (), {"sample": "decoded image"})() class StubTraceImageProcessor: def postprocess(self, decoded, output_type): return [StubGeneratedImage()] class StubTracePipeline(StubPipeline): def __init__(self) -> None: super().__init__() self.vae = StubTraceVae() self.image_processor = StubTraceImageProcessor() def __call__(self, **kwargs): self.calls.append(kwargs) callback = kwargs.get("callback_on_step_end") if callback is not None: callback(self, 0, "timestep", {"latents": StubTraceLatents()}) return type("PipelineResult", (), {"images": [self.generated_image]})() class StubTorch: float16 = "float16" float32 = "float32" class cuda: @staticmethod def is_available(): return False class Generator: def __init__(self, device): self.device = device self.seed = None def manual_seed(self, seed): self.seed = seed return self class no_grad: def __enter__(self): return None def __exit__(self, exc_type, exc, traceback): return False class SdTurboRunnerTests(TestCase): def test_real_runner_generates_selected_outcome_file_with_prompt_and_metadata(self): with TemporaryDirectory() as runtime_dir: pipeline = StubPipeline() loaded = [] box_image_loads = [] def load_pipeline(model_id, **kwargs): loaded.append({"model_id": model_id, "kwargs": kwargs}) return pipeline def load_box_image(path, config): box_image_loads.append({"width": config.width, "height": config.height}) return "box image" runner = SdTurboImageToImageModelRunner( runtime_dir=runtime_dir, pipeline_loader=load_pipeline, box_image_loader=load_box_image, torch_module=StubTorch, now=lambda: "2026-06-02T12:00:00Z", timer=iter([10.0, 10.25]).__next__, ) generated = runner.generate("living", seed=41100) self.assertTrue(runner.is_ready()) self.assertEqual(len(loaded), 1) self.assertEqual(loaded[0]["model_id"], "stabilityai/sd-turbo") self.assertEqual(generated["generation_seconds"], 0.25) self.assertTrue(Path(generated["image_ref"]).exists()) self.assertIn("large clearly visible living cat", pipeline.calls[0]["prompt"]) self.assertEqual(pipeline.calls[0]["num_inference_steps"], 6) self.assertEqual(pipeline.calls[0]["strength"], 0.78) self.assertEqual(box_image_loads, [{"width": 512, "height": 512}]) self.assertNotIn("negative_prompt", pipeline.calls[0]) self.assertEqual(generated["metadata"]["runner"], "sd_turbo_img2img") self.assertEqual(generated["metadata"]["device"], "cpu") def test_backend_response_includes_real_runner_metadata(self): with TemporaryDirectory() as runtime_dir: pipeline = StubPipeline() box_image_loads = [] runner = SdTurboImageToImageModelRunner( runtime_dir=runtime_dir, pipeline_loader=lambda model_id, **kwargs: pipeline, box_image_loader=lambda path, config: box_image_loads.append( {"width": config.width, "height": config.height} ) or "box image", torch_module=StubTorch, now=lambda: "2026-06-02T12:00:00Z", timer=iter([20.0, 20.5]).__next__, ) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "dead", clock=lambda: 123.0, ) response = backend.observe() self.assertEqual(response["status"], "generated") self.assertEqual(response["outcome"], "dead") self.assertTrue(Path(response["imageRef"]).exists()) self.assertEqual(response["metadata"]["seed"], 41100) self.assertEqual(response["metadata"]["generationSeconds"], 0.5) self.assertEqual(response["metadata"]["runner"], "sd_turbo_img2img") self.assertEqual(response["metadata"]["device"], "cpu") self.assertIn("large clearly visible motionless cat", pipeline.calls[0]["prompt"]) self.assertIn("no blood", pipeline.calls[0]["prompt"]) self.assertIn("no gore", pipeline.calls[0]["prompt"]) self.assertEqual(pipeline.calls[0]["num_inference_steps"], 6) self.assertEqual(pipeline.calls[0]["strength"], 0.7) self.assertEqual(box_image_loads, [{"width": 512, "height": 512}]) def test_real_runner_captures_selected_branch_trace_frames(self): with TemporaryDirectory() as runtime_dir: pipeline = StubTracePipeline() observed_trace_refs = [] runner = SdTurboImageToImageModelRunner( runtime_dir=runtime_dir, pipeline_loader=lambda model_id, **kwargs: pipeline, box_image_loader=lambda path, config: "box image", torch_module=StubTorch, now=lambda: "2026-06-02T12:00:00Z", timer=iter([22.0, 22.4]).__next__, ) generated = runner.generate( "living", seed=41100, trace_callback=observed_trace_refs.append, ) self.assertEqual(len(generated["trace_refs"]), 1) self.assertEqual(observed_trace_refs, generated["trace_refs"]) self.assertTrue(Path(generated["trace_refs"][0]).exists()) self.assertIn("callback_on_step_end", pipeline.calls[0]) self.assertEqual( pipeline.calls[0]["callback_on_step_end_tensor_inputs"], ["latents"], ) def test_backend_reports_starting_and_failure_when_real_runner_preload_fails(self): with TemporaryDirectory() as runtime_dir: def fail_to_load(model_id, **kwargs): raise RuntimeError("model weights unavailable") runner = SdTurboImageToImageModelRunner( runtime_dir=runtime_dir, pipeline_loader=fail_to_load, box_image_loader=lambda path, config: "box image", torch_module=StubTorch, ) backend = CatboxModelBackend( model_runner=runner, seed_source=lambda: 41100, outcome_source=lambda: "living", ) readiness = backend.readiness() response = backend.observe() self.assertEqual(readiness["status"], "starting") self.assertEqual(response["status"], "generation_failed") self.assertEqual(response["error"]["type"], "RuntimeError") self.assertIn("model weights unavailable", response["error"]["message"]) self.assertNotIn("imageRef", response) def test_real_runner_applies_generation_config_overrides(self): with TemporaryDirectory() as runtime_dir: pipeline = StubPipeline() runner = SdTurboImageToImageModelRunner( runtime_dir=runtime_dir, pipeline_loader=lambda model_id, **kwargs: pipeline, box_image_loader=lambda path, config: "box image", torch_module=StubTorch, now=lambda: "2026-06-02T12:00:00Z", timer=iter([30.0, 30.75]).__next__, ) generated = runner.generate( "dead", seed=41100, config={"steps": 6, "strength": 0.42, "ignored": "not forwarded"}, ) self.assertEqual(pipeline.calls[0]["num_inference_steps"], 6) self.assertEqual(pipeline.calls[0]["strength"], 0.42) self.assertEqual(generated["metadata"]["config"]["steps"], 6) self.assertEqual(generated["metadata"]["config"]["strength"], 0.42) self.assertNotIn("ignored", generated["metadata"]["config"])