"""
将 .safetensors 文件转换为 Diffusers 格式的脚本
用于 fal.ai 等需要完整 Diffusers 格式的平台
"""

from diffusers import StableDiffusionXLPipeline
import torch
import os

def convert_to_diffusers(
    input_path: str,
    output_dir: str,
    torch_dtype=torch.float16
):
    """
    转换 .safetensors 文件为 Diffusers 格式
    
    Args:
        input_path: 输入的 .safetensors 文件路径
        output_dir: 输出的 Diffusers 格式目录
        torch_dtype: 模型精度，默认 float16
    """
    
    print(f"开始转换: {input_path}")
    print(f"目标目录: {output_dir}")
    
    # 检查输入文件是否存在
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"找不到文件: {input_path}")
    
    # 加载单个文件
    print("\n步骤 1/2: 加载 .safetensors 文件...")
    pipe = StableDiffusionXLPipeline.from_single_file(
        input_path,
        torch_dtype=torch_dtype,
        use_safetensors=True
    )
    
    # 保存为 Diffusers 格式
    print("\n步骤 2/2: 保存为 Diffusers 格式...")
    pipe.save_pretrained(
        output_dir,
        safe_serialization=True
    )
    
    print("\n✓ 转换完成！")
    print(f"\n生成的目录结构位于: {output_dir}")
    print("\n包含的文件：")
    print("  ├── model_index.json")
    print("  ├── scheduler/")
    print("  ├── text_encoder/")
    print("  ├── text_encoder_2/")
    print("  ├── tokenizer/")
    print("  ├── tokenizer_2/")
    print("  ├── unet/")
    print("  └── vae/")


if __name__ == "__main__":
    # 配置参数
    INPUT_FILE = "./novaAsianXL_illustriousV50.safetensors"
    OUTPUT_DIR = "./novaAsianXL-diffusers"
    
    # 执行转换
    convert_to_diffusers(
        input_path=INPUT_FILE,
        output_dir=OUTPUT_DIR,
        torch_dtype=torch.float16
    )