lrm/lrm_15/test.py

import os

os.environ["HF_DATASETS_OFFLINE"] = "1"
os.environ["HF_METRICS_OFFLINE"] = "1"
os.environ["HF_MODULES_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["DIFFUSERS_OFFLINE"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"
import json
import sys
import tempfile
from io import BytesIO
from glob import glob
from pathlib import Path

import torch
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from datasets import load_dataset
from PIL import Image
from torchvision import transforms
from transformers import CLIPTokenizer
from accelerate.state import PartialState

from trainer.models.sd15_preference_model import SD15PreferenceModel, SD15PreferenceModelConfig

# Needed for accelerate.logging.get_logger calls used inside model.load().
_ = PartialState()


# -----------------
# Config
# -----------------
PROJECT_ROOT = Path('/g/data/rr81/LPO/lrm/lrm_15').resolve()
LOCAL_LRM_SD15_DIR = PROJECT_ROOT / 'LRM' / 'lrm_sd15'
BASE_SD15_ID = 'stable-diffusion-v1-5/stable-diffusion-v1-5'
DATASET_NAME = 'pickapic-anonymous/pickapic_v1'
SPLIT = 'test_unique'
BATCH_SIZE = 1
NUM_WORKERS = 2
MAX_BATCHES = None  # e.g. set 50 for quick check

os.chdir(PROJECT_ROOT)
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Project root:', PROJECT_ROOT)
print('Python:', sys.executable)
print('Torch:', torch.__version__)
print('CUDA available:', torch.cuda.is_available())
print('Device:', DEVICE)
print('Local SD1.5 LRM dir:', LOCAL_LRM_SD15_DIR)


# -----------------
# Load local SD1.5 LRM weights
# -----------------
if not ((LOCAL_LRM_SD15_DIR / 'state_dict.pt').exists() and (LOCAL_LRM_SD15_DIR / 'unet').exists() and (LOCAL_LRM_SD15_DIR / 'text_encoder').exists()):
    raise FileNotFoundError(
        f'Local SD1.5 LRM path is incomplete: {LOCAL_LRM_SD15_DIR}. '
        'Expected: state_dict.pt, unet/, text_encoder/'
    )

# SD15PreferenceModel expects a local torch checkpoint containing text_projection.weight at init time.
# Create a tiny placeholder; model.load(...) below replaces with actual LRM weights.
tmp_clip_file = Path(tempfile.gettempdir()) / 'lrm_sd15_dummy_clip_projection.pt'
if not tmp_clip_file.exists():
    torch.save({'text_projection.weight': torch.eye(768, dtype=torch.float32)}, tmp_clip_file)

model_cfg = SD15PreferenceModelConfig(
    pretrained_model_name_or_path=BASE_SD15_ID,
    clip_ckpt_path=str(tmp_clip_file),
    freeze_text_encoder=False,
)

model = SD15PreferenceModel(model_cfg)
model.load(str(LOCAL_LRM_SD15_DIR))
model.to(DEVICE).eval()

print('Model loaded from local LRM successfully.')
print('logit_scale(exp):', float(model.logit_scale.exp().detach().cpu().item()))


# -----------------
# Eval helpers (same metric style as training)
# -----------------
def features2probs(model_obj, text_features, image_0_features, image_1_features):
    image_0_scores = model_obj.logit_scale.exp() * torch.diag(torch.einsum('bd,cd->bc', text_features, image_0_features))
    image_1_scores = model_obj.logit_scale.exp() * torch.diag(torch.einsum('bd,cd->bc', text_features, image_1_features))
    scores = torch.stack([image_0_scores, image_1_scores], dim=-1)
    probs = torch.softmax(scores, dim=-1)
    return probs[:, 0], probs[:, 1]


def get_features(model_obj, input_ids, pixels_0_values, pixels_1_values, timesteps):
    all_pixel_values = torch.cat([pixels_0_values, pixels_1_values], dim=0)
    timesteps = timesteps.reshape(-1, 2)
    timesteps = torch.cat([timesteps[:, 0], timesteps[:, 1]], dim=0)
    text_features, all_image_features = model_obj(text_inputs=input_ids, image_inputs=all_pixel_values, time_cond=timesteps)
    all_image_features = all_image_features / all_image_features.norm(dim=-1, keepdim=True)
    text_features = text_features / text_features.norm(dim=-1, keepdim=True)
    image_0_features, image_1_features = all_image_features.chunk(2, dim=0)
    return image_0_features, image_1_features, text_features


def load_dataset_split_like_sana(dataset_name: str, split: str):
    offline_mode = os.getenv("HF_HUB_OFFLINE", "0").strip().lower() in {"1", "true", "yes", "on"}
    if not offline_mode:
        return load_dataset(dataset_name, split=split)

    if "/" not in dataset_name:
        return load_dataset(dataset_name, split=split)

    org, name = dataset_name.split("/", 1)

    # Follow lrm_sana behavior, but also probe common cache roots when env vars are unset.
    cache_candidates = []
    for p in [
        os.getenv("HF_HUB_CACHE"),
        os.getenv("HUGGINGFACE_HUB_CACHE"),
        (os.path.join(os.getenv("HF_HOME"), "hub") if os.getenv("HF_HOME") else None),
        os.path.expanduser("~/.cache/huggingface/hub"),
        "/scratch/rr81/ma5430/.cache/huggingface/hub",
    ]:
        if p and p not in cache_candidates:
            cache_candidates.append(p)

    repo_cache_dirs = [
        os.path.join(cache_root, f"datasets--{org}--{name}")
        for cache_root in cache_candidates
        if os.path.isdir(os.path.join(cache_root, f"datasets--{org}--{name}"))
    ]

    for repo_cache_dir in repo_cache_dirs:
        snapshot_dir = None
        ref_main = os.path.join(repo_cache_dir, "refs", "main")
        if os.path.isfile(ref_main):
            revision = open(ref_main, "r", encoding="utf-8").read().strip()
            candidate = os.path.join(repo_cache_dir, "snapshots", revision)
            if os.path.isdir(candidate):
                snapshot_dir = candidate

        if snapshot_dir is None:
            snapshots = sorted(glob(os.path.join(repo_cache_dir, "snapshots", "*")))
            if snapshots:
                snapshot_dir = snapshots[-1]

        if snapshot_dir is None:
            continue

        data_dir = os.path.join(snapshot_dir, "data")
        if not os.path.isdir(data_dir):
            continue

        selected_split = split
        parquet_files = sorted(glob(os.path.join(data_dir, f"{selected_split}-*.parquet")))
        if not parquet_files and split.startswith("validation"):
            for alt_split in ("test_unique", "test"):
                alt_files = sorted(glob(os.path.join(data_dir, f"{alt_split}-*.parquet")))
                if alt_files:
                    selected_split = alt_split
                    parquet_files = alt_files
                    print(f"Offline cache missing split '{split}', falling back to '{selected_split}'")
                    break

        if parquet_files:
            print(
                f"Loading cached offline split '{selected_split}' from {len(parquet_files)} parquet shards\n"
                f"cache={repo_cache_dir}"
            )
            return load_dataset("parquet", data_files=parquet_files, split="train")

    raise RuntimeError(
        "Offline mode is enabled and cached parquet dataset was not found. "
        f"Searched cache roots: {cache_candidates}. "
        "Set HF_HUB_CACHE/HF_HOME to your predownloaded cache root or disable offline mode."
    )


image_transform = transforms.Compose([
    transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.BILINEAR),
    transforms.CenterCrop(512),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5]),
])

