#### inference code drived from https://civitai.com/models/833294 and https://huggingface.co/spaces/nyanko7/toaru-xl-model

import torch
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
from PIL import Image
import random

# 初始化 pipe，只执行一次
ckpt_path = "nyaflow-xl-alpha.safetensors"  # https://huggingface.co/nyanko7/nyaflow-xl-alpha
ckpt_path = "noobaiXLNAIXL_vPred10Version.safetensors" #### https://civitai.com/models/833294
pipe = StableDiffusionXLPipeline.from_single_file(
    ckpt_path,
    use_safetensors=True,
    torch_dtype=torch.float16,
)
scheduler_args = {"prediction_type": "v_prediction", "rescale_betas_zero_snr": True}
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, **scheduler_args)
pipe.enable_xformers_memory_efficient_attention()
pipe = pipe.to("cuda")

# 定义默认参数
PRESET_Q = "year_2022, best quality, high quality, very aesthetic"
NEGATIVE_PROMPT = "lowres, worst quality, displeasing, bad anatomy, text, error, extra digit, cropped, error, fewer, extra, missing, worst quality, jpeg artifacts, censored, ai-generated worst quality displeasing, bad quality"

def generate_image(
    prompt: str,
    preset: str = PRESET_Q,
    height: int = 1216,
    width: int = 832,
    negative_prompt: str = NEGATIVE_PROMPT,
    guidance_scale: float = 4.0,
    randomize_seed: bool = True,
    seed: int = 42,
    inference_steps: int = 25,
) -> Image:
    # 合并 prompt 和 preset
    prompt = prompt.strip() + ", " + preset.strip()
    negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None

    # 随机化种子
    if randomize_seed:
        seed = random.randint(0, 9007199254740991)

    # 设置生成器
    generator = torch.Generator(device="cuda").manual_seed(seed)

    # 限制推理步数
    if inference_steps > 50:
        inference_steps = 50

    # 生成图像
    image = pipe(
        prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        generator=generator,
        num_inference_steps=inference_steps
    ).images[0]

    return image


# 示例调用
if __name__ == "__main__":
    prompt = "zhongli"
    image = generate_image(prompt)
    image

    prompt = "Neuvillette"
    image = generate_image(prompt)
    image