🤗 Serve any model with Inference Endpoints + Custom Handlers

Serving LoRA Adapters for Diffusion Models

Say that you want to serve a fine-tuned LoRA adapter for a Diffusers model such as alvarobartt/ghibli-characters-flux-lora which is a LoRA fine-tune of black-forest-labs/FLUX.1-dev. When trying to deploy it on Inference Endpoints, the following error will show:

As the error says, you need to make sure that the model repository with the LoRA adapter contains a handler.py file that will load the model first and then the adapter, as explained in the Diffusers Documentation on How to load adapters.

Note that since the base model here i.e. not the adapter within the repository, is gated, you need to make sure that you create and set the HF_TOKEN environment variable value with a valid Hugging Face Hub token with read access over the gated model, in this case being black-forest-labs/FLUX.1-dev.

import os
from typing import Any, Dict

from diffusers import DiffusionPipeline  # type: ignore
from PIL.Image import Image
import torch

from huggingface_inference_toolkit.logging import logger


class EndpointHandler:
    def __init__(self, model_dir: str, **kwargs: Any) -> None:  # type: ignore
        """The current `EndpointHandler` works with any FLUX.1-dev LoRA Adapter."""
        if os.getenv("HF_TOKEN") is None:
            raise ValueError(
                "Since `black-forest-labs/FLUX.1-dev` is a gated model, you will need to provide a valid "
                "`HF_TOKEN` as an environment variable for the handler to work properly."
            )

        self.pipeline = DiffusionPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            torch_dtype=torch.bfloat16,
            token=os.getenv("HF_TOKEN"),
        )
        self.pipeline.load_lora_weights(model_dir)
        self.pipeline.to("cuda")

    def __call__(self, data: Dict[str, Any]) -> Image:
        logger.info(f"Received incoming request with {data=}")

        if "inputs" in data and isinstance(data["inputs"], str):
            prompt = data.pop("inputs")
        elif "prompt" in data and isinstance(data["prompt"], str):
            prompt = data.pop("prompt")
        else:
            raise ValueError(
                "Provided input body must contain either the key `inputs` or `prompt` with the"
                " prompt to use for the image generation, and it needs to be a non-empty string."
            )

        parameters = data.pop("parameters", {})

        num_inference_steps = parameters.get("num_inference_steps", 30)
        width = parameters.get("width", 1024)
        height = parameters.get("height", 768)
        guidance_scale = parameters.get("guidance_scale", 3.5)

        # seed generator (seed cannot be provided as is but via a generator)
        seed = parameters.get("seed", 0)
        generator = torch.manual_seed(seed)

        return self.pipeline(  # type: ignore
            prompt,
            height=height,
            width=width,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            generator=generator,
        ).images[0]

The code above can be reused and included as a handler.py file within any available LoRA adapter for black-forest-labs/FLUX.1-dev without any code modification required; and minimal modifications when changing the base model to e.g. stabilityai/stable-diffusion-3.5-large, as most of the code is shared among the different text-to-image use cases.

Which when deployed on Inference Endpoints looks like this:

Find the custom handler at alvarobartt/ghibli-characters-flux-lora.

Serving Models from the Hub Not Supported Natively

Say that you want to serve nvidia/NVLM-D-72B which is an image-text-to-text model, i.e. a Visual Language Model (VLM), that's not supported on Text Generation Inference (TGI), neither on the default PyTorch container (since image-text-to-text doesn't have a pre-defined AutoPipeline implementation for that task yet, but should soon have it as per https://github.com/huggingface/transformers/pull/34170).

Then you would need to define a custom handler that runs the pre-processing, inference, and post-processing for that task in the handler.py file; including any other requirement in requirements.txt, which shouldn't be needed in most of the cases, since the default PyTorch container already comes with most of the Hugging Face dependencies installed for Transformers, Sentence-Transformers and Diffusers; as well as some commonly used extra dependencies of those.

Note that using custom handlers, in this case, is not just to cover an unsupported model but also to define a custom device mapping, add custom pre-processing code, and add some custom logging messages, among other things.

import math
from typing import Any, Dict, List

import torch
import torchvision.transforms as T
from torchvision.transforms.functional import InterpolationMode

import requests
from io import BytesIO
from PIL import Image

from transformers import AutoTokenizer, AutoModel

from huggingface_inference_toolkit.logging import logger


def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float("inf")
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio


def dynamic_preprocess(
    image, min_num=1, max_num=12, image_size=448, use_thumbnail=False
):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height

    # calculate the existing image aspect ratio
    target_ratios = set(
        (i, j)
        for n in range(min_num, max_num + 1)
        for i in range(1, n + 1)
        for j in range(1, n + 1)
        if i * j <= max_num and i * j >= min_num
    )
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio,
        target_ratios,
        orig_width,
        orig_height,
        image_size,
    )

    # calculate the target width and height
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

    # resize the image
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size,
        )
        # split the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images


def load_image(image_url, input_size=448, max_num=12):
    response = requests.get(image_url)
    image = Image.open(BytesIO(response.content)).convert("RGB")
    transform = build_transform(input_size=input_size)
    images = dynamic_preprocess(
        image, image_size=input_size, use_thumbnail=True, max_num=max_num
    )
    pixel_values = [transform(image) for image in images]
    pixel_values = torch.stack(pixel_values)
    return pixel_values


