code/smoke_test.py

"""CPU smoke test of the PhotoAgent MCTS closed-loop logic (no GPU models).

Validates the algorithm structure: perceiver -> MCTS planner -> executor ->
evaluator -> memory -> closed-loop. Uses a deterministic toy evaluator so we
can run it on CPU in seconds and confirm the tree search + early-stopping work.
"""
import sys, os, time
sys.path.insert(0, os.path.dirname(__file__))
import numpy as np
from PIL import Image
import photoagent_core as P


class ToyEvaluator(P.Evaluator):
    def __init__(self):
        self.device = "cpu"
        self.use_ugc = False
        self.clip = self.aes = self.ugc = None

    def score_ugc(self, image, instruction):
        return 0.0

    def score_aes(self, image):
        return 0.0

    def score_clip(self, image, instruction):
        return 0.0

    def evaluate(self, image, instruction):
        # Deterministic reward: improves with number of edits applied,
        # with diminishing returns -> exercises MCTS preference + early stop.
        h = hash((image.tobytes()[:256], instruction)) & 0xFFFF
        reward = 0.5 + 0.3 * np.random.RandomState(h).rand() 
        return {"ugc": reward, "clip": reward, "reward": reward}


class ToyExecutor(P.IdentityExecutor):
    def execute(self, image, instruction):
        # Slight deterministic perturbation so each edit yields a new image.
        arr = np.asarray(image).astype(np.int16)
        arr = np.clip(arr + np.random.RandomState(hash(instruction) & 0xFFFF).randint(-10, 10), 0, 255).astype(np.uint8)
        return Image.fromarray(arr)


def main():
    rng = np.random.RandomState(0)
    exe = ToyExecutor()
    ev = ToyEvaluator()
    planner = P.MCTSPlanner(exe, ev, depth=2, simulations=8, top_k=2, seed=0)
    images = [Image.fromarray(np.uint8(rng.randint(0,255,(256,256,3)))) for _ in range(3)]
    results = []
    for i, img in enumerate(images):
        out = P.run_photoagent(img, planner, max_iters=3, seed=i)
        results.append({"img": i, "best_score": float(out["best_score"]),
                        "scores": [float(s) for s in out["scores"]],
                        "n_memory": len(out["memory"]),
                        "plan_len": len(out["plan"])})
    print("=== CPU smoke test: closed-loop MCTS PhotoAgent ===")
    for r in results:
        print(r)
    assert all(r["plan_len"] >= 1 for r in results), "planner returned empty plan"
    assert all(r["n_memory"] >= 1 for r in results), "closed-loop produced no memory"
    print("SMOKE TEST OK — closed-loop MCTS agent produced plans + memory.")


if __name__ == "__main__":
    main()