Initial import

a19a224 unverified over 2 years ago

971 Bytes

from typing import Dict, List, Any
import torch
from transformers import BitsAndBytesConfig, pipeline

device = "cuda" if torch.cuda.is_available() else "cpu"


class EndpointHandler():
    def __init__(self, path=""):
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16
        )
        # self.pipeline = pipeline("image-to-text", model="llava-hf/llava-1.5-7b-hf",
        #                          model_kwargs={"quantization_config": quantization_config})

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        print(data)
        inputs = data.pop("inputs", data)
        image = data.pop("image", None)
        prompt = data.pop("prompt", None)
        # outputs = self.pipeline(image, prompt=prompt, generate_kwargs={
        # "max_new_tokens": 200})
        # return {"caption": outputs[0]["generated_text"]}
        return {"image": image, "prompt": prompt}

	from typing import Dict, List, Any
	import torch
	from transformers import BitsAndBytesConfig, pipeline

	device = "cuda" if torch.cuda.is_available() else "cpu"


	class EndpointHandler():
	def __init__(self, path=""):
	quantization_config = BitsAndBytesConfig(
	load_in_4bit=True,
	bnb_4bit_compute_dtype=torch.float16
	)
	# self.pipeline = pipeline("image-to-text", model="llava-hf/llava-1.5-7b-hf",
	# model_kwargs={"quantization_config": quantization_config})

	def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
	print(data)
	inputs = data.pop("inputs", data)
	image = data.pop("image", None)
	prompt = data.pop("prompt", None)
	# outputs = self.pipeline(image, prompt=prompt, generate_kwargs={
	# "max_new_tokens": 200})
	# return {"caption": outputs[0]["generated_text"]}
	return {"image": image, "prompt": prompt}