tokenizer = CLIPTokenizer.from_pretrained(BASE_SD15_ID, subfolder='tokenizer')
raw_test = load_dataset_split_like_sana(DATASET_NAME, SPLIT)
# Match training behavior: keep only labeled examples in non-train splits.
raw_test = raw_test.filter(lambda x: x['has_label'])


def to_image(x):
    if isinstance(x, dict):
        x = x['bytes']
    if isinstance(x, bytes):
        x = Image.open(BytesIO(x))
    return x.convert('RGB')


def preprocess_example(example):
    input_ids = tokenizer(
        example['caption'],
        max_length=tokenizer.model_max_length,
        padding='max_length',
        truncation=True,
        return_tensors='pt',
    ).input_ids.squeeze(0)

    pixel_0 = image_transform(to_image(example['jpg_0']))
    pixel_1 = image_transform(to_image(example['jpg_1']))

    # Non-train split uses timestep=1 in existing pipeline.
    timestep = torch.tensor([1, 1], dtype=torch.long)

    return {
        'input_ids': input_ids,
        'pixel_values_0': pixel_0,
        'pixel_values_1': pixel_1,
        'label_0': torch.tensor(example['label_0'], dtype=torch.long),
        'label_1': torch.tensor(example['label_1'], dtype=torch.long),
        'timestep': timestep,
    }


def collate_fn(batch):
    return {
        'input_ids': torch.stack([x['input_ids'] for x in batch], dim=0),
        'pixel_values_0': torch.stack([x['pixel_values_0'] for x in batch], dim=0),
        'pixel_values_1': torch.stack([x['pixel_values_1'] for x in batch], dim=0),
        'label_0': torch.stack([x['label_0'] for x in batch], dim=0),
        'label_1': torch.stack([x['label_1'] for x in batch], dim=0),
        'timestep': torch.stack([x['timestep'] for x in batch], dim=0),
    }


class EvalDataset(torch.utils.data.Dataset):
    def __init__(self, hf_ds):
        self.hf_ds = hf_ds

    def __len__(self):
        return len(self.hf_ds)

    def __getitem__(self, idx):
        return preprocess_example(self.hf_ds[idx])


eval_ds = EvalDataset(raw_test)
loader = DataLoader(
    eval_ds,
    shuffle=False,
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    collate_fn=collate_fn,
)


# -----------------
# Run evaluation
# -----------------
all_correct = []
num_batches = 0

with torch.no_grad():
    for batch in tqdm(loader, desc=f'Evaluating {SPLIT}'):
        num_batches += 1

        for k, v in list(batch.items()):
            if torch.is_tensor(v):
                batch[k] = v.to(DEVICE)

        image_0_features, image_1_features, text_features = get_features(
            model,
            batch['input_ids'],
            batch['pixel_values_0'],
            batch['pixel_values_1'],
            batch['timestep'],
        )

        image_0_probs, image_1_probs = features2probs(model, text_features, image_0_features, image_1_features)

        agree_on_0 = (image_0_probs > image_1_probs) * batch['label_0']
        agree_on_1 = (image_0_probs < image_1_probs) * batch['label_1']
        is_correct = (agree_on_0 + agree_on_1).detach().cpu()
        all_correct.append(is_correct)

        if MAX_BATCHES is not None and num_batches >= MAX_BATCHES:
            break

correct_tensor = torch.cat(all_correct).float() if all_correct else torch.tensor([], dtype=torch.float32)
accuracy = float(correct_tensor.mean().item()) if correct_tensor.numel() > 0 else float('nan')
num_samples = int(correct_tensor.numel())

metrics = {
    'split': SPLIT,
    'accuracy': accuracy,
    'num_samples': num_samples,
    f'{SPLIT}_accuracy': accuracy,
    f'{SPLIT}_num_samples': num_samples,
    'logit_scale': float(model.logit_scale.exp().detach().cpu().item()),
    'evaluated_batches': num_batches,
}

print(json.dumps(metrics, indent=2))