Views
No views yet

1import torch
2from PIL import Image
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5MID = "SVECTOR-CORPORATION/Fal-2-500M"
6IMAGE_TOKEN_INDEX = -200
7tok = AutoTokenizer.from_pretrained(MID, trust_remote_code=True)
8model = AutoModelForCausalLM.from_pretrained(
9 MID,
10 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
11 device_map="auto",
12 trust_remote_code=True,
13)
14messages = [
15 {"role": "user", "content": "<image>\nDescribe me this image."}
16]
17rendered = tok.apply_chat_template(
18 messages, add_generation_prompt=True, tokenize=False
19)
20
21pre, post = rendered.split("<image>", 1)
22pre_ids = tok(pre, return_tensors="pt", add_special_tokens=False).input_ids
23post_ids = tok(post, return_tensors="pt", add_special_tokens=False).input_ids
24
25img_tok = torch.tensor([[IMAGE_TOKEN_INDEX]], dtype=pre_ids.dtype)
26input_ids = torch.cat([pre_ids, img_tok, post_ids], dim=1).to(model.device)
27attention_mask = torch.ones_like(input_ids, device=model.device)
28
29img = Image.open("photo.jpg").convert("RGB")
30px = model.get_vision_tower().image_processor(images=img, return_tensors="pt")["pixel_values"]
31px = px.to(model.device, dtype=model.dtype)
32
33# Generate
34with torch.no_grad():
35 out = model.generate(
36 inputs=input_ids,
37 attention_mask=attention_mask,
38 images=px,
39 max_new_tokens=128,
40 )
41
42print(tok.decode(out[0], skip_special_tokens=True))