import torch
from lyrasd_model import LyraSdXLTxt2ImgPipeline
import time
import GPUtil
import os
from glob import glob
import random

# 存放模型文件的路径，应该包含一下结构：
#   1. clip 模型
#   2. 转换好的优化后的 unet 模型，放入其中的 unet_bins 文件夹
#   3. vae 模型
#   4. scheduler 配置

# LyraSD 的 C++ 编译动态链接库，其中包含 C++ CUDA 计算的细节
lib_path = "./lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm80.so"
model_path = "./models/helloworldSDXL20Fp16"
lora_path = "./models/dissolve_sdxl.safetensors"
torch.classes.load_library(lib_path)

# 构建 Txt2Img 的 Pipeline
model = LyraSdXLTxt2ImgPipeline()

model.reload_pipe(model_path)

# load lora
# lora model path, name，lora strength
model.load_lora_v2(lora_path, "dissolve_sdxl", 0.4)

# 准备应用的输入和超参数
prompt = "a cat, ral-dissolve"
negative_prompt = "nswf, watermark"
height, width = 1024, 1024
steps = 20
guidance_scale = 7.5
generator = torch.Generator().manual_seed(8788800)

start = time.perf_counter()
# 推理生成
images = model(prompt,
               height=height,
               width=width,
               num_inference_steps=steps,
               num_images_per_prompt=1,
               guidance_scale=guidance_scale,
               negative_prompt=negative_prompt,
               generator=generator
               )
print("image gen cost: ", time.perf_counter() - start)
# 存储生成的图片
for i, image in enumerate(images):
    image.save(f"outputs/res_txt2img_xl_lora_{i}.png")

# unload lora，参数为 lora 的名字，是否清除 lora 缓存
model.unload_lora_v2("dissolve_sdxl", True)