import torch
from diffusers import Lumina2Pipeline


MODEL_ID = "issai/Beynele"
PROMPTS = [
    "A Kazakh dombra resting on a patterned felt carpet.",
    "A cinematic aerial photo of Astana's Baiterek Tower at golden hour.",
    'The Kazakh Cyrillic word "бейнеле" sculpted from soft white clouds in a bright blue sky.',
]


def load_pipeline():
    pipe = Lumina2Pipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
    )
    pipe.enable_model_cpu_offload()
    return pipe


if __name__ == "__main__":
    pipe = load_pipeline()
    for idx, prompt in enumerate(PROMPTS, start=1):
        image = pipe(
            prompt,
            height=1024,
            width=1024,
            guidance_scale=4.0,
            num_inference_steps=40,
            cfg_trunc_ratio=0.25,
            cfg_normalization=True,
            generator=torch.Generator("cpu").manual_seed(42 + idx),
        ).images[0]
        image.save(f"beynele_example_{idx}.png")