JCo-MVTON/inference_dresscode_batch.py

"""
DressCodeDataset 多卡批量推理脚本
基于 inference_viton_batch.py 重写，支持多 GPU 分布式推理
基于 DressCodeDataset 下的 dresses, lower_body, upper_body 子目录中的 test_pairs_unpaired.txt 文件批量处理

使用方法（单卡，自动选择checkpoint）：
python inference_dresscode_batch.py \
    --dresscode_root /filesdir/DressCodeDataset \
    --output_dir /filesdir/JCo-MVTON/output_dresscode \
    --gpu_id 0 \
    --category dresses
# dresses 自动使用 try_on_dress.pt
# lower_body 自动使用 try_on_lower.pt
# upper_body 自动使用 try_on_upper.pt

使用方法（单卡，处理所有类别）：
python inference_dresscode_batch.py \
    --dresscode_root /filesdir/DressCodeDataset \
    --output_dir /filesdir/JCo-MVTON/output_dresscode \
    --gpu_id 0 \
    --category all
# 每个类别自动使用对应的checkpoint

使用方法（多卡，使用 torchrun）：
torchrun --nproc_per_node=4 inference_dresscode_batch.py \
    --dresscode_root /filesdir/DressCodeDataset \
    --output_dir /filesdir/JCo-MVTON/output_dresscode \
    --category dresses

使用方法（多卡，使用 torchrun + nohup 后台运行）：
nohup torchrun --nproc_per_node=4 inference_dresscode_batch.py \
    --dresscode_root /filesdir/DressCodeDataset \
    --output_dir /filesdir/JCo-MVTON/output_dresscode \
    --category upper_body \
    > nohup_dresscode_upper_body.out 2>&1 &

使用方法（手动多进程，每个类别一个GPU）：
CUDA_VISIBLE_DEVICES=0 python inference_dresscode_batch.py --gpu_id 0 --category dresses &
CUDA_VISIBLE_DEVICES=1 python inference_dresscode_batch.py --gpu_id 1 --category lower_body &
CUDA_VISIBLE_DEVICES=2 python inference_dresscode_batch.py --gpu_id 2 --category upper_body &

使用方法（手动指定checkpoint，覆盖自动选择）：
python inference_dresscode_batch.py \
    --dresscode_root /filesdir/DressCodeDataset \
    --checkpoint /filesdir/JCo-MVTON/try_on_upper.pt \
    --output_dir /filesdir/JCo-MVTON/output_dresscode \
    --gpu_id 0 \
    --category dresses
# 如果指定了 --checkpoint，所有类别都使用这个checkpoint
"""

import torch
import torch.distributed as dist
from flux.pipeline_flux import FluxPipeline
from flux.transformer_flux import FluxTransformer2DModel
import os
import argparse
from pathlib import Path
from PIL import Image
from torchvision import transforms
import torchvision.utils as vutils
from tqdm import tqdm


def setup_distributed():
    """初始化分布式训练环境"""
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        rank = int(os.environ['RANK'])
        world_size = int(os.environ['WORLD_SIZE'])
        local_rank = int(os.environ.get('LOCAL_RANK', 0))
        
        torch.cuda.set_device(local_rank)
        dist.init_process_group(backend='nccl')
        
        return rank, world_size, local_rank
    else:
        return None, 1, 0


