File size: 5,167 Bytes
7470c71 f6f8735 01fb00e f6f8735 d597361 7470c71 927af23 7470c71 927af23 01fb00e 7470c71 927af23 7470c71 f6f8735 927af23 f6f8735 927af23 f6f8735 28306cb f6f8735 01fb00e 321843f 01fb00e 927af23 f6f8735 321843f 01fb00e 321843f f6f8735 321843f 01fb00e 321843f 01fb00e 321843f 01fb00e f6f8735 01fb00e f6f8735 01fb00e f6f8735 01fb00e f6f8735 01fb00e f6f8735 01fb00e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import subprocess
import sys
import torch
import base64
from io import BytesIO
from PIL import Image
import requests
from transformers import AutoModelForCausalLM, AutoProcessor
def install(package):
subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-warn-script-location", package])
class EndpointHandler:
def __init__(self, path=""):
required_packages = ['timm', 'einops', 'flash-attn', 'Pillow']
for package in required_packages:
try:
install(package)
print(f"Successfully installed {package}")
except Exception as e:
print(f"Failed to install {package}: {str(e)}")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {self.device}")
self.model_name = "microsoft/Florence-2-base-ft"
self.model = AutoModelForCausalLM.from_pretrained(
self.model_name,
trust_remote_code=True,
revision='refs/pr/6'
).to(self.device)
self.processor = AutoProcessor.from_pretrained(
self.model_name,
trust_remote_code=True,
revision='refs/pr/6'
)
if torch.cuda.is_available():
torch.cuda.empty_cache()
def process_image(self, image_path):
try:
with open(image_path, 'rb') as image_file:
image = Image.open(image_file)
return image
except Exception as e:
print(f"Error processing image: {str(e)}")
return None
def __call__(self, data):
try:
# Extract inputs from the expected Hugging Face format
inputs = data.pop("inputs", data)
# Check if inputs is a dict or string
if isinstance(inputs, dict):
image_path = inputs.get("image", None)
text_input = inputs.get("text", "")
else:
# If inputs is not a dict, assume it's the image path
image_path = inputs
text_input = "What is in this image?"
# Process image
image = self.process_image(image_path) if image_path else None
# Prepare inputs for the model
model_inputs = self.processor(
images=image if image else None,
text=text_input,
return_tensors="pt"
)
# Move inputs to device
model_inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v
for k, v in model_inputs.items()}
# Generate output
with torch.no_grad():
outputs = self.model.generate(**model_inputs)
# Decode outputs
decoded_outputs = self.processor.batch_decode(outputs, skip_special_tokens=True)
return {"generated_text": decoded_outputs[0]}
except Exception as e:
return {"error": str(e)}
# import subprocess
# import sys
# import torch
# from transformers import AutoModelForCausalLM, AutoProcessor
# def install(package):
# subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-warn-script-location", package])
# class EndpointHandler:
# def __init__(self, path=""):
# required_packages = ['timm', 'einops', 'flash-attn']
# for package in required_packages:
# try:
# install(package)
# print(f"Successfully installed {package}")
# except Exception as e:
# print(f"Failed to install {package}: {str(e)}")
# self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# print(f"Using device: {self.device}")
# self.model_name = "microsoft/Florence-2-base-ft"
# self.model = AutoModelForCausalLM.from_pretrained(
# self.model_name,
# trust_remote_code=True,
# revision='refs/pr/6'
# ).to(self.device)
# self.processor = AutoProcessor.from_pretrained(
# self.model_name,
# trust_remote_code=True,
# revision='refs/pr/6'
# )
# if torch.cuda.is_available():
# torch.cuda.empty_cache()
# def __call__(self, data):
# try:
# inputs = data.pop("inputs", data)
# processed_inputs = self.processor(inputs, return_tensors="pt")
# processed_inputs = {k: v.to(self.device) for k, v in processed_inputs.items()}
# with torch.no_grad():
# outputs = self.model.generate(**processed_inputs)
# decoded_outputs = self.processor.batch_decode(outputs, skip_special_tokens=True)
# return {"outputs": decoded_outputs}
# except Exception as e:
# return {"error": str(e)} |