#!/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())