"""Inference helpers for the native Diffusers FLUX.2 Klein pipeline."""

from __future__ import annotations

import torch
from diffusers import Flux2KleinPipeline


def load_pipeline(model_path, device="cuda", dtype=torch.bfloat16):
    pipe = Flux2KleinPipeline.from_pretrained(
        model_path,
        torch_dtype=dtype,
        local_files_only=True,
    )
    return pipe.to(device)


def generate(pipe, prompt, *, height=512, width=512, steps=20, cfg=4.0, seed=0):
    generator = torch.Generator(device=pipe.device).manual_seed(seed)
    return pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=steps,
        guidance_scale=cfg,
        generator=generator,
    ).images[0]