def load_test_pairs(test_unpairs_file):
    """加载测试配对文件"""
    pairs = []
    with open(test_unpairs_file, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            # 支持制表符和空格分隔
            parts = line.split('\t') if '\t' in line else line.split()
            if len(parts) >= 2:
                person_img = parts[0].strip()
                cloth_img = parts[1].strip()
                pairs.append((person_img, cloth_img))
    return pairs


def load_model(checkpoint_path, device, mode=2, extra_branch_num=2, torch_dtype=torch.bfloat16):
    """加载模型（基于 inference.py 的逻辑）"""
    model_id = "black-forest-labs/FLUX.1-dev"
    
    print(f"[GPU {device}] 加载基础模型...")
    transformer = FluxTransformer2DModel.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        subfolder="transformer",
        extra_branch_num=extra_branch_num,
        local_files_only=True,
        low_cpu_mem_usage=False,
    ).to(device)
    
    # 初始化 extra branch (mode 2)
    print(f"[GPU {device}] 初始化 extra branch (mode {mode})...")
    with torch.no_grad():
        for j in range(extra_branch_num):
            if mode == 1:
                transformer.extra_embedder[j].load_state_dict(transformer.x_embedder.state_dict())
            for i in range(transformer.config.num_layers):
                transformer.transformer_blocks[i].attn.extra_to_q[j].load_state_dict(
                    transformer.transformer_blocks[i].attn.to_q.state_dict())
                transformer.transformer_blocks[i].attn.extra_to_k[j].load_state_dict(
                    transformer.transformer_blocks[i].attn.to_k.state_dict())
                transformer.transformer_blocks[i].attn.extra_to_v[j].load_state_dict(
                    transformer.transformer_blocks[i].attn.to_v.state_dict())
                if mode == 1:
                    transformer.transformer_blocks[i].extra_norm1[j].load_state_dict(
                        transformer.transformer_blocks[i].norm1.state_dict())
                    transformer.transformer_blocks[i].extra_norm2[j].load_state_dict(
                        transformer.transformer_blocks[i].norm2.state_dict())
                    transformer.transformer_blocks[i].extra_ff[j].load_state_dict(
                        transformer.transformer_blocks[i].ff.state_dict())
                    transformer.transformer_blocks[i].attn.extra_to_out[0][j].load_state_dict(
                        transformer.transformer_blocks[i].attn.to_out[0].state_dict())
                    transformer.transformer_blocks[i].attn.extra_to_out[1][j].load_state_dict(
                        transformer.transformer_blocks[i].attn.to_out[1].state_dict())
                    transformer.transformer_blocks[i].attn.extra_norm_q[j].load_state_dict(
                        transformer.transformer_blocks[i].attn.norm_q.state_dict())
                    transformer.transformer_blocks[i].attn.extra_norm_k[j].load_state_dict(
                        transformer.transformer_blocks[i].attn.norm_k.state_dict())
            if mode == 2:
                for i in range(transformer.config.num_single_layers):
                    transformer.single_transformer_blocks[i].attn.extra_to_q[j].load_state_dict(
                        transformer.single_transformer_blocks[i].attn.to_q.state_dict())
                    transformer.single_transformer_blocks[i].attn.extra_to_k[j].load_state_dict(
                        transformer.single_transformer_blocks[i].attn.to_k.state_dict())
                    transformer.single_transformer_blocks[i].attn.extra_to_v[j].load_state_dict(
                        transformer.single_transformer_blocks[i].attn.to_v.state_dict())
    
    # 创建 pipeline
    print(f"[GPU {device}] 创建 pipeline...")
    pipe = FluxPipeline.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        transformer=transformer,
        local_files_only=True,
    ).to(device)
    
    # 加载训练好的权重
    print(f"[GPU {device}] 加载 checkpoint: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    transformer.load_state_dict(checkpoint['module'], strict=False)
    del checkpoint  # 释放内存
    torch.cuda.empty_cache()  # 清理 GPU 缓存
    
    # 重新创建 pipeline（权重更新后）
    pipe = FluxPipeline.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        transformer=transformer,
        local_files_only=True,
    ).to(device)
    
    return pipe, transformer


def process_pair(person_path, cloth_path, pipe, device, output_dir, resolution=1024, mode=2, seed=0):
    """处理一对图像（基于 inference.py 的逻辑）"""
    height = resolution
    width = resolution * 3 // 4
    
    transform_person = transforms.Compose([
        transforms.Resize(size=(height, width)),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])
    
    transform_cloth = transforms.Compose([
        transforms.Resize(size=(height, height)),
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5])
    ])
    
    transform_output = transforms.Compose([
        transforms.ToTensor(),
    ])
    
    try:
        # 加载图像
        person = Image.open(person_path).convert("RGB").resize((width, height))
        cloth = Image.open(cloth_path).convert("RGB").resize((height, height))
        
        person_tensor = transform_person(person)
        cloth_tensor = transform_cloth(cloth)
        
        prompt = "A fashion model wearing stylish clothing, high-resolution 8k, detailed textures, realistic lighting, fashion photography style."
        
        # 生成图像
        with torch.inference_mode():
            generated_image = pipe(
                generator=torch.Generator(device="cpu").manual_seed(seed),
                prompt=prompt,
                num_inference_steps=28,
                guidance_scale=3.5,
                height=height,
                width=width,
                cloth_img=cloth_tensor,
                person_img=person_tensor,
                extra_branch_num=2,
                mode=mode,
                max_sequence_length=77,
            ).images[0]
        
        # 生成的图像本来就是 768 宽度，直接使用
        generated_tensor = transform_output(generated_image)
        
        # 只保存生成的图像（人穿衣服的部分），不拼接
        output_filename = Path(person_path).stem + "_" + Path(cloth_path).stem + ".png"
        output_path = os.path.join(output_dir, output_filename)
        vutils.save_image(generated_tensor, output_path)
        
        return True, None
    except Exception as e:
        return False, str(e)


