Buckets:
| """Evaluate autonomous image-boundary emission and decode teacher-forced BSQ predictions.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from peft import PeftModel | |
| from transformers import AutoProcessor | |
| from uniar import UniARForConditionalGeneration, UniARVisualDecoder, configure_processor | |
| from .dataset import ZebraTrajectoryDataset, image_gen_token_id | |
| from .objective import unwrap_ar_model | |
| from .trajectory import render_messages, resolve_image_path | |
| def _load_rgb(root: Path, relative: str): | |
| from PIL import Image | |
| with Image.open(resolve_image_path(root, relative)) as image: | |
| return image.convert("RGB").copy() | |
| def evaluate(args: argparse.Namespace) -> dict: | |
| from qwen_vl_utils import process_vision_info | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("evaluation requires CUDA") | |
| output_dir = args.run_dir / "evaluation" | |
| output_dir.mkdir(exist_ok=False) | |
| processor = configure_processor( | |
| AutoProcessor.from_pretrained(args.model_path, use_fast=True) | |
| ) | |
| marker_id = image_gen_token_id(processor) | |
| im_end_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>") | |
| base = UniARForConditionalGeneration.from_pretrained( | |
| args.model_path, | |
| torch_dtype=torch.bfloat16, | |
| attn_implementation=args.attn, | |
| low_cpu_mem_usage=True, | |
| ).to("cuda") | |
| model = PeftModel.from_pretrained(base, args.run_dir / "adapter").eval() | |
| dataset = ZebraTrajectoryDataset(args.samples_root) | |
| rows = [] | |
| switch_state = None | |
| with (output_dir / "autonomous_boundaries.jsonl").open("w", encoding="utf-8") as handle: | |
| for index, trajectory in enumerate(dataset.records): | |
| messages, _ = render_messages( | |
| trajectory, lambda path: _load_rgb(dataset.root, path) | |
| ) | |
| prompt_messages = messages[:-1] | |
| text = processor.apply_chat_template( | |
| prompt_messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| images, videos = process_vision_info(prompt_messages) | |
| inputs = processor( | |
| text=[text], images=images, videos=videos, return_tensors="pt" | |
| ).to("cuda") | |
| inputs.pop("mm_token_type_ids", None) | |
| output_ids = model.generate( | |
| **inputs, | |
| max_new_tokens=args.max_new_tokens, | |
| do_sample=False, | |
| eos_token_id=[marker_id, im_end_id], | |
| pad_token_id=processor.tokenizer.pad_token_id, | |
| ) | |
| generated = output_ids[0, inputs.input_ids.shape[1] :] | |
| emitted = bool((generated == marker_id).any()) | |
| decoded = processor.tokenizer.decode( | |
| generated, skip_special_tokens=False, clean_up_tokenization_spaces=False | |
| ) | |
| row = { | |
| "sample_id": trajectory["sample_id"], | |
| "emitted_image_gen": emitted, | |
| "generated_tokens": generated.tolist(), | |
| "decoded": decoded, | |
| } | |
| rows.append(row) | |
| if emitted and switch_state is None: | |
| switch_state = { | |
| "sample_id": trajectory["sample_id"], | |
| "output_ids": output_ids.detach(), | |
| "pixel_values": inputs.pixel_values.detach(), | |
| "image_grid_thw": inputs.image_grid_thw.detach(), | |
| } | |
| handle.write(json.dumps(row, ensure_ascii=False) + "\n") | |
| print( | |
| f"boundary {index + 1}/{len(dataset)} {trajectory['sample_id']} " | |
| f"image_gen={emitted}", | |
| flush=True, | |
| ) | |
| if switch_state is None: | |
| raise RuntimeError("no trajectory emitted <image_gen>; cannot test modality switch") | |
| ar_model = unwrap_ar_model(model) | |
| tail_ids = processor.tokenizer.encode( | |
| f" generate image {args.image_size // 32} {args.image_size // 32}<|vision_start|>", | |
| add_special_tokens=False, | |
| return_tensors="pt", | |
| ).to("cuda") | |
| visual_prefix_ids = torch.cat([switch_state["output_ids"], tail_ids], dim=1) | |
| prefix_attention = torch.ones_like(visual_prefix_ids) | |
| image_token_num = (args.image_size // 32) ** 2 | |
| output_placeholders = torch.full( | |
| (1, image_token_num), | |
| ar_model.config.image_token_id, | |
| dtype=torch.long, | |
| device="cuda", | |
| ) | |
| full_ids = torch.cat([visual_prefix_ids, output_placeholders], dim=1) | |
| full_attention = torch.ones_like(full_ids) | |
| output_grid = torch.tensor( | |
| [[1, args.image_size // 16, args.image_size // 16]], | |
| dtype=torch.long, | |
| device="cuda", | |
| ) | |
| combined_grid = torch.cat([switch_state["image_grid_thw"], output_grid], dim=0) | |
| position_ids, _ = ar_model.model.get_rope_index( | |
| full_ids, image_grid_thw=combined_grid, attention_mask=full_attention | |
| ) | |
| torch.manual_seed(args.seed) | |
| autonomous_indices = ar_model.generate_visual( | |
| prefix_input_ids=visual_prefix_ids, | |
| attention_mask=prefix_attention, | |
| pos_ids_image=position_ids[:, :, -image_token_num:], | |
| pos_ids_all=position_ids, | |
| image_token_num=image_token_num, | |
| temperature=args.visual_temperature, | |
| cfg=1.0, | |
| pixel_values=switch_state["pixel_values"], | |
| input_image_grid_thw=switch_state["image_grid_thw"], | |
| ) | |
| torch.save( | |
| autonomous_indices.cpu(), output_dir / "autonomous_first_visual_indices.pt" | |
| ) | |
| bits = torch.load( | |
| args.run_dir / "teacher_forced_first_image_bits.pt", | |
| map_location="cpu", | |
| weights_only=True, | |
| ) | |
| predicted_indices = ar_model.model.visual.bsq.codes_to_indexes( | |
| bits["predicted_bits"].to("cuda", torch.bfloat16) * 2 - 1 | |
| ) | |
| target_indices = ar_model.model.visual.bsq.codes_to_indexes( | |
| bits["target_bits"].to("cuda", torch.bfloat16) * 2 - 1 | |
| ) | |
| decoder = UniARVisualDecoder.from_pretrained(args.model_path, device=torch.device("cuda")) | |
| for filename, indices in [ | |
| ("teacher_forced_first_predicted.png", predicted_indices), | |
| ("teacher_forced_first_target.png", target_indices), | |
| ("autonomous_first_generated.png", autonomous_indices), | |
| ]: | |
| image = decoder.decode( | |
| indices.unsqueeze(0) if indices.ndim == 3 else indices, | |
| ar_height=args.image_size, | |
| ar_width=args.image_size, | |
| upsampling_ratio=1.0, | |
| num_inference_steps=args.decoder_steps, | |
| guidance_scale=1.0, | |
| )[0] | |
| image.save(output_dir / filename) | |
| summary = { | |
| "samples": len(rows), | |
| "autonomous_image_gen": sum(row["emitted_image_gen"] for row in rows), | |
| "rate": sum(row["emitted_image_gen"] for row in rows) / len(rows), | |
| "max_new_tokens": args.max_new_tokens, | |
| "decoded_images": [ | |
| "teacher_forced_first_predicted.png", | |
| "teacher_forced_first_target.png", | |
| "autonomous_first_generated.png", | |
| ], | |
| "autonomous_visual_sample": switch_state["sample_id"], | |
| } | |
| (output_dir / "summary.json").write_text( | |
| json.dumps(summary, indent=2) + "\n", encoding="utf-8" | |
| ) | |
| print(json.dumps(summary), flush=True) | |
| return summary | |
| def build_parser() -> argparse.ArgumentParser: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--model-path", type=Path, default=Path("models/UniAR-SFT")) | |
| parser.add_argument("--samples-root", type=Path, default=Path("data/Zebra-CoT-samples")) | |
| parser.add_argument("--run-dir", type=Path, required=True) | |
| parser.add_argument("--attn", default="flash_attention_2") | |
| parser.add_argument("--max-new-tokens", type=int, default=512) | |
| parser.add_argument("--image-size", type=int, default=256) | |
| parser.add_argument("--decoder-steps", type=int, default=8) | |
| parser.add_argument("--visual-temperature", type=float, default=0.1) | |
| parser.add_argument("--seed", type=int, default=17) | |
| return parser | |
| def main() -> None: | |
| evaluate(build_parser().parse_args()) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.17 kB
- Xet hash:
- eaec50876673620c38d1ccabab5ab62bc21e208fbf28d62eecc662328c4b4851
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.