analysis_tools/debug_h20_single_sample_eval.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import importlib.util
import json
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageDraw

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from flow_grpo.server_profiles import apply_server_profile_defaults


apply_server_profile_defaults()
OUT_DIR = REPO_ROOT / "analysis_outputs" / "h20_eval_corruption" / "single_sample_eval"
DEFAULT_OLD_RL_LORA = REPO_ROOT / "logs/radiomics/img-only-r32-a64-bs32-evalbs24-kl-beta0p005-scratch-15k/checkpoints/checkpoint-190/lora"


def load_config(entry: str):
    module_path, function_name = entry.split(":", 1)
    spec = importlib.util.spec_from_file_location("debug_h20_eval_config", Path(module_path).resolve())
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Could not load config module {module_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return getattr(module, function_name)()


def stats(name: str, value: Any) -> dict[str, Any]:
    try:
        import torch
        if isinstance(value, torch.Tensor):
            arr = value.detach().float().cpu().numpy()
        else:
            arr = np.asarray(value)
    except Exception:
        arr = np.asarray(value)
    return {
        "name": name,
        "shape": list(arr.shape),
        "dtype": str(arr.dtype),
        "min": float(np.nanmin(arr)),
        "max": float(np.nanmax(arr)),
        "mean": float(np.nanmean(arr)),
        "std": float(np.nanstd(arr)),
    }


def side_by_side(paths: list[tuple[str, Image.Image]], out_path: Path) -> None:
    cell_w, cell_h = 256, 286
    sheet = Image.new("RGB", (cell_w * len(paths), cell_h), "white")
    draw = ImageDraw.Draw(sheet)
    for i, (label, image) in enumerate(paths):
        img = image.convert("RGB")
        img.thumbnail((cell_w, cell_h - 30), Image.Resampling.BILINEAR)
        x = i * cell_w + (cell_w - img.width) // 2
        y = 28 + (cell_h - 30 - img.height) // 2
        draw.text((i * cell_w + 4, 6), label[:34], fill=(0, 0, 0))
        sheet.paste(img, (x, y))
    sheet.save(out_path)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=os.environ.get("CONFIG_ENTRY", "config/grpo.py:general_radiomics_omnigen_4gpu_kl"))
    parser.add_argument("--sample_index", type=int, default=0)
    parser.add_argument("--max_cases", type=int, default=4)
    parser.add_argument("--old_rl_lora_path", default=os.environ.get("OLD_RL_LORA_PATH") or os.environ.get("EVAL_LORA_PATH") or str(DEFAULT_OLD_RL_LORA))
    parser.add_argument("--output_dir", default=str(OUT_DIR))
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    import torch
    from peft import PeftModel
    from scripts.train_omnigen import RadiomicsEditDataset, _to_rgb_pil, load_omnigen_components, merge_lora_into_base_model, requires_grad
    from flow_grpo.omnigen_patch.omnigen_pipeline_with_logprob import pipeline_with_logprob

    config = load_config(args.config)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    if device.type != "cuda":
        raise RuntimeError("CUDA is not available; single-sample eval needs the real H20 GPU environment.")

    weight_dtype = torch.bfloat16 if getattr(config, "mixed_precision", "bf16") == "bf16" else torch.float16
    dataset = RadiomicsEditDataset(config.dataset, "test")
    sample = dataset[args.sample_index]
    input_image = Image.open(sample["input_image_paths"][0]).convert("RGB")
    gt_image = Image.open(sample["metadata"].get("output_image") or sample["metadata"].get("gt_image")).convert("RGB")

    model, vae, processor = load_omnigen_components(config, device, weight_dtype)
    requires_grad(vae, False)
    if getattr(config, "use_lora", False) and getattr(config.train, "merge_lora_path", None):
        model = merge_lora_into_base_model(model, config.train.merge_lora_path, weight_dtype, trainable=False)
    requires_grad(model, False)
    model.eval()

    cases = [("sft_only", None, 0.0), ("sft_only", None, 0.01), ("sft_plus_old190", args.old_rl_lora_path, 0.0), ("sft_plus_old190", args.old_rl_lora_path, 0.01)]
    results = {
        "config": args.config,
        "device": str(device),
        "weight_dtype": str(weight_dtype),
        "sample_index": args.sample_index,
        "instruction": sample["instruction"],
        "input_path": sample["input_image_paths"][0],
        "gt_path": sample["metadata"].get("output_image") or sample["metadata"].get("gt_image"),
        "cases": [],
    }

    active_model = model
    for case_index, (label, lora_path, noise_level) in enumerate(cases[: args.max_cases]):
        current_model = active_model
        if lora_path:
            if not Path(lora_path).exists():
                results["cases"].append({"label": label, "noise_level": noise_level, "error": f"missing lora path: {lora_path}"})
                continue
            current_model = PeftModel.from_pretrained(active_model, lora_path, is_trainable=False).to(dtype=weight_dtype)
            if hasattr(current_model, "set_adapter"):
                current_model.set_adapter("default")
            current_model.eval()

        generator = torch.Generator(device=device).manual_seed(int(getattr(config, "seed", 0)))
        with torch.no_grad():
            collected = pipeline_with_logprob(
                current_model,
                vae,
                processor,
                [sample["instruction"]],
                [sample["input_image_paths"]],
                height=config.resolution,
                width=config.resolution,
                num_inference_steps=config.sample.eval_num_steps,
                guidance_scale=config.sample.eval_guidance_scale,
                img_guidance_scale=config.sample.eval_img_guidance_scale,
                max_input_image_size=config.sample.max_input_image_size,
                use_img_guidance=config.sample.use_img_guidance,
                use_input_image_size_as_output=config.sample.use_input_image_size_as_output,
                dtype=weight_dtype,
                generator=generator,
                output_type="pt",
                noise_level=noise_level,
                sde_type=config.sample.sde_type,
            )
        output_tensor = collected["images"][0]
        output_pil = _to_rgb_pil(output_tensor)
        case_name = f"{case_index}_{label}_noise{str(noise_level).replace('.', 'p')}"
        output_pil.save(out_dir / f"{case_name}.png")
        side_by_side([("input", input_image), ("generated", output_pil), ("gt", gt_image)], out_dir / f"{case_name}_panel.png")
        results["cases"].append({
            "label": label,
            "lora_path": lora_path,
            "noise_level": noise_level,
            "initial_latent": stats("initial_latent", collected["all_latents"][0]),
            "final_latent": stats("final_latent", collected["all_latents"][-1]),
            "decoded_tensor": stats("decoded_tensor", output_tensor),
            "final_pil": stats("final_pil", np.asarray(output_pil)),
            "output": str(out_dir / f"{case_name}.png"),
        })

    (out_dir / "single_sample_eval_stats.json").write_text(json.dumps(results, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(results, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())