| 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 |
|
|
