|
import os |
|
from typing import Any, Dict |
|
|
|
from diffusers import DiffusionPipeline |
|
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: |
|
"""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 = 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] |
