diffusion-dpo-test/test_flops.py

import os
os.environ["HF_HOME"] = "/home/wanghongbo06/.cache/huggingface"

import torch
import torch.multiprocessing as mp
from diffusers.pipelines import FluxPipeline
from src.flux.condition import Condition
from src.flux.generate import generate, seed_everything
from color_fix import wavelet_color_fix, adain_color_fix
from PIL import Image
from tqdm import tqdm
import time
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Any


# ============== 配置 ==============
input_folder = "/home/wanghongbo06/baipudui/DATA/DIV2K/DIV2K-val-epoch1/lr"
output_folder = "/home/wanghongbo06/baipurui/results/flops"

# LoRA 路径
# SR_LORA_PATH = "/home/wanghongbo06/baipurui/CKPTs/FLUX_SR/pytorch_lora_weights_v2.safetensors"
SR_LORA_PATH = "/home/wanghongbo06/baipurui/OminiControl/runs/20260105-171922/ckpt/800/pytorch_lora_weights.safetensors"
# DPO_LORA_PATH = "/home/wanghongbo06/diffusion-dpo-adv/results/results_sobolev_20251212_2/checkpoint-400/lora_dpo/adapter_model.safetensors"
DPO_LORA_PATH = "/home/wanghongbo06/diffusion-dpo-adv/results/results_sobolev_20260107_1356/checkpoint-500/lora_dpo/adapter_model.safetensors"

# 可视化 Adversarial Sample
# ADV_LORA_PATH = '/home/wanghongbo06/diffusion-dpo/results/final_lora/adapter_model.safetensors'
# DPO_LORA_PATH = ADV_LORA_PATH



# LoRA scale（可以调整权重强度）
SR_LORA_SCALE = 1.0
DPO_LORA_SCALE = 1.0

# 多卡配置
NUM_GPUS = 1
# 同时加载模型的最大进程数（设为1表示串行加载，避免I/O瓶颈）
MAX_CONCURRENT_LOAD = 1

# 性能测试配置
WARMUP_IMAGES = 10  # 预热图片数（不计入统计）
ENABLE_PROFILING = True  # 是否启用详细性能分析
# ================================


@dataclass
class PerformanceMetrics:
    """性能指标数据类"""
    gpu_id: int
    inference_times: List[float]  # 每张图的推理时间（不含预热）
    warmup_time: float  # 预热时间
    peak_memory_mb: float  # 显存峰值 (MB)
    allocated_memory_mb: float  # 实际分配显存 (MB)
    reserved_memory_mb: float  # 保留显存 (MB)
    total_images: int  # 处理的总图片数
    
    @property
    def avg_inference_time(self) -> float:
        """平均推理时间（不含预热）"""
        if len(self.inference_times) == 0:
            return 0.0
        return np.mean(self.inference_times)
    
    @property
    def std_inference_time(self) -> float:
        """推理时间标准差"""
        if len(self.inference_times) < 2:
            return 0.0
        return np.std(self.inference_times)
    
    @property
    def throughput(self) -> float:
        """吞吐量（图片/秒）"""
        if len(self.inference_times) == 0:
            return 0.0
        total_time = sum(self.inference_times)
        return len(self.inference_times) / total_time if total_time > 0 else 0.0
    
    @property
    def memory_efficiency(self) -> float:
        """显存效率 = 实际分配 / 保留显存"""
        if self.reserved_memory_mb == 0:
            return 0.0
        return self.allocated_memory_mb / self.reserved_memory_mb * 100


