mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-02-02 20:07:53 +00:00
67 lines
2.8 KiB
Python
67 lines
2.8 KiB
Python
from transformers import CLIPImageProcessor, BitsAndBytesConfig, AutoTokenizer
|
|
|
|
from .caption import default_long_prompt, default_short_prompt, default_replacements, clean_caption
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
class FuyuImageProcessor:
|
|
def __init__(self, device='cuda'):
|
|
from transformers import FuyuProcessor, FuyuForCausalLM
|
|
self.device = device
|
|
self.model: FuyuForCausalLM = None
|
|
self.processor: FuyuProcessor = None
|
|
self.dtype = torch.bfloat16
|
|
self.tokenizer: AutoTokenizer
|
|
self.is_loaded = False
|
|
|
|
def load_model(self):
|
|
from transformers import FuyuProcessor, FuyuForCausalLM
|
|
model_path = "adept/fuyu-8b"
|
|
kwargs = {"device_map": self.device}
|
|
kwargs['load_in_4bit'] = True
|
|
kwargs['quantization_config'] = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_compute_dtype=self.dtype,
|
|
bnb_4bit_use_double_quant=True,
|
|
bnb_4bit_quant_type='nf4'
|
|
)
|
|
self.processor = FuyuProcessor.from_pretrained(model_path)
|
|
self.model = FuyuForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
|
|
self.is_loaded = True
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
self.model = FuyuForCausalLM.from_pretrained(model_path, torch_dtype=self.dtype, **kwargs)
|
|
self.processor = FuyuProcessor(image_processor=FuyuImageProcessor(), tokenizer=self.tokenizer)
|
|
|
|
def generate_caption(
|
|
self, image: Image,
|
|
prompt: str = default_long_prompt,
|
|
replacements=default_replacements,
|
|
max_new_tokens=512
|
|
):
|
|
# prepare inputs for the model
|
|
# text_prompt = f"{prompt}\n"
|
|
|
|
# image = image.convert('RGB')
|
|
model_inputs = self.processor(text=prompt, images=[image])
|
|
model_inputs = {k: v.to(dtype=self.dtype if torch.is_floating_point(v) else v.dtype, device=self.device) for k, v in
|
|
model_inputs.items()}
|
|
|
|
generation_output = self.model.generate(**model_inputs, max_new_tokens=max_new_tokens)
|
|
prompt_len = model_inputs["input_ids"].shape[-1]
|
|
output = self.tokenizer.decode(generation_output[0][prompt_len:], skip_special_tokens=True)
|
|
output = clean_caption(output, replacements=replacements)
|
|
return output
|
|
|
|
# inputs = self.processor(text=text_prompt, images=image, return_tensors="pt")
|
|
# for k, v in inputs.items():
|
|
# inputs[k] = v.to(self.device)
|
|
|
|
# # autoregressively generate text
|
|
# generation_output = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
|
|
# generation_text = self.processor.batch_decode(generation_output[:, -max_new_tokens:], skip_special_tokens=True)
|
|
# output = generation_text[0]
|
|
#
|
|
# return clean_caption(output, replacements=replacements)
|