""" 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()