Buckets:
| """Replay the first autonomous visual state and resume UniAR language decoding.""" | |
| 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, 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 continue_after_visual(args: argparse.Namespace) -> dict: | |
| from qwen_vl_utils import process_vision_info | |
| 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() | |
| ar_model = unwrap_ar_model(model) | |
| dataset = ZebraTrajectoryDataset(args.samples_root) | |
| trajectory = dataset[0] | |
| messages, _ = render_messages( | |
| trajectory, lambda path: _load_rgb(dataset.root, path) | |
| ) | |
| prompt_messages = messages[:-1] | |
| prompt_text = processor.apply_chat_template( | |
| prompt_messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| images, videos = process_vision_info(prompt_messages) | |
| inputs = processor( | |
| text=[prompt_text], images=images, videos=videos, return_tensors="pt" | |
| ).to("cuda") | |
| inputs.pop("mm_token_type_ids", None) | |
| with (args.run_dir / "evaluation" / "autonomous_boundaries.jsonl").open() as handle: | |
| boundary = json.loads(next(handle)) | |
| generated_ids = torch.tensor( | |
| [boundary["generated_tokens"]], dtype=torch.long, device="cuda" | |
| ) | |
| grid_size = args.image_size // 32 | |
| tail_ids = processor.tokenizer.encode( | |
| f" generate image {grid_size} {grid_size}<|vision_start|>", | |
| add_special_tokens=False, | |
| return_tensors="pt", | |
| ).to("cuda") | |
| visual_prefix_ids = torch.cat([inputs.input_ids, generated_ids, tail_ids], dim=1) | |
| visual_indices = torch.load( | |
| args.run_dir / "evaluation" / "autonomous_first_visual_indices.pt", | |
| map_location="cuda", | |
| weights_only=True, | |
| ) | |
| image_token_num = visual_indices.shape[1] | |
| output_placeholders = torch.full( | |
| (1, image_token_num), | |
| ar_model.config.image_token_id, | |
| dtype=torch.long, | |
| device="cuda", | |
| ) | |
| vision_end = torch.tensor( | |
| [[ar_model.config.vision_end_token_id]], dtype=torch.long, device="cuda" | |
| ) | |
| full_ids = torch.cat([visual_prefix_ids, output_placeholders, vision_end], dim=1) | |
| attention_mask = 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([inputs.image_grid_thw, output_grid], dim=0) | |
| position_ids, _ = ar_model.model.get_rope_index( | |
| full_ids, image_grid_thw=combined_grid, attention_mask=attention_mask | |
| ) | |
| input_main_parts, input_deep = ar_model.model.get_image_features( | |
| inputs.pixel_values, inputs.image_grid_thw | |
| ) | |
| input_main = torch.cat(input_main_parts, dim=0) | |
| signed_codes = ar_model.model.visual.bsq.indexes_to_codes(visual_indices) | |
| raw_codes = (signed_codes + 1) // 2 | |
| output_main, output_deep = ar_model.model.visual.bsq_indices_to_features( | |
| raw_codes, with_merger=True | |
| ) | |
| visual_main = torch.cat([input_main, output_main], dim=0) | |
| visual_deep = [ | |
| torch.cat([left, right], dim=0) | |
| for left, right in zip(input_deep, output_deep) | |
| ] | |
| embeddings = ar_model.get_input_embeddings()(full_ids) | |
| visual_mask = full_ids == ar_model.config.image_token_id | |
| embeddings = embeddings.masked_scatter( | |
| visual_mask.unsqueeze(-1).expand_as(embeddings), | |
| visual_main.to(embeddings.dtype), | |
| ) | |
| text_model = ar_model.model.language_model | |
| outputs = text_model( | |
| inputs_embeds=embeddings, | |
| attention_mask=attention_mask, | |
| position_ids=position_ids, | |
| visual_pos_masks=visual_mask, | |
| deepstack_visual_embeds=[value.to(embeddings.dtype) for value in visual_deep], | |
| use_cache=True, | |
| ) | |
| cache = outputs.past_key_values | |
| logits = ar_model.lm_head(outputs.last_hidden_state[:, -1, :])[0] | |
| next_position = position_ids[:, :, -1:] + 1 | |
| continuation: list[int] = [] | |
| for _ in range(args.max_new_tokens): | |
| token_id = int(logits.argmax()) | |
| continuation.append(token_id) | |
| if token_id in {marker_id, im_end_id, processor.tokenizer.pad_token_id}: | |
| break | |
| token = torch.tensor([[token_id]], dtype=torch.long, device="cuda") | |
| attention_mask = torch.cat([attention_mask, torch.ones_like(token)], dim=1) | |
| outputs = text_model( | |
| input_ids=token, | |
| attention_mask=attention_mask, | |
| position_ids=next_position, | |
| past_key_values=cache, | |
| use_cache=True, | |
| ) | |
| cache = outputs.past_key_values | |
| logits = ar_model.lm_head(outputs.last_hidden_state[:, -1, :])[0] | |
| next_position = next_position + 1 | |
| decoded = processor.tokenizer.decode( | |
| continuation, skip_special_tokens=False, clean_up_tokenization_spaces=False | |
| ) | |
| result = { | |
| "sample_id": trajectory["sample_id"], | |
| "generated_tokens": continuation, | |
| "decoded": decoded, | |
| "emitted_next_image_gen": marker_id in continuation, | |
| } | |
| path = args.run_dir / "evaluation" / "post_image_continuation.json" | |
| path.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") | |
| print(json.dumps(result, ensure_ascii=False), flush=True) | |
| return result | |
| 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("--image-size", type=int, default=256) | |
| parser.add_argument("--max-new-tokens", type=int, default=512) | |
| return parser | |
| def main() -> None: | |
| continue_after_visual(build_parser().parse_args()) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.87 kB
- Xet hash:
- b32606201def4975cd266981b5e5d616b242ee48cbbaec19b2bc463960268cfb
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.