#!/usr/bin/env python3
"""测试 Z-Image-Turbo Tiny 模型"""

import torch
from pathlib import Path

print("测试 Z-Image-Turbo Tiny 模型")
print("=" * 60)

# 使用当前目录（测试脚本所在目录）
model_path = Path(__file__).parent

# 1. 测试加载组件
print("\n[1] 测试加载各个组件...")

try:
    from diffusers import AutoencoderKL
    vae = AutoencoderKL.from_pretrained(
        model_path / "vae",
        torch_dtype=torch.float32,
        local_files_only=True
    )
    print(f"✓ VAE 加载成功: {sum(p.numel() for p in vae.parameters()):,} 参数")
except Exception as e:
    print(f"✗ VAE 加载失败: {e}")

try:
    from transformers import Qwen3Model
    text_encoder = Qwen3Model.from_pretrained(
        model_path / "text_encoder",
        torch_dtype=torch.float32,
        local_files_only=True
    )
    print(f"✓ Text Encoder 加载成功: {sum(p.numel() for p in text_encoder.parameters()):,} 参数")
except Exception as e:
    print(f"✗ Text Encoder 加载失败: {e}")

try:
    from transformers import Qwen2Tokenizer
    tokenizer = Qwen2Tokenizer.from_pretrained(
        model_path / "tokenizer",
        local_files_only=True
    )
    print(f"✓ Tokenizer 加载成功: vocab_size={tokenizer.vocab_size}")
except Exception as e:
    print(f"✗ Tokenizer 加载失败: {e}")

try:
    from diffusers import ZImageTransformer2DModel
    transformer = ZImageTransformer2DModel.from_pretrained(
        model_path / "transformer",
        torch_dtype=torch.float32,
        local_files_only=True
    )
    print(f"✓ Transformer 加载成功: {sum(p.numel() for p in transformer.parameters()):,} 参数")
except Exception as e:
    print(f"✗ Transformer 加载失败: {e}")
    print("  提示: 可能需要最新版 diffusers 或自定义模型类")

# 2. 测试加载完整 Pipeline
print("\n[2] 测试加载完整 Pipeline...")
try:
    from diffusers import ZImagePipeline
    pipe = ZImagePipeline.from_pretrained(
        str(model_path),
        torch_dtype=torch.float32,
        local_files_only=True
    )
    print("✓ Pipeline 加载成功!")
    print(f"  - Transformer: {sum(p.numel() for p in pipe.transformer.parameters()):,} 参数")
    print(f"  - VAE: {sum(p.numel() for p in pipe.vae.parameters()):,} 参数")
    print(f"  - Text Encoder: {sum(p.numel() for p in pipe.text_encoder.parameters()):,} 参数")
    
    # 3. 测试简单推理（使用 512x512）
    print("\n[3] 测试简单推理...")
    try:
        with torch.no_grad():
            # 简单的前向传播测试，使用 512x512
            test_prompt = "a cat"
            result = pipe(
                prompt=test_prompt,
                num_inference_steps=2,  # 极少步数仅测试
                height=512,  # Z-Image-Turbo 要求 512x512
                width=512,
            )
            print(f"✓ 推理成功! 生成图像 shape: {result.images[0].size}")
    except Exception as e:
        print(f"✗ 推理失败: {e}")
        print("\n详细错误信息:")
        import traceback
        traceback.print_exc()
        
except Exception as e:
    print(f"✗ Pipeline 加载失败: {e}")
    print("  提示: 需要安装支持 ZImagePipeline 的 diffusers 版本")

print("\n" + "=" * 60)
print("测试完成!")