Spaces:
Running
Running
| import argparse | |
| import json | |
| import os | |
| import torch | |
| from PIL import Image | |
| from transformers import AutoTokenizer | |
| from .rope import precompute_freqs_cis | |
| from .text import lm_head, text_decoder, text_encoder | |
| from .vision import encode_image | |
| from .weights import load_from_pt, load_from_safetensors | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--image", "-i", type=str, required=True) | |
| parser.add_argument("--prompt", "-p", type=str, required=True) | |
| parser.add_argument("--model", "-m", type=str, required=True) | |
| parser.add_argument("--config", "-c", type=str, default="{}") | |
| parser.add_argument("--max-tokens", "-t", type=int, default=200) | |
| parser.add_argument("--sampler", "-s", type=str, default="greedy") | |
| args = parser.parse_args() | |
| if torch.cuda.is_available(): | |
| torch.set_default_device("cuda") | |
| elif torch.backends.mps.is_available(): | |
| torch.set_default_device("mps") | |
| # Load config. | |
| config = json.loads(args.config) | |
| text_n_heads = config.get("text_n_heads", 32) | |
| # Load model. | |
| model_path = args.model | |
| if not os.path.exists(model_path): | |
| raise FileNotFoundError(f"Model not found at {model_path}") | |
| if model_path.endswith(".pt"): | |
| model = load_from_pt(model_path, **config) | |
| elif model_path.endswith(".safetensors"): | |
| model = load_from_safetensors(model_path, **config) | |
| else: | |
| raise ValueError(f"Invalid model format: {model_path}") | |
| # Encode image. | |
| image_path = args.image | |
| if not os.path.exists(image_path): | |
| raise FileNotFoundError(f"Image not found at {image_path}") | |
| image = Image.open(image_path) | |
| image = image.resize((378, 378)) | |
| image_tensor = encode_image(image, model.vision) | |
| # Encode text, and create inputs_embeds. | |
| tokenizer = AutoTokenizer.from_pretrained("vikhyatk/moondream2") | |
| prompt = f"\n\nQuestion: {args.prompt}\n\nAnswer:" | |
| input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"] | |
| input_ids = torch.cat([torch.tensor([[tokenizer.eos_token_id]]), input_ids], dim=1) | |
| inputs_embeds = text_encoder(input_ids, model.text) | |
| inputs_embeds = torch.cat( | |
| [ | |
| inputs_embeds[:, 0:1, :], | |
| image_tensor.unsqueeze(0), | |
| inputs_embeds[:, 1:, :], | |
| ], | |
| dim=1, | |
| ) | |
| kv_cache = torch.empty(24, 2, 1, text_n_heads, 2048, 64, dtype=torch.float16) | |
| freqs_cis = precompute_freqs_cis(32, 2048) | |
| pos = 0 | |
| for _ in range(args.max_tokens): | |
| with torch.no_grad(): | |
| hidden, kv_cache_update = text_decoder( | |
| inputs_embeds, model.text, kv_cache[:, :, :, :, :pos, :], freqs_cis | |
| ) | |
| logits = lm_head(hidden, model.text) | |
| kv_cache[:, :, :, :, pos : pos + kv_cache_update.size(-2), :] = ( | |
| kv_cache_update | |
| ) | |
| pos += kv_cache_update.size(-2) | |
| if args.sampler == "multinomial": | |
| next_token = torch.multinomial( | |
| torch.softmax(logits, dim=-1), num_samples=1 | |
| ).squeeze(0) | |
| elif args.sampler == "greedy": | |
| next_token = torch.argmax(logits, dim=-1) | |
| else: | |
| raise ValueError(f"Invalid sampler: {args.sampler}") | |
| if next_token == tokenizer.eos_token_id: | |
| print() | |
| break | |
| input_ids = next_token.unsqueeze(0) | |
| inputs_embeds = text_encoder(input_ids, model.text) | |
| output_text = tokenizer.batch_decode(input_ids)[0] | |
| print(output_text, end="", flush=True) | |