import os import numpy as np import torch import argparse from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler from diffusers.image_processor import VaeImageProcessor from tqdm import tqdm from PIL import Image, ImageFilter from model.pipeline import CatVTONPipeline, CatVTONPix2PixPipeline """ torchrun --nproc_per_node=2 inference.py --mask_free --resume_path zhengchong/CatVTON-MaskFree --dataset_name vitonhd --data_root_path /filesdir/VITONDataset --output_dir output1 --batch_size 8 --num_inference_steps 50 --guidance_scale 2.5 --height 1024 --width 768 --mixed_precision bf16""" """nohup torchrun --nproc_per_node=2 inference.py \ --mask_free \ --dataset_name dresscode \ --data_root_path /filesdir/DressCodeDataset \ --output_dir output1 \ --batch_size 8 \ --num_inference_steps 50 \ --guidance_scale 2.5 \ --height 1024 \ --width 768 \ --mixed_precision bf16 > dresscode_inference.log 2>&1 &""" class InferenceDataset(Dataset): def __init__(self, args): self.args = args self.vae_processor = VaeImageProcessor(vae_scale_factor=8) self.mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True) self.data = self.load_data() def load_data(self): return [] def __len__(self): return len(self.data) def __getitem__(self, idx): data = self.data[idx] person = Image.open(data['person']) cloth = Image.open(data['cloth']) result = { 'index': idx, 'person_name': data['person_name'], 'person': self.vae_processor.preprocess(person, self.args.height, self.args.width)[0], 'cloth': self.vae_processor.preprocess(cloth, self.args.height, self.args.width)[0], } # 只有在非mask-free模式下才加载mask if not self.args.mask_free and 'mask' in data: mask = Image.open(data['mask']) result['mask'] = self.mask_processor.preprocess(mask, self.args.height, self.args.width)[0] return result class VITONHDTestDataset(InferenceDataset): def load_data(self): # 查找test目录下的test_unpairs.txt test_dir = os.path.join(self.args.data_root_path, "test") pair_txt = os.path.join(test_dir, "test_unpairs.txt") if not os.path.exists(pair_txt): raise FileNotFoundError(f"Pair file not found: {pair_txt}") with open(pair_txt, 'r') as f: lines = f.readlines() self.args.data_root_path = os.path.join(self.args.data_root_path, "test") output_dir = os.path.join(self.args.output_dir, "vitonhd", 'unpaired' if not self.args.eval_pair else 'paired') data = [] for line in lines: line = line.strip() if not line: # 跳过空行 continue parts = line.split() if len(parts) < 2: continue person_img, cloth_img = parts[0], parts[1] if os.path.exists(os.path.join(output_dir, person_img)): continue if self.args.eval_pair: cloth_img = person_img item = { 'person_name': person_img, 'person': os.path.join(self.args.data_root_path, 'image', person_img), 'cloth': os.path.join(self.args.data_root_path, 'cloth', cloth_img), } # 只有在非mask-free模式下才添加mask路径 if not self.args.mask_free: # 支持两种mask文件名格式 mask_name = person_img.replace('.jpg', '_mask.png') mask_path = os.path.join(self.args.data_root_path, 'agnostic-mask', mask_name) if not os.path.exists(mask_path): # 尝试另一种格式 mask_name = person_img.replace('.jpg', '.png') mask_path = os.path.join(self.args.data_root_path, 'agnostic-mask', mask_name) item['mask'] = mask_path data.append(item) return data class DressCodeTestDataset(InferenceDataset): def load_data(self): data = [] for sub_folder in ['upper_body', 'lower_body', 'dresses']: assert os.path.exists(os.path.join(self.args.data_root_path, sub_folder)), f"Folder {sub_folder} does not exist." pair_txt = os.path.join(self.args.data_root_path, sub_folder, 'test_pairs_paired.txt' if self.args.eval_pair else 'test_pairs_unpaired.txt') assert os.path.exists(pair_txt), f"File {pair_txt} does not exist." with open(pair_txt, 'r') as f: lines = f.readlines() output_dir = os.path.join(self.args.output_dir, f"dresscode-{self.args.height}", 'unpaired' if not self.args.eval_pair else 'paired', sub_folder) sub_folder_count = 0 for line in lines: line = line.strip() if not line: # 跳过空行 continue parts = line.split() if len(parts) < 2: # 跳过格式不正确的行 continue person_img, cloth_img = parts[0], parts[1] if os.path.exists(os.path.join(output_dir, person_img)): continue person_path = os.path.join(self.args.data_root_path, sub_folder, 'images', person_img) cloth_path = os.path.join(self.args.data_root_path, sub_folder, 'images', cloth_img) # 打印前5个示例 if sub_folder_count < 5: print(f"[{sub_folder}] Pair {sub_folder_count + 1}:") print(f" Person: {person_img} -> {person_path} (exists: {os.path.exists(person_path)})") print(f" Cloth: {cloth_img} -> {cloth_path} (exists: {os.path.exists(cloth_path)})") sub_folder_count += 1 item = { 'person_name': os.path.join(sub_folder, person_img), 'person': person_path, 'cloth': cloth_path, } # 只有在非mask-free模式下才添加mask路径 if not self.args.mask_free: item['mask'] = os.path.join(self.args.data_root_path, sub_folder, 'agnostic_masks', person_img.replace('.jpg', '.png')) data.append(item) return data def parse_args(): parser = argparse.ArgumentParser(description="Simple example of a training script.") parser.add_argument( "--base_model_path", type=str, default="booksforcharlie/stable-diffusion-inpainting", # Change to a copy repo as runawayml delete original repo help=( "The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub." ), ) parser.add_argument( "--resume_path", type=str, default="zhengchong/CatVTON", help=( "The Path to the checkpoint of trained tryon model. Use 'zhengchong/CatVTON-MaskFree' for mask-free version." ), ) parser.add_argument( "--dataset_name", type=str, required=True, help="The datasets to use for evaluation.", ) parser.add_argument( "--data_root_path", type=str, required=True, help="Path to the dataset to evaluate." ) parser.add_argument( "--output_dir", type=str, default="output", help="The output directory where the model predictions will be written.", ) parser.add_argument( "--seed", type=int, default=555, help="A seed for reproducible evaluation." ) parser.add_argument( "--batch_size", type=int, default=8, help="The batch size for evaluation." ) parser.add_argument( "--num_inference_steps", type=int, default=50, help="Number of inference steps to perform.", ) parser.add_argument( "--guidance_scale", type=float, default=2.5, help="The scale of classifier-free guidance for inference.", ) parser.add_argument( "--width", type=int, default=768, help=( "The resolution for input images, all the images in the train/validation dataset will be resized to this" " resolution" ), ) parser.add_argument( "--height", type=int, default=1024, help=( "The resolution for input images, all the images in the train/validation dataset will be resized to this" " resolution" ), ) parser.add_argument( "--repaint", action="store_true", help="Whether to repaint the result image with the original background." ) parser.add_argument( "--eval_pair", action="store_true", help="Whether or not to evaluate the pair.", ) parser.add_argument( "--concat_eval_results", action="store_true", help="Whether or not to concatenate the all conditions into one image.", ) parser.add_argument( "--allow_tf32", action="store_true", default=True, help=( "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" ), ) parser.add_argument( "--dataloader_num_workers", type=int, default=8, help=( "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." ), ) parser.add_argument( "--mixed_precision", type=str, default="bf16", choices=["no", "fp16", "bf16"], help=( "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" " 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." ), ) parser.add_argument( "--concat_axis", type=str, choices=["x", "y", 'random'], default="y", help="The axis to concat the cloth feature, select from ['x', 'y', 'random'].", ) parser.add_argument( "--enable_condition_noise", action="store_true", default=True, help="Whether or not to enable condition noise.", ) parser.add_argument( "--mask_free", action="store_true", help="Whether to use mask-free version (CatVTON-MaskFree). If enabled, mask will not be required.", ) parser.add_argument( "--p2p_base_model_path", type=str, default="timbrooks/instruct-pix2pix", help="The base model path for mask-free version (pix2pix model).", ) parser.add_argument( "--local_rank", type=int, default=-1, help="Local rank for distributed training/inference.", ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) if env_local_rank != -1 and env_local_rank != args.local_rank: args.local_rank = env_local_rank return args def repaint(person, mask, result): _, h = result.size kernal_size = h // 50 if kernal_size % 2 == 0: kernal_size += 1 mask = mask.filter(ImageFilter.GaussianBlur(kernal_size)) person_np = np.array(person) result_np = np.array(result) mask_np = np.array(mask) / 255 repaint_result = person_np * (1 - mask_np) + result_np * mask_np repaint_result = Image.fromarray(repaint_result.astype(np.uint8)) return repaint_result def to_pil_image(images): images = (images / 2 + 0.5).clamp(0, 1) images = images.cpu().permute(0, 2, 3, 1).float().numpy() if images.ndim == 3: images = images[None, ...] images = (images * 255).round().astype("uint8") if images.shape[-1] == 1: # special case for grayscale (single channel) images pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] else: pil_images = [Image.fromarray(image) for image in images] return pil_images @torch.no_grad() def main(): args = parse_args() # 初始化分布式环境（如果使用多GPU） if args.local_rank != -1: torch.distributed.init_process_group(backend='nccl') torch.cuda.set_device(args.local_rank) device = f"cuda:{args.local_rank}" world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() print(f"Initialized distributed inference: rank {rank}/{world_size-1}, device {device}") else: device = "cuda" world_size = 1 rank = 0 # Pipeline if args.mask_free: # Mask-free version uses CatVTONPix2PixPipeline base_model = args.p2p_base_model_path # 如果使用默认的 resume_path，自动改为 mask-free 版本 if args.resume_path == "zhengchong/CatVTON": checkpoint = "zhengchong/CatVTON-MaskFree" else: checkpoint = args.resume_path pipeline = CatVTONPix2PixPipeline( attn_ckpt_version="mix-48k-1024", attn_ckpt=checkpoint, base_ckpt=args.p2p_base_model_path, weight_dtype={ "no": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16, }[args.mixed_precision], device=device, skip_safety_check=True ) if rank == 0: print(f"Model Configuration (Mask-Free):") print(f" Pipeline: CatVTONPix2PixPipeline") print(f" Base Model: {base_model}") print(f" Checkpoint/LoRA: {checkpoint}") print(f" Weight Dtype: {args.mixed_precision}") print(f" Device: {device}") else: # Original version uses CatVTONPipeline (requires mask) base_model = args.base_model_path checkpoint = args.resume_path pipeline = CatVTONPipeline( attn_ckpt_version=args.dataset_name, attn_ckpt=args.resume_path, base_ckpt=args.base_model_path, weight_dtype={ "no": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16, }[args.mixed_precision], device=device, skip_safety_check=True ) if rank == 0: print(f"Model Configuration (Original):") print(f" Pipeline: CatVTONPipeline") print(f" Base Model: {base_model}") print(f" Checkpoint/LoRA: {checkpoint}") print(f" Weight Dtype: {args.mixed_precision}") print(f" Device: {device}") # Dataset if args.dataset_name == "vitonhd": dataset = VITONHDTestDataset(args) elif args.dataset_name == "dresscode": dataset = DressCodeTestDataset(args) else: raise ValueError(f"Invalid dataset name {args.dataset}.") print(f"Dataset {args.dataset_name} loaded, total {len(dataset)} pairs.") # 使用DistributedSampler来分配数据到不同的GPU if world_size > 1: sampler = DistributedSampler( dataset, num_replicas=world_size, rank=rank, shuffle=False ) else: sampler = None dataloader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.dataloader_num_workers, sampler=sampler ) # Inference generator = torch.Generator(device=device).manual_seed(args.seed) args.output_dir = os.path.join(args.output_dir, f"{args.dataset_name}-{args.height}", "paired" if args.eval_pair else "unpaired") if not os.path.exists(args.output_dir): os.makedirs(args.output_dir, exist_ok=True) # 只在主进程显示进度条 if rank == 0: pbar = tqdm(dataloader, desc=f"Processing on {world_size} GPU(s)") else: pbar = dataloader for batch in pbar: person_images = batch['person'] cloth_images = batch['cloth'] # 根据是否使用mask-free模式调用不同的pipeline if args.mask_free: # Mask-free版本不需要mask参数 results = pipeline( image=person_images, condition_image=cloth_images, num_inference_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, height=args.height, width=args.width, generator=generator, ) else: # 原始版本需要mask参数 masks = batch['mask'] results = pipeline( person_images, cloth_images, masks, num_inference_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, height=args.height, width=args.width, generator=generator, ) if args.concat_eval_results or args.repaint: person_images = to_pil_image(person_images) cloth_images = to_pil_image(cloth_images) if not args.mask_free: masks = to_pil_image(masks) for i, result in enumerate(results): person_name = batch['person_name'][i] output_path = os.path.join(args.output_dir, person_name) if not os.path.exists(os.path.dirname(output_path)): os.makedirs(os.path.dirname(output_path)) if args.repaint and not args.mask_free: # repaint功能只在有mask时可用 person_path, mask_path = dataset.data[batch['index'][i]]['person'], dataset.data[batch['index'][i]]['mask'] person_image= Image.open(person_path).resize(result.size, Image.LANCZOS) mask = Image.open(mask_path).resize(result.size, Image.NEAREST) result = repaint(person_image, mask, result) if args.concat_eval_results: w, h = result.size concated_result = Image.new('RGB', (w*3, h)) concated_result.paste(person_images[i], (0, 0)) concated_result.paste(cloth_images[i], (w, 0)) concated_result.paste(result, (w*2, 0)) result = concated_result result.save(output_path) if __name__ == "__main__": main()