|
--- |
|
license: apache-2.0 |
|
base_model: |
|
- rhymes-ai/Aria |
|
--- |
|
|
|
This repository offers int8 quantized weights of the [rhymes-ai/Aria](https://huggingface.co/rhymes-ai/Aria) model utilizing the [TorchAO](https://github.com/pytorch/ao) quantization framework. It now supports inference within 30GB of GPU memory. |
|
|
|
|
|
## Quick Start |
|
### Installation |
|
``` |
|
pip install transformers==4.45.0 accelerate==0.34.1 sentencepiece==0.2.0 torch==2.5.0 torchao==0.6.1 torchvision requests Pillow |
|
pip install flash-attn --no-build-isolation |
|
``` |
|
|
|
### Inference |
|
|
|
```python |
|
import requests |
|
import torch |
|
from PIL import Image |
|
from transformers import AutoModelForCausalLM, AutoProcessor |
|
|
|
model_id_or_path = "rhymes-ai/Aria-torchao-int8wo" |
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_id_or_path, |
|
device_map="auto", |
|
torch_dtype=torch.bfloat16, |
|
trust_remote_code=True, |
|
attn_implementation="flash_attention_2", |
|
) |
|
|
|
processor = AutoProcessor.from_pretrained(model_id_or_path, trust_remote_code=True) |
|
|
|
image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" |
|
|
|
image = Image.open(requests.get(image_path, stream=True).raw) |
|
|
|
messages = [ |
|
{ |
|
"role": "user", |
|
"content": [ |
|
{"text": None, "type": "image"}, |
|
{"text": "what is the image?", "type": "text"}, |
|
], |
|
} |
|
] |
|
|
|
text = processor.apply_chat_template(messages, add_generation_prompt=True) |
|
inputs = processor(text=text, images=image, return_tensors="pt") |
|
inputs["pixel_values"] = inputs["pixel_values"].to(model.dtype) |
|
inputs = {k: v.to(model.device) for k, v in inputs.items()} |
|
|
|
with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.bfloat16): |
|
output = model.generate( |
|
**inputs, |
|
max_new_tokens=500, |
|
stop_strings=["<|im_end|>"], |
|
tokenizer=processor.tokenizer, |
|
do_sample=True, |
|
temperature=0.9, |
|
) |
|
output_ids = output[0][inputs["input_ids"].shape[1] :] |
|
result = processor.decode(output_ids, skip_special_tokens=True) |
|
|
|
print(result) |
|
``` |