scripts/eval_omnigen.py

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)