|
|
import os |
|
|
from typing import Any, Dict |
|
|
from PIL import Image |
|
|
import torch |
|
|
from diffusers import FluxPipeline |
|
|
from huggingface_inference_toolkit.logging import logger |
|
|
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe |
|
|
import time |
|
|
|
|
|
|
|
|
class EndpointHandler: |
|
|
def __init__(self, path=""): |
|
|
self.pipe = FluxPipeline.from_pretrained( |
|
|
"NoMoreCopyrightOrg/flux-dev", |
|
|
torch_dtype=torch.bfloat16, |
|
|
).to("cuda") |
|
|
apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.12) |
|
|
self.pipe.transformer = torch.compile( |
|
|
self.pipe.transformer, mode="max-autotune-no-cudagraphs", |
|
|
) |
|
|
self.pipe.vae = torch.compile( |
|
|
self.pipe.vae, mode="max-autotune-no-cudagraphs", |
|
|
) |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> str: |
|
|
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", 28) |
|
|
width = parameters.get("width", 1024) |
|
|
height = parameters.get("height", 1024) |
|
|
guidance_scale = parameters.get("guidance_scale", 3.5) |
|
|
|
|
|
|
|
|
seed = parameters.get("seed", 0) |
|
|
generator = torch.manual_seed(seed) |
|
|
start_time = time.time() |
|
|
result = self.pipe( |
|
|
prompt, |
|
|
height=height, |
|
|
width=width, |
|
|
guidance_scale=guidance_scale, |
|
|
num_inference_steps=num_inference_steps, |
|
|
generator=generator, |
|
|
).images[0] |
|
|
end_time = time.time() |
|
|
time_taken = end_time - start_time |
|
|
print(f"Time taken: {time_taken:.2f} seconds") |
|
|
return result |
|
|
|
