train/usage_example.py

#!/usr/bin/env python3
"""
Usage Example for Trained QwenIllustrious Model
展示如何使用训练后的QwenIllustrious模型
"""

import os
import sys
from pathlib import Path

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from arch import QwenIllustriousInference


def example_usage_lora_weights():
    """
    Example: Using LoRA weights (separate from base model)
    示例：使用LoRA权重（与基础模型分离）
    """
    print("🚀 示例：使用LoRA权重进行推理")
    print("=" * 50)
    
    # 假设训练输出在这个目录
    output_dir = "./qwen_illustrious_output"
    
    pipeline = QwenIllustriousInference(
        # 基础模型路径
        qwen_model_path="models/Qwen3-Embedding-0.6B",
        unet_path="models/sdxl_base",  # 原始SDXL模型路径
        vae_path="models/sdxl_base",
        
        # 训练后的组件
        adapter_path=f"{output_dir}/adapter/adapter.safetensors",
        lora_weights_path=f"{output_dir}/lora_weights/lora_weights.safetensors",
        lora_config_path=f"{output_dir}/lora_weights",
        
        device="cuda",
        dtype="bfloat16"
    )
    
    if pipeline.is_ready:
        # 生成图像
        images = pipeline.generate(
            prompt="A beautiful anime girl in a magical garden, high quality",
            negative_prompt="low quality, blurry, distorted",
            height=1024,
            width=1024,
            num_inference_steps=50,
            guidance_scale=7.5
        )
        
        if images:
            images[0].save("example_lora_output.png")
            print("✅ 图像已保存到 example_lora_output.png")
    else:
        print("❌ 管道未准备就绪")


def example_usage_fused_model():
    """
    Example: Using fused model (LoRA merged into UNet)
    示例：使用融合模型（LoRA已合并到UNet中）
    """
    print("\n🚀 示例：使用融合模型进行推理")
    print("=" * 50)
    
    # 假设训练输出在这个目录
    output_dir = "./qwen_illustrious_output"
    
    pipeline = QwenIllustriousInference(
        # 基础模型路径
        qwen_model_path="models/Qwen3-Embedding-0.6B",
        vae_path="models/sdxl_base",
        
        # 训练后的组件
        adapter_path=f"{output_dir}/adapter/adapter.safetensors",
        use_fused_unet=True,
        fused_unet_path=f"{output_dir}/unet_fused",
        
        device="cuda",
        dtype="bfloat16"
    )
    
    if pipeline.is_ready:
        # 生成图像
        images = pipeline.generate(
            prompt="Two anime characters having a conversation in a cozy room",
            negative_prompt="low quality, blurry, distorted",
            height=1024,
            width=1024,
            num_inference_steps=50,
            guidance_scale=7.5
        )
        
        if images:
            images[0].save("example_fused_output.png")
            print("✅ 图像已保存到 example_fused_output.png")
    else:
        print("❌ 管道未准备就绪")


def training_command_examples():
    """
    Show example training commands
    显示训练命令示例
    """
    print("\n📚 训练命令示例")
    print("=" * 50)
    
    print("1. 基础训练命令:")
    print("python train/train_qwen_illustrious.py \\")
    print("  --qwen_model_path models/Qwen3-Embedding-0.6B \\")
    print("  --sdxl_model_path models/sdxl_base \\")
    print("  --dataset_path illustrious_generated \\")
    print("  --output_dir qwen_illustrious_output \\")
    print("  --train_batch_size 4 \\")
    print("  --num_train_epochs 10 \\")
    print("  --learning_rate 1e-4 \\")
    print("  --lora_rank 64 \\")
    print("  --mixed_precision fp16")
    
    print("\n2. 使用预计算嵌入加速训练:")
    print("# 第一步：预计算嵌入")
    print("python train/precompute_embeddings.py \\")
    print("  --qwen_model_path models/Qwen3-Embedding-0.6B \\")
    print("  --sdxl_model_path models/sdxl_base \\")
    print("  --dataset_path illustrious_generated \\")
    print("  --cache_dir cache \\")
    print("  --batch_size 8")
    
    print("\n# 第二步：使用预计算嵌入训练")
    print("python train/train_qwen_illustrious.py \\")
    print("  --qwen_model_path models/Qwen3-Embedding-0.6B \\")
    print("  --sdxl_model_path models/sdxl_base \\")
    print("  --dataset_path illustrious_generated \\")
    print("  --precompute_embeddings \\")
    print("  --cache_dir cache \\")
    print("  --output_dir qwen_illustrious_output \\")
    print("  --train_batch_size 8 \\")  # 可以使用更大的batch size
    print("  --num_train_epochs 10")


def inference_command_examples():
    """
    Show example inference commands
    显示推理命令示例
    """
    print("\n🎨 推理命令示例")
    print("=" * 50)
    
    print("1. 使用LoRA权重推理:")
    print("python inference_updated.py \\")
    print("  --prompt 'A beautiful anime girl in a garden' \\")
    print("  --adapter_path qwen_illustrious_output/adapter/adapter.safetensors \\")
    print("  --lora_weights_path qwen_illustrious_output/lora_weights/lora_weights.safetensors \\")
    print("  --lora_config_path qwen_illustrious_output/lora_weights \\")
    print("  --output my_image.png")
    
    print("\n2. 使用融合模型推理:")
    print("python inference_updated.py \\")
    print("  --prompt 'Two characters having a conversation' \\")
    print("  --adapter_path qwen_illustrious_output/adapter/adapter.safetensors \\")
    print("  --use_fused_unet \\")
    print("  --fused_unet_path qwen_illustrious_output/unet_fused \\")
    print("  --output my_image.png")


if __name__ == "__main__":
    print("🎯 QwenIllustrious 使用示例")
    print("=" * 60)
    
    training_command_examples()
    inference_command_examples()
    
    # 如果模型文件存在，运行实际示例
    output_dir = Path("./qwen_illustrious_output")
    if output_dir.exists():
        print("\n" + "=" * 60)
        print("🧪 运行实际推理示例 (检测到训练输出)")
        print("=" * 60)
        
        try:
            example_usage_lora_weights()
            example_usage_fused_model()
        except Exception as e:
            print(f"❌ 运行示例时出错: {e}")
    else:
        print(f"\n⚠️ 未找到训练输出目录: {output_dir}")
        print("请先运行训练脚本生成模型文件")