diff --git a/hemm/models/qwen_vl.py b/hemm/models/qwen_vl.py new file mode 100644 index 0000000..edf6434 --- /dev/null +++ b/hemm/models/qwen_vl.py @@ -0,0 +1,34 @@ +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.generation import GenerationConfig +import torch +torch.manual_seed(1234) + +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-VL", trust_remote_code=True) + +# use bf16 +# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL", device_map="auto", trust_remote_code=True, bf16=True).eval() +# use fp16 +# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL", device_map="auto", trust_remote_code=True, fp16=True).eval() +# use cpu only +# model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL", device_map="cpu", trust_remote_code=True).eval() +# use cuda device +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL", device_map="cuda", trust_remote_code=True).eval() + +# Specify hyperparameters for generation (No need to do this if you are using transformers>=4.32.0) +# model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-VL", trust_remote_code=True) + +query = tokenizer.from_list_format([ + {'image': 'https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg'}, + {'text': 'Generate the caption in English with grounding:'}, +]) +inputs = tokenizer(query, return_tensors='pt') +inputs = inputs.to(model.device) +pred = model.generate(**inputs) +response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False) +print(response) +# https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpegGenerate the caption in English with grounding: Woman(451,379),(731,806) and her dog(219,424),(576,896) playing on the beach<|endoftext|> +image = tokenizer.draw_bbox_on_latest_picture(response) +if image: + image.save('2.jpg') +else: + print("no box")