def split_model():
    device_map = {}
    world_size = torch.cuda.device_count()
    num_layers = 80
    # Since the first GPU will be used for ViT, treat it as half a GPU.
    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
    num_layers_per_gpu = [num_layers_per_gpu] * world_size
    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
    layer_cnt = 0
    for i, num_layer in enumerate(num_layers_per_gpu):
        for j in range(num_layer):
            device_map[f"language_model.model.layers.{layer_cnt}"] = i
            layer_cnt += 1
    device_map["vision_model"] = 0
    device_map["mlp1"] = 0
    device_map["language_model.model.tok_embeddings"] = 0
    device_map["language_model.model.embed_tokens"] = 0
    device_map["language_model.output"] = 0
    device_map["language_model.model.norm"] = 0
    device_map["language_model.lm_head"] = 0
    device_map[f"language_model.model.layers.{num_layers - 1}"] = 0

    return device_map


IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose(
        [
            T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
            T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(mean=MEAN, std=STD),
        ]
    )
    return transform


class EndpointHandler:
    def __init__(self, model_dir: str, **kwargs: Any) -> None:
        self.model = AutoModel.from_pretrained(
            model_dir,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            use_flash_attn=False,
            trust_remote_code=True,
            device_map=split_model(),
        ).eval()

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_dir, trust_remote_code=True, use_fast=False
        )

    def __call__(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
        logger.info(f"Received incoming request with {data=}")
        
        if "instances" in data:
            logger.warning("Using `instances` instead of `inputs` is deprecated.")
            data["inputs"] = data.pop("instances")

        if "inputs" not in data:
            raise ValueError(
                "The request body must contain a key 'inputs' with a list of inputs."
            )

        if not isinstance(data["inputs"], list):
            raise ValueError(
                "The request inputs must be a list of dictionaries with either the key"
                " 'prompt' or 'prompt' + 'image_url', and optionally including the key"
                " 'generation_config'."
            )

        if not all(isinstance(input, dict) and "prompt" in input.keys() for input in data["inputs"]):
            raise ValueError(
                "The request inputs must be a list of dictionaries with either the key"
                " 'prompt' or 'prompt' + 'image_url', and optionally including the key"
                " 'generation_config'."
            )

        predictions = []
        for input in data["inputs"]:
            if "prompt" not in input:
                raise ValueError(
                    "The request input body must contain at least the key 'prompt' with the prompt to use."
                )

            generation_config = input.get("generation_config", dict(max_new_tokens=1024, do_sample=False))

            if "image_url" not in input:
                # pure-text conversation
                response, history = self.model.chat(
                    self.tokenizer,
                    None,
                    input["prompt"],
                    generation_config,
                    history=None,
                    return_history=True,
                )
            else:
                # single-image single-round conversation
                pixel_values = load_image(input["image_url"], max_num=6).to(
                    torch.bfloat16
                )
                response = self.model.chat(
                    self.tokenizer,
                    pixel_values,
                    f"<image>\n{input['prompt']}",
                    generation_config,
                )

            predictions.append(response)
        return {"predictions": predictions}

Which when deployed on Inference Endpoints looks like this:

Find the custom handler at alvarobartt/NVLM-D-72B-IE-compatible.

Defining Custom Specifications for I/O Payloads

Note that when using custom specifications for the I/O payloads, the "Task" that runs within the "Default" container within the Inference Endpoint needs to be set to "Custom", otherwise the playground in the UI will be created for that given task, which will fail due to the pre-defined output parsing; whilst the custom task will print out the raw response in JSON format.

Say that you have a UI or SDK that is expecting an API to receive or produce a given payload, but the default Inference Endpoints payload formatting for either or both input and output is not compliant with that, but you still want to leverage Hugging Face Inference Endpoints to use those within your application seamlessly.

Then you would need to implement a custom handler that given a task such as e.g. zero-shot-classification, expects an input different from the default one:

{"inputs": "I have a problem with my iphone that needs to be resolved asap!", "parameters": {"candidate_labels": ["urgent", "not urgent", "phone", "tablet", "computer"]}}

But you want it to expect the following:

{"sequence": "I have a problem with my iphone that needs to be resolved asap!", "labels": ["urgent", "not urgent", "phone", "tablet", "computer"]}

And by default producing the output:

{"sequence": "I have a problem with my iphone that needs to be resolved asap!!", "labels": ["urgent", "phone", "computer", "not urgent", "tablet"], "scores": [0.504, 0.479, 0.013, 0.003, 0.002]}

But you want it to produce:

{"sequence": "I have a problem with my iphone that needs to be resolved asap!!", "label": "urgent", "timestamp": 1732028280}

Then the custom handler would look similar to the following:

import os
from typing import Any, Dict
import time

from transformers import pipeline
import torch

from huggingface_inference_toolkit.logging import logger


class EndpointHandler:
    def __init__(self, model_dir: str, **kwargs: Any) -> None:
        """Initialize the EndpointHandler for zero-shot classification."""
        self.classifier = pipeline(
            "zero-shot-classification",
            model=model_dir,
            device=0 if torch.cuda.is_available() else -1
        )

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        logger.info(f"Received incoming request with {data=}")

        if "sequence" not in data or not isinstance(data["sequence"], str):
            raise ValueError(
                "Provided input body must contain the key `sequence` with the text to classify, "
                "and it needs to be a non-empty string."
            )

        if "labels" not in data or not isinstance(data["labels"], list):
            raise ValueError(
                "Provided input body must contain the key `labels` with a list of classification labels."
            )

        sequence = data["sequence"]
        labels = data["labels"]

        output = self.classifier(sequence, candidate_labels=labels)

        return {
            "sequence": sequence,
            "label": output["labels"][0],
            "timestamp": int(time.time())
        }