In this blog we will cover various methods that are employed for fast inference on diffusion models
For gated repo make sure you have HuggingFace Token for access, Like FLUX is
Now select a gpu from marketplace, for our purpose a 48GB VRAM gpu will work: https://console.aquanode.io/marketplace?vram=48GB&computeType=all
It's simple process, just have a ssh key, put it etc. A straight process, but if you still need help:
You can checkout here how to deploy a VM https://docs.aquanode.io/docs/virtual-machines/running-vm
Now in the VM
export HF_TOKEN=your-hf-token
Setup development cuda: nvcc; if not installed on your VM.
For ubuntu22/24 this will work
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
sudo apt install -y cuda-toolkit-12-8
echo 'export PATH=/usr/local/cuda-12.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
# Verify
nvcc --version
Install uv
wget -qO- https://astral.sh/uv/install.sh | sh
For our VM of cuda12.8, torch 2.9.0 will be faster as it's wheels are available, you can go with updated version of torch,if wheels available for that, otherwise it will automatically build yours(sometimes take time).
uv venv
source .venv/bin/activate
# Number of cpus to allocate for flash attn build, more will use more RAM, refrain from using more than 32, even in high config VM
export MAX_JOBS=8
export CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -n1)
export TORCH_CUDA_ARCH_LIST="${CUDA_ARCH}"
export FLASH_ATTN_CUDA_ARCHS="${CUDA_ARCH}"
uv pip install torch==2.9.0 psutils setuptools ninja huggingface_hub
uv pip install "xfuser[flash-attn]" --no-build-isolation
uv pip install fastapi pydantic ray uvicorn
git clone https://github.com/huggingface/diffusers
cd diffusers
uv pip install -e .
cd ..
Save the following as flux_xfuser.py. It wraps the xDiT pipeline in a small
FastAPI service with a single POST /generate endpoint, returning the image
either as base64 or written to disk.
import os
import io
import time
import base64
import torch
import ray
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from transformers import T5EncoderModel
import uvicorn
from xfuser import xFuserArgs, xFuserFluxPipeline
from xfuser.config import FlexibleArgumentParser
from xfuser.ray.pipeline.pipeline_utils import RayDiffusionPipeline
class GenerateRequest(BaseModel):
prompt: str
num_inference_steps: Optional[int] = 28
seed: Optional[int] = 42
height: Optional[int] = 1024
width: Optional[int] = 1024
guidance_scale: Optional[float] = 0.0
max_sequence_length: Optional[int] = 256
save_disk_path: Optional[str] = None
class Config:
json_schema_extra = {
"example": {
"prompt": "a beautiful landscape",
"num_inference_steps": 28,
"seed": 42,
"height": 1024,
"width": 1024,
"guidance_scale": 0.0,
"max_sequence_length": 256,
}
}
app = FastAPI(title="xDiT Flux HTTP Service")
class Engine:
def __init__(self, engine_args: xFuserArgs):
if not ray.is_initialized():
ray.init(runtime_env={
"env_vars": {
"MASTER_ADDR": "localhost",
"MASTER_PORT": "29500",
# "NCCL_P2P_LEVEL": "PHB", # Check if works without this
"LOCAL_RANK": "0",
}
})
engine_args.use_ray = True
engine_config, input_config = engine_args.create_config()
engine_config.runtime_config.dtype = torch.bfloat16
# Load T5 encoder inside each Ray worker (same pattern as ray_flux_example.py)
encoder_kwargs = {
"text_encoder_2": {
"model_class": T5EncoderModel,
"pretrained_model_name_or_path": engine_config.model_config.model,
"subfolder": "text_encoder_2",
"torch_dtype": torch.bfloat16,
}
}
self.pipe = RayDiffusionPipeline.from_pretrained(
PipelineClass=xFuserFluxPipeline,
pretrained_model_name_or_path=engine_config.model_config.model,
engine_config=engine_config,
torch_dtype=torch.bfloat16,
**encoder_kwargs,
)
self.pipe.prepare_run(input_config)
async def generate(self, request: GenerateRequest):
start_time = time.time()
results = self.pipe(
height=request.height,
width=request.width,
prompt=request.prompt,
num_inference_steps=request.num_inference_steps,
output_type="pil",
max_sequence_length=request.max_sequence_length,
guidance_scale=request.guidance_scale,
generator=torch.Generator(device="cuda").manual_seed(request.seed),
)
elapsed = time.time() - start_time
# Only the dp-last-group worker returns images; others return None
image = None
for images in results:
if images is not None:
image = images[0]
break
if image is None:
raise RuntimeError("No worker returned an image.")
if request.save_disk_path:
os.makedirs(request.save_disk_path, exist_ok=True)
filename = f"generated_image_{time.strftime('%Y%m%d-%H%M%S')}.png"
path = os.path.join(request.save_disk_path, filename)
image.save(path)
return {
"message": "Image generated successfully",
"elapsed_time": f"{elapsed:.2f}s",
"output": path,
"save_to_disk": True,
}
else:
buf = io.BytesIO()
image.save(buf, format="PNG")
img_b64 = base64.b64encode(buf.getvalue()).decode()
return {
"message": "Image generated successfully",
"elapsed_time": f"{elapsed:.2f}s",
"output": img_b64,
"save_to_disk": False,
}
@app.post("/generate")
async def generate_image(request: GenerateRequest):
if not request.prompt:
raise HTTPException(status_code=400, detail="prompt cannot be empty")
if request.height <= 0 or request.width <= 0:
raise HTTPException(status_code=400, detail="height and width must be positive")
if request.num_inference_steps <= 0:
raise HTTPException(status_code=400, detail="num_inference_steps must be positive")
try:
return await engine.generate(request)
except Exception as e:
if isinstance(e, HTTPException):
raise
raise HTTPException(status_code=500, detail=str(e))
if __name__ == "__main__":
parser = FlexibleArgumentParser(description="xDiT FLUX HTTP Service")
xFuserArgs.add_cli_args(parser)
parser.add_argument("--host", type=str, default="0.0.0.0")
parser.add_argument("--port", type=int, default=6000)
args = parser.parse_args()
engine_args = xFuserArgs.from_cli_args(args)
engine = Engine(engine_args=engine_args)
uvicorn.run(app, host=args.host, port=args.port)
For a single-GPU server (48GB VRAM), run:
python flux_xfuser.py --model black-forest-labs/FLUX.1-dev --ray_world_size 1 --use_teacache --use_fbcache
OR, to also compile the model (slower first request, faster afterwards):
python flux_xfuser.py --model black-forest-labs/FLUX.1-dev --ray_world_size 1 --use_teacache --use_fbcache --use_torch_compile
If you are going to use it as inference server, you can just use startup scripts at https://console.aquanode.io/workloads/startup-scripts
Inference time
curl -X POST http://localhost:6000/generate \
-H "Content-Type: application/json" \
-d '{
"prompt": "a beautiful mountain landscape at sunset",
"num_inference_steps": 28,
"height": 1024,
"width": 1024,
"seed": 42,
"save_disk_path": "/tmp/outputs"
}'
Not providing save_disk_path, will default the api to return base64 encoded image
Results like:
Have more traffic? Just rent a multi-gpu VM follow same process and instead run
python flux_xfuser.py --model black-forest-labs/FLUX.1-dev --ray_world_size 1 --use_teacache --use_fbcache --data_parallel_degree $NUM_GPU
And now server will distribute the requests across GPUs
