import os
import random
import numpy as np
import torch
from diffusers import ErnieImagePipeline
from tqdm import tqdm


seed = random.randint(0, 100000)
print(f"seed: {seed}")
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

torch.cuda.manual_seed_all(seed)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


pipe = ErnieImagePipeline.from_pretrained(
    "baidu/ERNIE-Image",
    torch_dtype=torch.bfloat16 
)


pipe = pipe.to("cuda")

pipe.transformer.eval() 
pipe.vae.eval()
pipe.text_encoder.eval()
pipe.pe.eval()


generator = torch.Generator(device="cuda").manual_seed(seed)


os.makedirs("../tests", exist_ok=True)

# 生成图片
prompt_list = [
    "A highly detailed biological pathway diagram in BioRender style. Depicting the viral infection process of human immune cells. Showing a virus particle attaching to a T-cell receptor, viral RNA replicating inside the cell nucleus, and the cell transforming into a malignant tumor cell. Includes molecular signaling pathways, proteins, and epigenetic modification symbols. Scientific flat vector style, soft pastel medical color palette, clean white background, educational graphic, crisp lines, professional scientific journal illustration. --ar 16:9 --v 6.0",
    "建筑坐落在住宅小区道路旁，被高大浓密的绿色乔木包围，树冠形成自然遮荫空间，严格保持图中所有元素的一致性。 午后自然阳光透过树叶形成斑驳光影（dappled sunlight），光线柔和且具有层次，地面呈现清晰树影，空气通透，微风感 建筑下方为开放式咖啡空间，人群自然分布，轻松社交状态，室内外边界模糊，空间通透流动 摄影级真实渲染，超高动态范围（HDR），自然曝光，真实反射与折射，玻璃微反射环境，极致细节。SANAA建筑风格，极简主义，轻盈漂浮感，日式当代建筑语言，自然主义建筑融合 高级建筑摄影风格，类似Dezeen / ArchDaily封面级别，纪实但理想化 电影级光影（cinematic natural lighting），柔和高光不过曝，阴影细腻。tree canopy filtered sunlight, soft shadows, volumetric light subtle, natural ambient occlusion 阳光穿过树叶产生细碎光斑，地面光影随机分布，光影边界柔和不过硬 色反射光（green bounce light）轻微影响建筑底部色温。广角镜头 24mm，低机位轻微仰视，前景草地，中景建筑，背景树林 画面有树枝作为前景遮挡（frame foreground leaves），增加空间层次 景深适中，整体清晰但有空气透视。--quality 2 --style raw --ar 3:4 --lighting natural --render photorealistic --detail ultra",
    "A photograph of the Straw Hat Pirates drawn on a glass whiteboard with a faded green marker, front view, 4K resolution."
]

for idx, prompt in enumerate(prompt_list):
    output = pipe(
        prompt=prompt,
        height=1024,
        width=1024,
        num_inference_steps=50,
        guidance_scale=5.0,
        generator=generator,
    )
    revised_prompt = output.revised_prompts
    images = output.images
    images[0].save(f"../tests/hf_output{idx+1}.png")
    print(f"Prompt {idx+1} revised: {revised_prompt}")