Views
No views yet
1from peft import PeftModel
2from transformers import AutoProcessor, LlavaForConditionalGeneration
3import torch
4
5# Load base model
6base_model = LlavaForConditionalGeneration.from_pretrained(
7 "llava-hf/llava-1.5-7b-hf",
8 torch_dtype=torch.bfloat16
9).to('cuda')
10processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
11
12# Load LoRA adapter
13model = PeftModel.from_pretrained(
14 base_model,
15 "ZinengTang/llava-lora-spatial"
16)
17
18from PIL import Image
19init_prompt_instruct = "Describe the location of the blue sphere relative to the environment features."
20conversation = [
21 {
22 "role": "user",
23 "content": [
24 {"type": "text", "text": init_prompt_instruct},
25 {"type": "image"}, # This will be replaced with the actual image
26 ],
27 },
28]
29speaker_image = Image.open('your_image_path')
30prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
31# print(prompt)
32# Process the input image and prompt
33inputs = processor(
34 images=speaker_image,
35 text=prompt,
36 return_tensors="pt",
37 max_length=256,
38).to('cuda')
39
40with torch.no_grad():
41 generated = model.generate(
42 input_ids=inputs["input_ids"],
43 attention_mask=inputs["attention_mask"],
44 pixel_values=inputs["pixel_values"],
45 max_length=512,
46 num_beams=1,
47 do_sample=True,
48 temperature=0.7
49 )
50 generated_message = processor.batch_decode(
51 generated,
52 skip_special_tokens=True
53 )
54 print(generated_message)
55 generated_message = generated_message[0].split('ASSISTANT: ')[-1][:100]
56