code/run_gpu_job.py

"""GPU job: run PhotoAgent (open-substitute) closed-loop MCTS agent end-to-end.

Runs:
  - Claim 1: full closed-loop PhotoAgent on 6 proxy benchmark photos.
  - Claim 5: MCTS ablation (depth=1 greedy vs depth=3; simulations=10 vs 20)
             reproducing the paper's Table 1 planning-strategy finding.
  - Claim 4: compare PhotoAgent(open) vs open baselines (Identity / single-step
             Step1X) on the same UGC + CLIP reward, a structured proxy for Table 1.

Outputs results.json + edited images to ./outputs, also prints a summary to stdout
(immutable in hf jobs logs).
"""
import sys, os, time, json
sys.path.insert(0, os.path.dirname(__file__))
import numpy as np
from PIL import Image
import photoagent_core as P


def load_images(folder):
    imgs = {}
    for fn in sorted(os.listdir(folder)):
        if fn.lower().endswith((".jpg", ".png", ".jpeg")):
            cat = fn.rsplit(".", 1)[0]
            imgs[cat] = Image.open(os.path.join(folder, fn)).convert("RGB").resize((512, 512))
    return imgs


def run_baseline_identity(images, evaluator):
    """Baseline: no editing (input only) — lower bound reward."""
    out = {}
    for k, img in images.items():
        out[k] = evaluator.evaluate(img, "")["reward"]
    return out


def run_single_step(images, executor, evaluator):
    """Baseline: single open-loop Step1X edit, greedy one action."""
    out = {}
    for k, img in images.items():
        acts = P.perceive(img, hash(k) & 0xFFFF)
        edited = executor.execute(img, acts[0])
        out[k] = evaluator.evaluate(edited, acts[0])["reward"]
    return out


def main():
    device = "cuda" if P.Evaluator is not None else "cpu"
    imgs = load_images("inputs")
    print(f"[job] loaded {len(imgs)} proxy benchmark photos: {list(imgs)}")

    executor = P.InstructPix2PixExecutor(device=device)
    evaluator = P.Evaluator(device=device, use_ugc=False)

    # ---- Claim 1: full closed-loop PhotoAgent ----
    planner = P.MCTSPlanner(executor, evaluator, depth=3, simulations=20, top_k=2, seed=0)
    pa_scores = {}
    pa_ugc = {}
    os.makedirs("outputs/photoagent", exist_ok=True)
    for k, img in imgs.items():
        res = P.run_photoagent(img, planner, max_iters=2, seed=0)
        res["final_image"].save(f"outputs/photoagent/{k}.png")
        pa_scores[k] = res["best_score"]
        pa_ugc[k] = evaluator.evaluate(res["final_image"], "")["aes"]
    pa_mean = float(np.mean(list(pa_scores.values())))

    # ---- Claim 5: MCTS ablation (depth / simulations) ----
    ablation = {}
    for cfg_name, depth, sims in [("depth1_greedy", 1, 20), ("depth3_sim10", 3, 10), ("depth3_sim20", 3, 20)]:
        pl = P.MCTSPlanner(executor, evaluator, depth=depth, simulations=sims, top_k=2, seed=0)
        sc = []
        for k, img in imgs.items():
            r = P.run_photoagent(img, pl, max_iters=2, seed=0)
            sc.append(r["best_score"])
        ablation[cfg_name] = float(np.mean(sc))
        print(f"[ablation] {cfg_name}: mean reward = {ablation[cfg_name]:.4f}")

    # ---- Claim 4: baselines vs PhotoAgent ---- 
    base_identity = run_baseline_identity(imgs, evaluator)
    base_single = run_single_step(imgs, executor, evaluator)
    base_identity_mean = float(np.mean(list(base_identity.values())))
    base_single_mean = float(np.mean(list(base_single.values())))

    # ---- Claim 2 (proxy reward model) + Claim 3 (benchmark composition) ----
    import subprocess
    try:
        subprocess.run([sys.executable, "claim2_reward_model.py"], check=False)
        print("[claim2] reward-model proxy done")
    except Exception as e:
        print("[claim2] failed:", e)
    try:
        subprocess.run([sys.executable, "claim3_benchmark.py"], check=False)
        print("[claim3] benchmark proxy done")
    except Exception as e:
        print("[claim3] failed:", e)

    summary = {
        "setting": "toy proxy (6 photos, open substitutes)",
        "photoagent_mean_reward": pa_mean,
        "baseline_identity_mean": base_identity_mean,
        "baseline_single_step_mean": base_single_mean,
        "ablation_depth_sim": ablation,
        "per_category_photoagent": pa_scores,
        "per_category_photoagent_ugc": pa_ugc,
        "per_category_identity": base_identity,
        "per_category_single_step": base_single,
        "executor": "timbrooks/instruct-pix2pix (open; paper's own baseline)",
        "evaluator": "shunk031/aesthetics-predictor-v2 (ImageReward proxy) + openai/clip-vit-base-patch32",
    }
    os.makedirs("outputs", exist_ok=True)
    with open("outputs/results.json", "w") as f:
        json.dump(summary, f, indent=2)
    print("=== GPU JOB SUMMARY ===")
    print(json.dumps({k: v for k, v in summary.items()
                      if not isinstance(v, dict)}, indent=2))
    print("PhotoAgent wins vs identity baseline:",
          pa_mean > base_identity_mean)
    print("PhotoAgent (depth3_sim20) vs depth1_greedy:",
          ablation["depth3_sim20"] > ablation["depth1_greedy"])


if __name__ == "__main__":
    main()