from pathlib import Path

import torch
from diffusers import AutoModel, ErnieImagePipeline
from transformers import AutoModel as TransformersAutoModel
from transformers import AutoModelForCausalLM


model_dir = Path(r"./")
output_image = model_dir / "demo_output.png"
prompt = "一只橙色的小猫坐在木桌上，前面放着一张写有 Hello Ernie-Image NF4 的白纸，柔和自然光，写实风格。"

if not model_dir.exists():
    raise FileNotFoundError(f"Model dir not found: {model_dir}")
if not torch.cuda.is_available():
    raise RuntimeError("CUDA is required to run this demo.")

transformer = AutoModel.from_pretrained(
    str(model_dir / "transformer"),
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    local_files_only=True,
)

text_encoder = TransformersAutoModel.from_pretrained(
    str(model_dir / "text_encoder"),
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    local_files_only=True,
)

pe = AutoModelForCausalLM.from_pretrained(
    str(model_dir / "pe"),
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    local_files_only=True,
)

pipe = ErnieImagePipeline.from_pretrained(
    str(model_dir),
    transformer=transformer,
    text_encoder=text_encoder,
    pe=pe,
    torch_dtype=torch.bfloat16,
    local_files_only=True,
)
pipe.vae.to("cuda", dtype=torch.bfloat16)

image = pipe(
    prompt=prompt,
    height=512,
    width=512,
    num_inference_steps=50,
    guidance_scale=4.0,
    use_pe=True,
    generator=torch.Generator(device="cuda").manual_seed(1234),
).images[0]

image.save(output_image)
print(f"Saved image to: {output_image}")