def main():
    parser = argparse.ArgumentParser(description='DressCodeDataset 多卡批量推理')
    parser.add_argument('--dresscode_root', type=str, required=True,
                        help='DressCodeDataset 根目录路径')
    parser.add_argument('--checkpoint', type=str, default=None,
                        help='模型 checkpoint 路径（如果指定，将用于所有类别；否则会根据类别自动选择）')
    parser.add_argument('--checkpoint_dir', type=str, default='/filesdir/JCo-MVTON',
                        help='Checkpoint 目录路径（用于自动选择checkpoint）')
    parser.add_argument('--output_dir', type=str, required=True,
                        help='输出目录路径')
    parser.add_argument('--category', type=str, required=True,
                        choices=['dresses', 'lower_body', 'upper_body', 'all'],
                        help='要处理的类别：dresses, lower_body, upper_body, 或 all（处理所有类别）')
    parser.add_argument('--gpu_id', type=int, default=0,
                        help='使用的 GPU ID（单卡模式）')
    parser.add_argument('--test_file', type=str, default='test_pairs_unpaired.txt',
                        help='测试文件名称（相对于 category 目录）')
    parser.add_argument('--mode', type=int, default=2,
                        help='模型模式 (1 或 2)')
    parser.add_argument('--extra_branch_num', type=int, default=2,
                        help='额外分支数量')
    parser.add_argument('--resolution', type=int, default=1024,
                        help='图像分辨率')
    parser.add_argument('--seed', type=int, default=0,
                        help='随机种子')
    parser.add_argument('--start_idx', type=int, default=None,
                        help='开始索引（手动分配任务时使用）')
    parser.add_argument('--end_idx', type=int, default=None,
                        help='结束索引（手动分配任务时使用）')
    
    args = parser.parse_args()
    
    # 初始化分布式环境
    rank, world_size, local_rank = setup_distributed()
    
    if rank is not None:
        # 多卡模式（使用 torchrun）
        device = f"cuda:{local_rank}"
        gpu_id = local_rank
        is_main_process = (rank == 0)
    else:
        # 单卡模式
        device = f"cuda:{args.gpu_id}"
        gpu_id = args.gpu_id
        is_main_process = True
        rank = 0
        world_size = 1
    
    # 确定要处理的类别
    if args.category == 'all':
        categories = ['dresses', 'lower_body', 'upper_body']
    else:
        categories = [args.category]
    
    # Checkpoint 映射：每个类别对应的checkpoint文件
    checkpoint_map = {
        'dresses': 'try_on_dress.pt',
        'lower_body': 'try_on_lower.pt',
        'upper_body': 'try_on_upper.pt',
    }
    
    # 处理每个类别
    for category in categories:
        if is_main_process:
            print(f"\n{'='*60}")
            print(f"处理类别: {category}")
            print(f"{'='*60}")
        
        # 确定该类别使用的checkpoint
        if args.checkpoint:
            # 如果用户指定了checkpoint，使用指定的
            category_checkpoint = args.checkpoint
        else:
            # 否则根据类别自动选择
            checkpoint_filename = checkpoint_map.get(category)
            if not checkpoint_filename:
                if is_main_process:
                    print(f"错误: 类别 {category} 没有对应的checkpoint映射")
                continue
            category_checkpoint = os.path.join(args.checkpoint_dir, checkpoint_filename)
        
        if not os.path.exists(category_checkpoint):
            if is_main_process:
                print(f"错误: Checkpoint 文件不存在: {category_checkpoint}")
            continue
        
        if is_main_process:
            print(f"使用 checkpoint: {category_checkpoint}")
        
        # 加载模型（每个类别加载对应的checkpoint）
        if is_main_process:
            print(f"加载模型...")
        
        pipe, transformer = load_model(
            category_checkpoint, device, args.mode, args.extra_branch_num
        )
        
        # 创建输出目录（每个类别一个子目录）
        category_output_dir = Path(args.output_dir) / category
        category_output_dir.mkdir(parents=True, exist_ok=True)
        
        # 加载测试对
        category_dir = os.path.join(args.dresscode_root, category)
        test_unpairs_file = os.path.join(category_dir, args.test_file)
        
        if not os.path.exists(test_unpairs_file):
            if is_main_process:
                print(f"警告: 找不到测试文件: {test_unpairs_file}，跳过类别 {category}")
            continue
        
        # 图像目录：DressCodeDataset 的图像都在 {category}/images/ 目录下
        images_dir = os.path.join(category_dir, "images")
        
        if not os.path.exists(images_dir):
            if is_main_process:
                print(f"警告: 找不到图像目录: {images_dir}，跳过类别 {category}")
            continue
        
        all_pairs = load_test_pairs(test_unpairs_file)
        total_pairs = len(all_pairs)
        
        # 确定任务分配
        if args.start_idx is not None and args.end_idx is not None:
            # 手动指定范围
            start_idx = args.start_idx
            end_idx = args.end_idx
        else:
            # 自动分配（基于 rank）
            per_gpu = total_pairs // world_size
            start_idx = rank * per_gpu
            if rank == world_size - 1:
                end_idx = total_pairs  # 最后一个 GPU 处理剩余的所有任务
            else:
                end_idx = start_idx + per_gpu
        
        pairs_to_process = all_pairs[start_idx:end_idx]
        
        # 预先检查有多少文件已存在
        existing_count = 0
        for person_img, cloth_img in pairs_to_process:
            output_filename = Path(person_img).stem + "_" + Path(cloth_img).stem + ".png"
            output_path = os.path.join(category_output_dir, output_filename)
            if os.path.exists(output_path):
                existing_count += 1
        
        if is_main_process:
            print(f"类别 {category}: 总共 {total_pairs} 对图像")
            print(f"GPU {gpu_id} (rank {rank}) 处理 {start_idx} 到 {end_idx} ({len(pairs_to_process)} 对)")
            print(f"图像目录: {images_dir}")
            print(f"输出目录: {category_output_dir}")
            if existing_count > 0:
                print(f"发现 {existing_count} 个已存在的文件，将跳过")
        
        # 处理图像对（模型已在循环外加载）
        success_count = 0
        fail_count = 0
        skip_count = 0
        
        if is_main_process:
            pbar = tqdm(pairs_to_process, desc=f"GPU {gpu_id} - {category}")
        else:
            pbar = pairs_to_process
        
        for person_img, cloth_img in pbar:
            person_path = os.path.join(images_dir, person_img)
            cloth_path = os.path.join(images_dir, cloth_img)
            
            # 检查输出文件是否已存在
            output_filename = Path(person_img).stem + "_" + Path(cloth_img).stem + ".png"
            output_path = os.path.join(category_output_dir, output_filename)
            
            if os.path.exists(output_path):
                # 输出文件已存在，跳过
                skip_count += 1
                if is_main_process and skip_count % 100 == 0:
                    # 每跳过100个文件输出一次日志，避免日志过多
                    print(f"[GPU {gpu_id}] 已跳过 {skip_count} 个已存在的文件...")
                continue
            
            if not os.path.exists(person_path):
                if is_main_process:
                    print(f"[GPU {gpu_id}] 警告: 找不到人物图像 {person_path}")
                fail_count += 1
                continue
            
            if not os.path.exists(cloth_path):
                if is_main_process:
                    print(f"[GPU {gpu_id}] 警告: 找不到衣服图像 {cloth_path}")
                fail_count += 1
                continue
            
            success, error = process_pair(
                person_path, cloth_path, pipe, device, category_output_dir,
                args.resolution, args.mode, args.seed
            )
            
            if success:
                success_count += 1
            else:
                fail_count += 1
                if is_main_process:
                    print(f"[GPU {gpu_id}] 处理失败 {person_img} + {cloth_img}: {error}")
        
        if is_main_process:
            print(f"\n类别 {category} - GPU {gpu_id} 完成:")
            print(f"  成功: {success_count}")
            print(f"  跳过: {skip_count} (已存在)")
            print(f"  失败: {fail_count}")
            print(f"  输出目录: {category_output_dir}")
    
    # 清理分布式环境
    if rank is not None:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()