Z-Image-Turbo-SDA on Ascend NPU

# 设置 NPU 设备
export ASCEND_RT_VISIBLE_DEVICES=0

# 设置内存分配优化
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

# 设置 HCCL 缓冲区大小
export HCCL_BUFFSIZE=512

# 设置线程绑定
export OMP_PROC_BIND=false
export OMP_NUM_THREADS=1

# 下载基础模型
HF_ENDPOINT=https://hf-mirror.com huggingface-cli download \
    --resume-download Tongyi-MAI/Z-Image-Turbo \
    --local-dir /tmp/models/Z-Image-Turbo

# 下载 SDA LoKr 适配器
HF_ENDPOINT=https://hf-mirror.com huggingface-cli download \
    --resume-download F16/z-image-turbo-sda \
    --local-dir /opt/atomgit/z_image_turbo_sda/model

python inference.py \
    --base_model_path /tmp/models/Z-Image-Turbo \
    --adapter_path /opt/atomgit/z_image_turbo_sda/model/adapter/zit_sda_v1.safetensors \
    --prompt "A young woman standing on a sunny coastline, white dress fluttering in the sea breeze." \
    --height 1024 \
    --width 1024 \
    --num_inference_steps 8 \
    --guidance_scale 1.0 \
    --adapter_weight 1.0 \
    --dtype bfloat16 \
    --device npu:0 \
    --output output

# 检查模型加载
python -c "
import torch
import torch_npu
from diffusers import ZImagePipeline
from safetensors.torch import load_file

pipe = ZImagePipeline.from_pretrained('/tmp/models/Z-Image-Turbo', torch_dtype=torch.bfloat16)
pipe = pipe.to('npu:0')

adapter = load_file('/opt/atomgit/z_image_turbo_sda/model/adapter/zit_sda_v1.safetensors')
pipe.load_lora_weights(adapter, adapter_name='sda_diversity')
pipe.set_adapters(['sda_diversity'], adapter_weights=[1.0])
print('Model loaded successfully!')
print(f'Device: {pipe.device}')
print(f'Dtype: {pipe.dtype}')
"

# 快速推理测试 (512x512, 4步)
python inference.py \
    --height 512 \
    --width 512 \
    --num_inference_steps 4 \
    --warmup_runs 1 \
    --benchmark_runs 1 \
    --output output/smoke_test

python perf_eval.py \
    --height 1024 \
    --width 1024 \
    --num_inference_steps 8 \
    --num_warmup 2 \
    --num_iterations 10 \
    --output perf_results.json