import torch
from diffusers import ZImagePipeline
import os

# 1. Load the pipeline
# Use bfloat16 for optimal performance on supported GPUs
pipe = ZImagePipeline.from_pretrained(
    "./Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)

# =================【这里是新增的加载 LoRA 代码】=================
# 指向你刚才训练输出的文件夹路径
lora_dir = "./feifei-zimage-lora" 
lora_file = "pytorch_lora_weights.safetensors"
full_path = os.path.join(lora_dir, lora_file)

if os.path.exists(full_path):
    print(f"正在加载 LoRA: {full_path}")
    try:
        # adapter_name 可以随意起，用来标记这个 LoRA
        pipe.load_lora_weights(lora_dir, weight_name=lora_file, adapter_name="feifei")
        print("✅ LoRA 加载成功！")
        
        # [可选] 设置 LoRA 的权重强度 (1.0 = 100% 强度, 0.5 = 50%)
        # pipe.set_adapters(["feifei"], adapter_weights=[1.0])
        
    except Exception as e:
        print(f"❌ LoRA 加载失败: {e}")
        print("可能是键名不匹配，或者文件损坏。")
else:
    print(f"⚠️ 找不到 LoRA 文件: {full_path}")
# ===============================================================

pipe.to("cuda")

# [Optional] Attention Backend
# pipe.transformer.set_attention_backend("flash")

prompt = "jpop model in bikini at sea"

# 2. Generate Image
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=9,
    guidance_scale=0.0,
    generator=torch.Generator("cuda").manual_seed(42),
    # cross_attention_kwargs={"scale": 1.0} # 另一种控制 LoRA 强度的方法
).images[0]

image.save("example_lora_test.png")
print("图像已保存为 example_lora_test.png")