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))