def estimate_model_flops(pipe, height=512, width=512, num_inference_steps=28):
    """
    修正后的 FLOPs 估算函数 (针对 Flux 架构优化)
    """
    try:
        from fvcore.nn import FlopCountAnalysis, flop_count_str
        
        # 1. 获取正确的维度信息
        transformer = pipe.transformer
        config = transformer.config
        
        # Flux 特定的维度参数
        num_heads = config.num_attention_heads
        head_dim = config.attention_head_dim
        hidden_size = num_heads * head_dim  # 通常是 3072
        
        # 2. 计算 Latent 空间的分辨率
        # Flux 使用的 VAE 通常由 8x 下采样，patch size 为 1 或 2
        # 这里假设 input 是 latent，Sequence length = (H/16) * (W/16) * Time_ids ? 
        # Flux 处理 patch 后的 latent。Standard latent is H/8, W/8. 
        # Then patched to 2x2? Let's assume standard packed sequence length.
        # 对于 512x512 图片 -> Latent 64x64 = 4096 tokens.
        packed_seq_len = (height // 8) * (width // 8) // 4  # Flux patch_size=2 implies /2 on each dim?
        # 更安全的做法：直接取 4096 (针对 512x512) 或根据实际 latent 形状
        # Flux 的 latent 是 H/8, W/8. Flatten 后是 4096.
        seq_len = (height // 8) * (width // 8)
        
        print(f"DEBUG: Estimating with Hidden Size: {hidden_size}, Seq Len: {seq_len}")

        device = next(transformer.parameters()).device
        dtype = next(transformer.parameters()).dtype

        # 3. 构造正确维度的 Dummy Inputs
        # 注意：Flux forward 需要正确的 img_ids 和 txt_ids 才能跑通，
        # 为了避免构造复杂的 IDs 导致报错，我们这里只针对主要的 Linear 层进行 Hook，
        # 或者尝试构造尽可能真实的输入。
        
        dummy_hidden_states = torch.randn(1, seq_len, hidden_size, device=device, dtype=dtype)
        
        # Encoder hidden states (T5/CLIP text embeddings)
        # Flux text context length is usually 512
        dummy_encoder_hidden_states = torch.randn(1, 512, hidden_size, device=device, dtype=dtype) 
        
        # Pooled projections
        dummy_pooled = torch.randn(1, 768, device=device, dtype=dtype)
        
        # Timestep
        dummy_timestep = torch.tensor([500], device=device, dtype=dtype) # half precision
        
        # Flux 需要 img_ids 来计算 RoPE，如果传 None 可能会报错或跳过计算
        # 这里尝试只传必要的 args。如果 fvcore 报错，可能需要手动计算 Linear 层的 flops
        
        inputs = (
            dummy_hidden_states,
            dummy_encoder_hidden_states,
            dummy_pooled,
            dummy_timestep,
            # img_ids, txt_ids, guidance 通常可以为 None 或跳过，取决于具体实现
            # 如果报错，需要补全这些参数
        )

        # 4. 运行分析
        # 忽略未调用的参数警告
        flops_analysis = FlopCountAnalysis(transformer, inputs)
        
        # 强制忽略未使用的算子警告
        flops_analysis.unsupported_ops_warnings(False)
        
        single_forward_flops = flops_analysis.total()
        
        # 5. 加上 VAE 的估算 (粗略估算，通常 VAE 约占总量的 5-10% 或更少，但在 SR 中不能完全忽略)
        # 这里为了保守，只算 Transformer，但在报告中注明 "Transformer Only"
        
        total_flops = single_forward_flops * num_inference_steps
        
        print(f"DEBUG: Single step FLOPs: {single_forward_flops/1e12:.4f} TFLOPs")
        return total_flops, "fvcore (Transformer Only)"

    except Exception as e:
        print(f"fvcore FLOPs 估算失败: {e}")
        # 回退到理论计算 (Theoretical Calculation for Transformer)
        # Kaplan Scaling Laws approx: 6 * N * D_model^2 * Seq_len ???
        # 这里的备用方案应该更科学一点
        
        # 简单的 Transformer FLOPs 理论公式：
        # FLOPs per token ≈ 72 * (d_model ^ 2) (包含 attention 和 FFN) ? 
        # 更准确的近似: 
        # FLOPs = 24 * B * S * H^2 + 4 * B * S^2 * H (Attention + FFN)
        
        try:
            config = pipe.transformer.config
            H = config.num_attention_heads * config.attention_head_dim
            L = config.num_layers
            S = (height // 8) * (width // 8)
            
            # 这是一个非常粗略的 Transformer 理论计算
            # 1. Linear Layers (Q,K,V, Out, MLP up, MLP down)
            # 每一层通常有 4个投影 (Attn) + 3个投影 (MLP)? Flux 是 MMDiT 结构更复杂
            # 保守估计：每层参数量 P_layer. FLOPs ≈ 2 * P_layer * S
            
            total_params = sum(p.numel() for p in pipe.transformer.parameters())
            # Transformer FLOPs ≈ 2 * Params * Sequence_Length
            theoretical_flops = 2 * total_params * S * num_inference_steps
            
            return theoretical_flops, "Theoretical (2*Params*SeqLen)"
        except:
            return 0, "failed"

def profile_single_inference(pipe, image, prompt, condition, device):
    """
    对单次推理进行详细的性能分析
    """
    # 确保在正确的设备上操作
    device_id = int(device.split(':')[1]) if isinstance(device, str) else device
    
    torch.cuda.reset_peak_memory_stats(device)
    
    # 使用 with torch.cuda.device 确保事件在正确的设备上创建
    with torch.cuda.device(device_id):
        torch.cuda.synchronize()
        
        # 使用 time.perf_counter 作为更可靠的计时方式（多GPU兼容）
        start_time = time.perf_counter()
        
        result_img = generate(
            pipe,
            prompt=prompt,
            conditions=[condition],
            default_lora=True,
        ).images[0]
        
        torch.cuda.synchronize()
        end_time = time.perf_counter()
        
        inference_time = end_time - start_time
    
    # 获取显存信息
    peak_memory = torch.cuda.max_memory_allocated(device) / (1024 ** 2)  # MB
    allocated_memory = torch.cuda.memory_allocated(device) / (1024 ** 2)  # MB
    reserved_memory = torch.cuda.memory_reserved(device) / (1024 ** 2)  # MB
    
    return result_img, inference_time, peak_memory, allocated_memory, reserved_memory


def load_pipeline(gpu_id, load_semaphore=None):
    """在指定 GPU 上加载 pipeline，使用信号量控制并发加载"""
    device = f"cuda:{gpu_id}"
    
    # 显式设置当前进程使用的 GPU
    torch.cuda.set_device(gpu_id)
    
    # 使用信号量控制同时加载模型的进程数
    if load_semaphore is not None:
        load_semaphore.acquire()
    
    try:
        print(f"[GPU {gpu_id}] 开始加载模型...")
        load_start = time.time()
        pipe = FluxPipeline.from_pretrained(
            '/home/wanghongbo06/baipurui/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev/snapshots/3de623fc3c33e44ffbe2bad470d0f45bccf2eb21',
            torch_dtype=torch.bfloat16,
            token="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",
            local_files_on=True,
            catch_dir=".cache/flux-sr"
        ).to(device)
        
        # 加载 LoRA
        pipe.load_lora_weights(SR_LORA_PATH, adapter_name="sr")
        pipe.load_lora_weights(DPO_LORA_PATH, adapter_name="dpo")
        pipe.set_adapters(["sr", "dpo"], adapter_weights=[SR_LORA_SCALE, DPO_LORA_SCALE])
        
        load_time = time.time() - load_start
        print(f"[GPU {gpu_id}] 模型加载完成，耗时 {load_time:.1f}s")
        
    finally:
        if load_semaphore is not None:
            load_semaphore.release()
    
    return pipe


def process_images(gpu_id, image_list, output_folder, load_semaphore, ready_event, start_barrier, metrics_dict=None):
    """
    单个 GPU 上的处理函数
    Args:
        gpu_id: GPU 编号
        image_list: 该 GPU 需要处理的图片文件名列表
        output_folder: 输出目录
        load_semaphore: 控制模型加载并发的信号量
        ready_event: 通知主进程模型已加载完成
        start_barrier: 同步所有进程开始推理
        metrics_dict: 用于存储性能指标的共享字典
    """
    try:
        if len(image_list) == 0:
            ready_event.set()
            start_barrier.wait()
            return
        
        device = f"cuda:{gpu_id}"
        
        # 显式设置当前进程使用的 GPU（在子进程开始时设置）
        torch.cuda.set_device(gpu_id)
        
        # 加载模型到指定 GPU（通过信号量控制并发）
        pipe = load_pipeline(gpu_id, load_semaphore)
        
        # 通知主进程该GPU模型已加载完成
        ready_event.set()
        
        # 等待所有GPU都加载完成后再开始推理
        start_barrier.wait()
        
        print(f"[GPU {gpu_id}] 开始处理 {len(image_list)} 张图片")
        
        prompt = ""
        
        # 性能统计变量
        inference_times = []
        warmup_time = 0.0
        peak_memory_mb = 0.0
        allocated_memory_mb = 0.0
        reserved_memory_mb = 0.0
        
        # 重置显存统计
        torch.cuda.reset_peak_memory_stats(device)
        
        # 只在 GPU 0 上显示主进度条
        pbar = tqdm(
            enumerate(image_list), 
            total=len(image_list),
            desc=f"GPU {gpu_id}",
            position=gpu_id,
            leave=True,
            ncols=120,
            bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]'
        )
        
        for idx, filename in pbar:
            image_path = os.path.join(input_folder, filename)
            image = Image.open(image_path).convert("RGB")
            
            # 居中裁剪 + resize 到 512x512
            w, h = image.size
            min_dim = min(w, h)
            image = image.crop(
                ((w - min_dim) // 2, (h - min_dim) // 2, (w + min_dim) // 2, (h + min_dim) // 2)
            ).resize((512, 512), Image.BICUBIC)
            
            # 构造条件并生成图像
            condition = Condition("sr", image)
            seed_everything(1)
            
            # 使用精确计时进行推理
            result_img, inf_time, peak_mem, alloc_mem, reserved_mem = profile_single_inference(
                pipe, image, prompt, condition, device
            )
            
            # 更新显存峰值
            peak_memory_mb = max(peak_memory_mb, peak_mem)
            allocated_memory_mb = alloc_mem
            reserved_memory_mb = reserved_mem
            
            # 区分预热和正式推理
            if idx < WARMUP_IMAGES:
                warmup_time += inf_time
                pbar.set_postfix({
                    'warmup': f'{inf_time:.2f}s',
                    'mem': f'{peak_mem:.0f}MB'
                })
            else:
                inference_times.append(inf_time)
                avg_time = np.mean(inference_times)
                pbar.set_postfix({
                    'time': f'{inf_time:.2f}s',
                    'avg': f'{avg_time:.2f}s',
                    'mem': f'{peak_mem:.0f}MB'
                })
            
            result_img = adain_color_fix(result_img, image)
            result_img.save(os.path.join(output_folder, filename))
        
        # 获取最终显存统计
        final_peak_memory = torch.cuda.max_memory_allocated(device) / (1024 ** 2)
        final_allocated = torch.cuda.memory_allocated(device) / (1024 ** 2)
        final_reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)
        
        # 创建性能指标对象
        metrics = PerformanceMetrics(
            gpu_id=gpu_id,
            inference_times=inference_times,
            warmup_time=warmup_time,
            peak_memory_mb=max(peak_memory_mb, final_peak_memory),
            allocated_memory_mb=final_allocated,
            reserved_memory_mb=final_reserved,
            total_images=len(image_list)
        )
        
        # 存储到共享字典
        if metrics_dict is not None:
            metrics_dict[gpu_id] = {
                'inference_times': inference_times,
                'warmup_time': warmup_time,
                'peak_memory_mb': metrics.peak_memory_mb,
                'allocated_memory_mb': metrics.allocated_memory_mb,
                'reserved_memory_mb': metrics.reserved_memory_mb,
                'total_images': len(image_list),
                'avg_inference_time': metrics.avg_inference_time,
                'std_inference_time': metrics.std_inference_time,
                'throughput': metrics.throughput,
                'memory_efficiency': metrics.memory_efficiency
            }
        
        # 打印单 GPU 性能摘要
        print(f"\n[GPU {gpu_id}] ✅ 完成！")
        print(f"  📊 性能摘要:")
        print(f"     - 处理图片数: {len(image_list)} (预热: {WARMUP_IMAGES}, 统计: {len(inference_times)})")
        print(f"     - 预热时间: {warmup_time:.2f}s")
        print(f"     - 平均推理时间: {metrics.avg_inference_time:.3f}s ± {metrics.std_inference_time:.3f}s")
        print(f"     - 吞吐量: {metrics.throughput:.2f} 图/秒")
        print(f"     - 显存峰值: {metrics.peak_memory_mb:.1f} MB")
        print(f"     - 显存效率: {metrics.memory_efficiency:.1f}%")
    
    except Exception as e:
        print(f"\n[GPU {gpu_id}] ❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        # 确保事件被设置，避免死锁
        ready_event.set()
        raise


def print_performance_report(metrics_dict: Dict[int, Dict], load_time: float, total_time: float, total_images: int):
    """
    打印详细的性能报告
    """
    print("\n" + "=" * 70)
    print("                       📊 详细性能报告")
    print("=" * 70)
    
    # 汇总所有 GPU 的数据
    all_inference_times = []
    total_warmup_time = 0.0
    max_peak_memory = 0.0
    total_allocated_memory = 0.0
    total_reserved_memory = 0.0
    
    for gpu_id, metrics in sorted(metrics_dict.items()):
        all_inference_times.extend(metrics['inference_times'])
        total_warmup_time += metrics['warmup_time']
        max_peak_memory = max(max_peak_memory, metrics['peak_memory_mb'])
        total_allocated_memory += metrics['allocated_memory_mb']
        total_reserved_memory += metrics['reserved_memory_mb']
    
    # ============== 1. 推理时间统计 ==============
    print("\n🕐 推理时间统计:")
    print("-" * 50)
    
    if len(all_inference_times) > 0:
        avg_time = np.mean(all_inference_times)
        std_time = np.std(all_inference_times)
        min_time = np.min(all_inference_times)
        max_time = np.max(all_inference_times)
        median_time = np.median(all_inference_times)
        p95_time = np.percentile(all_inference_times, 95)
        p99_time = np.percentile(all_inference_times, 99)
        
        print(f"   统计图片数: {len(all_inference_times)} (排除预热 {WARMUP_IMAGES * len(metrics_dict)} 张)")
        print(f"   平均推理时间: {avg_time:.4f} 秒/张")
        print(f"   标准差: {std_time:.4f} 秒")
        print(f"   最小值: {min_time:.4f} 秒")
        print(f"   最大值: {max_time:.4f} 秒")
        print(f"   中位数: {median_time:.4f} 秒")
        print(f"   P95: {p95_time:.4f} 秒")
        print(f"   P99: {p99_time:.4f} 秒")
        print(f"   预热总时间: {total_warmup_time:.2f} 秒")
    else:
        print("   ⚠️  没有有效的推理时间数据")
    
    # ============== 2. 吞吐量统计 ==============
    print("\n⚡ 吞吐量 (Throughput):")
    print("-" * 50)
    
    if len(all_inference_times) > 0:
        total_inference_time = sum(all_inference_times)
        throughput_per_sec = len(all_inference_times) / total_inference_time if total_inference_time > 0 else 0
        throughput_per_min = throughput_per_sec * 60
        throughput_per_hour = throughput_per_sec * 3600
        
        # 多卡并行吞吐量（wall-clock time）
        inference_wall_time = total_time - load_time
        parallel_throughput_sec = total_images / inference_wall_time if inference_wall_time > 0 else 0
        parallel_throughput_min = parallel_throughput_sec * 60
        
        print(f"   单 GPU 吞吐量:")
        print(f"      - {throughput_per_sec:.3f} 图/秒")
        print(f"      - {throughput_per_min:.1f} 图/分钟")
        print(f"      - {throughput_per_hour:.0f} 图/小时")
        print(f"   {len(metrics_dict)} GPU 并行吞吐量 (wall-clock):")
        print(f"      - {parallel_throughput_sec:.3f} 图/秒")
        print(f"      - {parallel_throughput_min:.1f} 图/分钟")
    
    # ============== 3. 显存统计 ==============
    print("\n💾 显存 (GPU Memory):")
    print("-" * 50)
    
    for gpu_id, metrics in sorted(metrics_dict.items()):
        print(f"   GPU {gpu_id}:")
        print(f"      - 显存峰值: {metrics['peak_memory_mb']:.1f} MB ({metrics['peak_memory_mb']/1024:.2f} GB)")
        print(f"      - 实际分配: {metrics['allocated_memory_mb']:.1f} MB")
        print(f"      - 保留显存: {metrics['reserved_memory_mb']:.1f} MB")
        print(f"      - 显存效率: {metrics['memory_efficiency']:.1f}%")
    
    if len(metrics_dict) > 1:
        print(f"   汇总:")
        print(f"      - 最大显存峰值: {max_peak_memory:.1f} MB ({max_peak_memory/1024:.2f} GB)")
        print(f"      - 总分配显存: {total_allocated_memory:.1f} MB")
    
    # ============== 4. FLOPs 估算 ==============
    print("\n🔢 计算量 (FLOPs) - 估算:")
    print("-" * 50)
    print("   ⚠️  FLOPs 估算需要在单 GPU 模式下单独运行")
    print("   💡 提示: 设置 NUM_GPUS=1 并运行 estimate_flops_standalone() 获取准确值")
    
    # ============== 5. 时间分解 ==============
    print("\n⏱️  时间分解:")
    print("-" * 50)
    inference_time = total_time - load_time
    print(f"   模型加载时间: {load_time:.1f} 秒 ({load_time/total_time*100:.1f}%)")
    print(f"   推理时间: {inference_time:.1f} 秒 ({inference_time/total_time*100:.1f}%)")
    print(f"   总时间: {total_time:.1f} 秒")
    
    # ============== 6. 汇总 ==============
    print("\n" + "=" * 70)
    print("                       📈 性能汇总")
    print("=" * 70)
    
    if len(all_inference_times) > 0:
        avg_time = np.mean(all_inference_times)
        print(f"""
┌─────────────────────────────────────────────────────────────────┐
│  指标                          │  值                            │
├─────────────────────────────────────────────────────────────────┤
│  平均推理时间 (不含预热)       │  {avg_time:.4f} 秒/张               │
│  吞吐量 (单GPU)               │  {throughput_per_sec:.3f} 图/秒              │
│  吞吐量 ({len(metrics_dict)}GPU 并行)            │  {parallel_throughput_sec:.3f} 图/秒              │
│  显存峰值                      │  {max_peak_memory:.1f} MB ({max_peak_memory/1024:.2f} GB)      │
│  总处理图片                    │  {total_images} 张                         │
└─────────────────────────────────────────────────────────────────┘
""")
    
    print("=" * 70)


def estimate_flops_standalone():
    """
    独立运行的 FLOPs 估算函数
    需要在单 GPU 上运行
    """
    print("=" * 60)
    print("🔢 正在估算模型 FLOPs...")
    print("=" * 60)
    
    device = "cuda:0"
    
    # 加载模型
    print("加载模型中...")
    pipe = FluxPipeline.from_pretrained(
        '/home/wanghongbo06/baipurui/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev/snapshots/3de623fc3c33e44ffbe2bad470d0f45bccf2eb21',
        torch_dtype=torch.bfloat16,
        token="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",
        local_files_on=True,
        catch_dir=".cache/flux-sr"
    ).to(device)
    
    pipe.load_lora_weights(SR_LORA_PATH, adapter_name="sr")
    pipe.load_lora_weights(DPO_LORA_PATH, adapter_name="dpo")
    pipe.set_adapters(["sr", "dpo"], adapter_weights=[SR_LORA_SCALE, DPO_LORA_SCALE])
    
    # 估算 FLOPs
    flops, method = estimate_model_flops(pipe)
    
    if flops > 0:
        print(f"\n📊 FLOPs 估算结果 (方法: {method}):")
        print(f"   - 每次推理 FLOPs: {flops:.2e}")
        print(f"   - 每次推理 TFLOPs: {flops / 1e12:.2f}")
        
        # 如果有推理时间，可以计算 FLOPS (每秒浮点运算数)
        # FLOPS = FLOPs / inference_time
    else:
        print("❌ FLOPs 估算失败")
    
    return flops


def save_metrics_to_json(metrics_dict: Dict, output_path: str, load_time: float, total_time: float, total_images: int):
    """
    将性能指标保存到 JSON 文件
    """
    import json
    
    # 计算汇总指标
    all_times = []
    for gpu_id, m in metrics_dict.items():
        all_times.extend(m['inference_times'])
    
    inference_wall_time = total_time - load_time
    
    summary = {
        'avg_inference_time_sec': float(np.mean(all_times)) if all_times else 0,
        'std_inference_time_sec': float(np.std(all_times)) if all_times else 0,
        'min_inference_time_sec': float(np.min(all_times)) if all_times else 0,
        'max_inference_time_sec': float(np.max(all_times)) if all_times else 0,
        'median_inference_time_sec': float(np.median(all_times)) if all_times else 0,
        'p95_inference_time_sec': float(np.percentile(all_times, 95)) if all_times else 0,
        'p99_inference_time_sec': float(np.percentile(all_times, 99)) if all_times else 0,
        'throughput_single_gpu_per_sec': float(len(all_times) / sum(all_times)) if all_times and sum(all_times) > 0 else 0,
        'throughput_parallel_per_sec': float(total_images / inference_wall_time) if inference_wall_time > 0 else 0,
        'peak_memory_mb': max([m['peak_memory_mb'] for m in metrics_dict.values()]) if metrics_dict else 0,
        'peak_memory_gb': max([m['peak_memory_mb'] for m in metrics_dict.values()]) / 1024 if metrics_dict else 0,
        'total_images': total_images,
        'warmup_images': WARMUP_IMAGES * len(metrics_dict),
        'measured_images': len(all_times),
        'model_load_time_sec': load_time,
        'inference_wall_time_sec': inference_wall_time,
        'total_time_sec': total_time,
        'num_gpus': len(metrics_dict),
    }
    
    result = {
        'summary': summary,
        'per_gpu_metrics': {str(k): v for k, v in metrics_dict.items()}
    }
    
    with open(output_path, 'w') as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
    
    print(f"📄 性能指标已保存到: {output_path}")


def main(save_metrics_path: str = None):
    """
    主函数
    Args:
        save_metrics_path: 可选，保存性能指标的 JSON 文件路径
    """
    os.makedirs(output_folder, exist_ok=True)
    
    # 获取所有待处理的图片
    all_images = sorted([
        f for f in os.listdir(input_folder)
        if f.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".webp"))
    ])
    
    total_images = len(all_images)
    print("=" * 70)
    print("                    🚀 Diffusion 超分性能测试")
    print("=" * 70)
    print(f"📁 输入目录: {input_folder}")
    print(f"📁 输出目录: {output_folder}")
    print(f"🖼️  总图片数: {total_images}")
    print(f"🎮 GPU 数量: {NUM_GPUS}")
    print(f"📦 每 GPU 处理: ~{total_images // NUM_GPUS} 张")
    print(f"⚙️  模型加载并发数: {MAX_CONCURRENT_LOAD}")
    print(f"🔥 预热图片数: {WARMUP_IMAGES} (每个GPU)")
    print(f"📊 性能分析: {'开启' if ENABLE_PROFILING else '关闭'}")
    print("=" * 70)
    
    # 将图片列表平均分配给各个 GPU
    image_chunks = [[] for _ in range(NUM_GPUS)]
    for i, img in enumerate(all_images):
        image_chunks[i % NUM_GPUS].append(img)
    
    # 记录开始时间
    start_time = time.time()
    
    # 使用多进程并行处理
    mp.set_start_method('spawn', force=True)
    
    # 创建信号量来限制同时加载模型的进程数（避免I/O瓶颈）
    load_semaphore = mp.Semaphore(MAX_CONCURRENT_LOAD)
    
    # 创建事件来追踪每个进程的模型加载状态
    ready_events = [mp.Event() for _ in range(NUM_GPUS)]
    
    # 创建屏障来同步所有进程在加载完成后开始推理
    start_barrier = mp.Barrier(NUM_GPUS)
    
    # 创建共享字典存储各 GPU 的性能指标
    manager = mp.Manager()
    metrics_dict = manager.dict()
    
    processes = []
    
    print(f"\n⏳ 开始加载模型（最多 {MAX_CONCURRENT_LOAD} 个并发，避免I/O瓶颈）...")
    
    for gpu_id in range(NUM_GPUS):
        p = mp.Process(
            target=process_images,
            args=(gpu_id, image_chunks[gpu_id], output_folder, 
                  load_semaphore, ready_events[gpu_id], start_barrier, metrics_dict)
        )
        p.start()
        processes.append(p)
    
    # 等待所有模型加载完成
    loaded_count = 0
    for i, event in enumerate(ready_events):
        event.wait()
        loaded_count += 1
        print(f"  ✅ GPU {i} 就绪 ({loaded_count}/{NUM_GPUS})")
    
    load_time = time.time() - start_time
    print(f"\n⏱️  模型加载总耗时: {load_time:.1f}s ({load_time/60:.1f} 分钟)")
    print("🚀 所有模型加载完成，开始并行推理...\n")
    
    # 等待所有进程完成
    for p in processes:
        p.join()
    
    # 计算总耗时
    total_time = time.time() - start_time
    
    # 将 manager.dict 转换为普通 dict
    metrics_dict_normal = dict(metrics_dict)
    
    # 打印详细性能报告
    print_performance_report(metrics_dict_normal, load_time, total_time, total_images)
    
    # 保存性能指标到文件
    if save_metrics_path:
        save_metrics_to_json(metrics_dict_normal, save_metrics_path, load_time, total_time, total_images)
    
    print(f"\n📁 结果保存在: {output_folder}")
    print("=" * 70)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Diffusion 超分性能测试')
    parser.add_argument('--mode', type=str, default='benchmark', choices=['benchmark', 'flops'],
                        help='运行模式: benchmark (默认) 或 flops (仅估算FLOPs)')
    parser.add_argument('--save-metrics', type=str, default=None,
                        help='保存性能指标到 JSON 文件的路径')
    parser.add_argument('--num-gpus', type=int, default=None,
                        help='使用的 GPU 数量 (覆盖默认值)')
    parser.add_argument('--warmup', type=int, default=None,
                        help='预热图片数量 (覆盖默认值)')
    
    args = parser.parse_args()
    
    # 覆盖配置 (使用 global)
    if args.num_gpus is not None:
        NUM_GPUS = args.num_gpus
    if args.warmup is not None:
        WARMUP_IMAGES = args.warmup
    
    if args.mode == 'flops':
        # 仅估算 FLOPs
        estimate_flops_standalone()
    else:
        # 运行完整的 benchmark
        main(save_metrics_path=args.save_metrics)


# pyiqa psnr ssim lpips musiq clipiqa+ --target /home/wanghongbo06/diffusion-dpo-test/DIV2K-val/sobolev-400  --r /home/wanghongbo06/baipurui/DATA/DIV2K-val/gt