|
import os |
|
from typing import Any, Dict |
|
|
|
from diffusers import FluxPipeline, FluxTransformer2DModel |
|
from torchao.quantization import int8_weight_only, quantize_ |
|
from PIL.Image import Image |
|
import torch |
|
|
|
from huggingface_inference_toolkit.logging import logger |
|
|
|
class EndpointHandler: |
|
def __init__(self, **kwargs: Any) -> None: |
|
repo_id = "camenduru/FLUX.1-dev-diffusers" |
|
dtype = torch.bfloat16 |
|
transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype) |
|
quantize_(transformer, int8_weight_only(), device="cuda") |
|
transformer.to(memory_format=torch.channels_last) |
|
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True) |
|
self.pipeline = FluxPipeline.from_pretrained(repo_id, transformer=transformer, torch_dtype=torch.bfloat16).to("cuda") |
|
self.pipeline.vae.to(memory_format=torch.channels_last) |
|
self.pipeline.vae.decode = torch.compile(self.pipeline.vae.decode, mode="max-autotune", fullgraph=True) |
|
|
|
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 = parameters.get("seed", 0) |
|
generator = torch.manual_seed(seed) |
|
|
|
return self.pipeline( |
|
prompt, |
|
height=height, |
|
width=width, |
|
guidance_scale=guidance_scale, |
|
num_inference_steps=num_inference_steps, |
|
generator=generator, |
|
).images[0] |
