import contextlib import datetime import json import os import sys from absl import app, flags from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import set_seed import traceback from ml_collections import config_flags import numpy as np import torch from torch.utils.data import DataLoader, Subset import tqdm from flow_grpo.omnigen_patch.omnigen_pipeline_with_logprob import pipeline_with_logprob, pipeline_with_logprob_joint_image_reward from flow_grpo.omnigen_patch.joint_model_loader import load_joint_omnigen_components_for_rl from scripts.train_omnigen import ( RadiomicsEditDataset, load_omnigen_components, merge_lora_into_base_model, unwrap_model, requires_grad, _to_rgb_pil, ) tqdm = tqdm.tqdm FLAGS = flags.FLAGS if "config" not in FLAGS: config_flags.DEFINE_config_file("config", "config/base.py", "Test configuration.") if "output_dir" not in FLAGS: flags.DEFINE_string("output_dir", "/data/wtchen/code/flow_grpo_cxr/outputs", "Output directory") if "resume_dir" not in FLAGS: flags.DEFINE_string("resume_dir", None, "If set, resume evaluation in this exact directory and skip exiting images.") if "eval_lora_path" not in FLAGS: flags.DEFINE_string("eval_lora_path", None, "Optional LoRA path specifically for evaluation to override the one in config") if "max_samples" not in FLAGS: flags.DEFINE_integer("max_samples", None, "If set, evaluate at most this many dataset samples globally.") logger = get_logger(__name__) def _target_image_path(output_dir, metadata): output_image = metadata.get("output_image") or metadata.get("gt_image") if output_image: norm = output_image.replace("\\", "/").rstrip("/") parts = norm.split("/") if len(parts) >= 2: return os.path.join(output_dir, parts[-2], parts[-1]) patient_id = metadata.get("patient_id", "unknown_patient") image_name = metadata.get("image_name", "unknown_image.png") return os.path.join(output_dir, patient_id, image_name) def _is_complete_image(path): return os.path.isfile(path) and os.path.getsize(path) > 0 def _write_jsonl(path, records): with open(path, "w", encoding="utf-8") as f: for record in records: f.write(json.dumps(record, ensure_ascii=False) + "\n") def main(_): config = FLAGS.config if FLAGS.resume_dir: output_dir = FLAGS.resume_dir run_name = os.path.basename(FLAGS.resume_dir) unique_id = "resumed" else: unique_id = datetime.datetime.now().strftime("%Y.%m.%d_%H.%M.%S") run_name = config.run_name if config.run_name else "eval" run_name = f"{run_name}_eval_{unique_id}" output_dir = os.path.join(FLAGS.output_dir, run_name) accelerator = Accelerator(mixed_precision=config.mixed_precision) os.makedirs(output_dir, exist_ok=True) logger.info(f"\n{config}") set_seed(config.seed, device_specific=True) weight_dtype = torch.float32 if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 logger.info("Loading base model components...") use_joint_mask = bool(getattr(config, "use_joint_mask", False)) current_eval_lora_path = getattr(FLAGS, "eval_lora_path", None) or getattr(config.train, "lora_path", None) if use_joint_mask: model, vae, processor, _ = load_joint_omnigen_components_for_rl( config, device=accelerator.device, weight_dtype=weight_dtype, attach_rl_lora=False, eval_lora_path=current_eval_lora_path, ) else: model, vae, processor = load_omnigen_components(config, accelerator.device, weight_dtype) requires_grad(vae, False) if config.use_lora: merge_lora_path = getattr(config.train, "merge_lora_path", None) if merge_lora_path: model = merge_lora_into_base_model( model, merge_lora_path, weight_dtype, trainable=False, ) requires_grad(model, False) lora_path_to_eval = current_eval_lora_path if lora_path_to_eval: from peft import PeftModel logger.info(f"Loading evaluation LoRA adapter from {lora_path_to_eval}") model = PeftModel.from_pretrained( model, lora_path_to_eval, is_trainable=False, ) if hasattr(model, "set_adapter"): model.set_adapter("default") model.to(dtype=weight_dtype) elif config.train.lora_path: lora_path_to_eval = current_eval_lora_path model = merge_lora_into_base_model(model, lora_path_to_eval, weight_dtype, trainable=False) else: requires_grad(model, False) # Persist evaluation LoRA path for reproducibility. eval_info_path = os.path.join(output_dir, "eval_lora_path.json") with open(eval_info_path, "w", encoding="utf-8") as f: json.dump({"eval_lora_path": current_eval_lora_path}, f, indent=2) logger.info(f"Saved eval LoRA metadata: {eval_info_path}") model.eval() test_dataset = RadiomicsEditDataset(config.dataset, "test") model = accelerator.prepare(model) logger.info("***** Running OmniGen Evaluation (All Samples) *****") logger.info(f" Test batch size per device = {config.sample.test_batch_size}") logger.info(f" Process rank = {accelerator.process_index}") logger.info(f" Number of processes = {accelerator.num_processes}") requested_eval_indices = list(range(len(test_dataset))) if FLAGS.max_samples is not None: requested_eval_indices = requested_eval_indices[: FLAGS.max_samples] all_eval_indices = requested_eval_indices if FLAGS.resume_dir: all_eval_indices = [ index for index in requested_eval_indices if not _is_complete_image(_target_image_path(output_dir, test_dataset.metadatas[index])) ] logger.info( f"Resume mode: {len(all_eval_indices)} / {len(requested_eval_indices)} requested images are missing or empty." ) rank_eval_indices = all_eval_indices[accelerator.process_index::accelerator.num_processes] rank_dataset = Subset(test_dataset, rank_eval_indices) test_dataloader = DataLoader( rank_dataset, batch_size=config.sample.test_batch_size, shuffle=False, num_workers=2, collate_fn=RadiomicsEditDataset.collate_fn, drop_last=False, ) autocast = accelerator.autocast eval_dataloader_len = len(test_dataloader) eval_iterable = test_dataloader eval_total = eval_dataloader_len logger.info(f" Total dataset samples = {len(test_dataset)}") logger.info(f" Samples assigned to this rank = {len(rank_eval_indices)}") logger.info(f" Total evaluation steps on this rank = {eval_total}") def run_generation(batch_instructions, batch_input_image_paths): nonlocal processor with autocast(): with torch.no_grad(): eval_pipeline_fn = pipeline_with_logprob_joint_image_reward if use_joint_mask else pipeline_with_logprob collected = eval_pipeline_fn( model, vae, processor, batch_instructions, batch_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, output_type="pt", noise_level=getattr(config.sample, "noise_level", 0.0), sde_type=config.sample.sde_type, mask_scale_factor=getattr(getattr(config, "joint", {}), "mask_scale_factor", 1.0), ) processor = collected["processor"] return collected["images"].float().cpu().numpy() def save_image(image_array, metadata): img_path = _target_image_path(output_dir, metadata) if _is_complete_image(img_path): return False patient_dir = os.path.dirname(img_path) os.makedirs(patient_dir, exist_ok=True) tmp_path = f"{img_path}.rank{accelerator.process_index}.pid{os.getpid()}.tmp" img = _to_rgb_pil(image_array) img.save(tmp_path, format="PNG") os.replace(tmp_path, img_path) return True saved_count = 0 skipped_count = 0 failed_records = [] for batch_index, test_batch in enumerate( tqdm( eval_iterable, desc="Eval", total=eval_total, disable=not accelerator.is_local_main_process, dynamic_ncols=True, ) ): prompts, instructions, prompt_metadata, input_image_paths, ref_images, _ = test_batch logger.info(f"[Rank {accelerator.process_index}] Processing batch_index {batch_index}, batch_size={len(prompt_metadata)}") try: local_images = run_generation(instructions, input_image_paths) if len(local_images) != len(prompt_metadata): raise RuntimeError( f"Pipeline returned {len(local_images)} images for {len(prompt_metadata)} metadata entries." ) except Exception as e: logger.error( f"Batch generation failed on rank {accelerator.process_index}, " f"batch_index {batch_index}: {e}. Retrying one sample at a time." ) traceback.print_exc() if torch.cuda.is_available(): torch.cuda.empty_cache() for sample_index, metadata in enumerate(prompt_metadata): try: if _is_complete_image(_target_image_path(output_dir, metadata)): skipped_count += 1 continue local_images = run_generation( [instructions[sample_index]], [input_image_paths[sample_index]], ) if len(local_images) != 1: raise RuntimeError(f"Single-sample retry returned {len(local_images)} images.") if save_image(local_images[0], metadata): saved_count += 1 logger.info(f"Saved: {_target_image_path(output_dir, metadata)}") else: skipped_count += 1 except Exception as sample_error: logger.error( f"Sample failed on rank {accelerator.process_index}, " f"batch_index {batch_index}, sample_index {sample_index}: {sample_error}" ) failed_records.append( { "rank": accelerator.process_index, "batch_index": batch_index, "sample_index": sample_index, "target_path": _target_image_path(output_dir, metadata), "metadata": metadata, "error": repr(sample_error), "traceback": traceback.format_exc(), } ) if torch.cuda.is_available(): torch.cuda.empty_cache() continue for sample_index, metadata in enumerate(prompt_metadata): try: if save_image(local_images[sample_index], metadata): saved_count += 1 logger.info(f"Saved: {_target_image_path(output_dir, metadata)}") else: skipped_count += 1 except Exception as save_error: logger.error( f"Save failed on rank {accelerator.process_index}, " f"batch_index {batch_index}, sample_index {sample_index}: {save_error}" ) failed_records.append( { "rank": accelerator.process_index, "batch_index": batch_index, "sample_index": sample_index, "target_path": _target_image_path(output_dir, metadata), "metadata": metadata, "error": repr(save_error), "traceback": traceback.format_exc(), } ) failed_path = os.path.join(output_dir, f"failed_rank_{accelerator.process_index}.jsonl") _write_jsonl(failed_path, failed_records) status_path = os.path.join(output_dir, f"status_rank_{accelerator.process_index}.json") with open(status_path, "w", encoding="utf-8") as f: json.dump( { "rank": accelerator.process_index, "assigned_samples": len(rank_eval_indices), "saved": saved_count, "skipped_existing": skipped_count, "failed": len(failed_records), }, f, indent=2, ) accelerator.wait_for_everyone() if accelerator.is_main_process: missing_records = [] for index in requested_eval_indices: metadata = test_dataset.metadatas[index] target_path = _target_image_path(output_dir, metadata) if not _is_complete_image(target_path): missing_records.append( { "index": index, "target_path": target_path, "metadata": metadata, } ) missing_path = os.path.join(output_dir, "missing_images.jsonl") _write_jsonl(missing_path, missing_records) summary_path = os.path.join(output_dir, "eval_summary.json") with open(summary_path, "w", encoding="utf-8") as f: json.dump( { "dataset_size": len(test_dataset), "requested_images": len(requested_eval_indices), "complete_images": len(requested_eval_indices) - len(missing_records), "missing_images": len(missing_records), "output_dir": output_dir, "resume_dir": FLAGS.resume_dir, }, f, indent=2, ) if missing_records: raise RuntimeError( f"Evaluation did not finish all images: {len(missing_records)} missing. " f"See {missing_path}" ) logger.info("Evaluation finished.") if __name__ == "__main__": app.run(main)