trq54y654y

gistfile1.txt

import torch
import sys
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoImageProcessor


def load_model_auto(model_path):
    dtype = torch.float16
    model = AutoModelForCausalLM.from_pretrained(
        model_path, torch_dtype=dtype, trust_remote_code=True, device_map="cuda"
    ).eval()
    model.to("cuda:0")
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    image_processor = AutoImageProcessor.from_pretrained(
        model_path, trust_remote_code=True
    )
    return tokenizer, model, image_processor


if __name__ == "__main__":
    load_model_auto(sys.argv[1])
    print('model loaded!')