import os import math import PIL import numpy as np import torch from PIL import Image from accelerate.state import AcceleratorState from packaging import version import accelerate from typing import List, Optional, Tuple from torch.nn import functional as F from diffusers import UNet2DConditionModel, SchedulerMixin # Compute DREAM and update latents for diffusion sampling def compute_dream_and_update_latents_for_inpaint( unet: UNet2DConditionModel, noise_scheduler: SchedulerMixin, timesteps: torch.Tensor, noise: torch.Tensor, noisy_latents: torch.Tensor, target: torch.Tensor, encoder_hidden_states: torch.Tensor, dream_detail_preservation: float = 1.0, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """ Implements "DREAM (Diffusion Rectification and Estimation-Adaptive Models)" from http://arxiv.org/abs/2312.00210. DREAM helps align training with sampling to help training be more efficient and accurate at the cost of an extra forward step without gradients. Args: `unet`: The state unet to use to make a prediction. `noise_scheduler`: The noise scheduler used to add noise for the given timestep. `timesteps`: The timesteps for the noise_scheduler to user. `noise`: A tensor of noise in the shape of noisy_latents. `noisy_latents`: Previously noise latents from the training loop. `target`: The ground-truth tensor to predict after eps is removed. `encoder_hidden_states`: Text embeddings from the text model. `dream_detail_preservation`: A float value that indicates detail preservation level. See reference. Returns: `tuple[torch.Tensor, torch.Tensor]`: Adjusted noisy_latents and target. """ alphas_cumprod = noise_scheduler.alphas_cumprod.to(timesteps.device)[timesteps, None, None, None] sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 # The paper uses lambda = sqrt(1 - alpha) ** p, with p = 1 in their experiments. dream_lambda = sqrt_one_minus_alphas_cumprod**dream_detail_preservation pred = None # b, 4, h, w with torch.no_grad(): pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample noisy_latents_no_condition = noisy_latents[:, :4] _noisy_latents, _target = (None, None) if noise_scheduler.config.prediction_type == "epsilon": predicted_noise = pred delta_noise = (noise - predicted_noise).detach() delta_noise.mul_(dream_lambda) _noisy_latents = noisy_latents_no_condition.add(sqrt_one_minus_alphas_cumprod * delta_noise) _target = target.add(delta_noise) elif noise_scheduler.config.prediction_type == "v_prediction": raise NotImplementedError("DREAM has not been implemented for v-prediction") else: raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}") _noisy_latents = torch.cat([_noisy_latents, noisy_latents[:, 4:]], dim=1) return _noisy_latents, _target # Prepare the input for inpainting model. def prepare_inpainting_input( noisy_latents: torch.Tensor, mask_latents: torch.Tensor, condition_latents: torch.Tensor, enable_condition_noise: bool = True, condition_concat_dim: int = -1, ) -> torch.Tensor: """ Prepare the input for inpainting model. Args: noisy_latents (torch.Tensor): Noisy latents. mask_latents (torch.Tensor): Mask latents. condition_latents (torch.Tensor): Condition latents. enable_condition_noise (bool): Enable condition noise. Returns: torch.Tensor: Inpainting input. """ if not enable_condition_noise: condition_latents_ = condition_latents.chunk(2, dim=condition_concat_dim)[-1] noisy_latents = torch.cat([noisy_latents, condition_latents_], dim=condition_concat_dim) noisy_latents = torch.cat([noisy_latents, mask_latents, condition_latents], dim=1) return noisy_latents # Compute VAE encodings def compute_vae_encodings(image: torch.Tensor, vae: torch.nn.Module) -> torch.Tensor: """ Args: images (torch.Tensor): image to be encoded vae (torch.nn.Module): vae model Returns: torch.Tensor: latent encoding of the image """ pixel_values = image.to(memory_format=torch.contiguous_format).float() pixel_values = pixel_values.to(vae.device, dtype=vae.dtype) with torch.no_grad(): model_input = vae.encode(pixel_values).latent_dist.sample() model_input = model_input * vae.config.scaling_factor return model_input # Init Accelerator from accelerate import Accelerator, DistributedDataParallelKwargs from accelerate.utils import ProjectConfiguration def init_accelerator(config): accelerator_project_config = ProjectConfiguration( project_dir=config.project_name, logging_dir=os.path.join(config.project_name, "logs"), ) accelerator_ddp_config = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator( mixed_precision=config.mixed_precision, log_with=config.report_to, project_config=accelerator_project_config, kwargs_handlers=[accelerator_ddp_config], gradient_accumulation_steps=config.gradient_accumulation_steps, ) # Disable AMP for MPS. if torch.backends.mps.is_available(): accelerator.native_amp = False if accelerator.is_main_process: accelerator.init_trackers( project_name=config.project_name, config={ "learning_rate": config.learning_rate, "train_batch_size": config.train_batch_size, "image_size": f"{config.width}x{config.height}", }, ) return accelerator def init_weight_dtype(wight_dtype): return { "no": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16, }[wight_dtype] def init_add_item_id(config): return torch.tensor( [ config.height, config.width * 2, 0, 0, config.height, config.width * 2, ] ).repeat(config.train_batch_size, 1) def prepare_eval_data(dataset_root, dataset_name, is_pair=True): assert dataset_name in ["vitonhd", "dresscode", "farfetch"], "Unknown dataset name {}.".format(dataset_name) if dataset_name == "vitonhd": data_root = os.path.join(dataset_root, "VITONHD-1024", "test") if is_pair: keys = os.listdir(os.path.join(data_root, "Images")) cloth_image_paths = [ os.path.join(data_root, "Images", key, key + "-0.jpg") for key in keys ] person_image_paths = [ os.path.join(data_root, "Images", key, key + "-1.jpg") for key in keys ] else: # read ../test_pairs.txt cloth_image_paths = [] person_image_paths = [] with open( os.path.join(dataset_root, "VITONHD-1024", "test_pairs.txt"), "r" ) as f: lines = f.readlines() for line in lines: cloth_image, person_image = ( line.replace(".jpg", "").strip().split(" ") ) cloth_image_paths.append( os.path.join( data_root, "Images", cloth_image, cloth_image + "-0.jpg" ) ) person_image_paths.append( os.path.join( data_root, "Images", person_image, person_image + "-1.jpg" ) ) elif dataset_name == "dresscode": data_root = os.path.join(dataset_root, "DressCode-1024") if is_pair: part = ["lower", "lower", "upper", "upper", "dresses", "dresses"] ids = ["013581", "051685", "000190", "050072", "020829", "053742"] cloth_image_paths = [ os.path.join(data_root, "Images", part[i], ids[i], ids[i] + "_1.jpg") for i in range(len(part)) ] person_image_paths = [ os.path.join(data_root, "Images", part[i], ids[i], ids[i] + "_0.jpg") for i in range(len(part)) ] else: raise ValueError("DressCode dataset does not support non-pair evaluation.") elif dataset_name == "farfetch": data_root = os.path.join(dataset_root, "FARFETCH-1024") cloth_image_paths = [ # TryOn "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Blouses/13732751/13732751-2.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Hoodies/14661627/14661627-4.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Vests & Tank Tops/16532697/16532697-4.jpg", "Images/men/Pants/Loose Fit Pants/14750720/14750720-6.jpg", # Garment Transfer "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Shirts/10889688/10889688-3.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Shorts/Leather & Faux Leather Shorts/20143338/20143338-1.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Jackets/Blazers/15541224/15541224-2.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/men/Polo Shirts/Polo Shirts/17652415/17652415-0.jpg" # "Images/men/Jackets/Hooded Jackets/12550261/12550261-1.jpg", # "Images/men/Shirts/Shirts/15614589/15614589-4.jpg", # "Images/women/Dresses/Day Dresses/10372515/10372515-3.jpg", # "Images/women/Dresses/Sundresses/18520992/18520992-4.jpg", # "Images/women/Skirts/Asymmetric & Draped Skirts/12404908/12404908-2.jpg", ] person_image_paths = [ # TryOn "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Blouses/13732751/13732751-0.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Hoodies/14661627/14661627-2.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Vests & Tank Tops/16532697/16532697-1.jpg", "Images/men/Pants/Loose Fit Pants/14750720/14750720-5.jpg", # Garment Transfer "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Tops/Shirts/10889688/10889688-1.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Shorts/Leather & Faux Leather Shorts/20143338/20143338-2.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/women/Jackets/Blazers/15541224/15541224-0.jpg", "/home/chongzheng/Projects/hivton/Datasets/FARFETCH-1024/Images/men/Polo Shirts/Polo Shirts/17652415/17652415-4.jpg", # "Images/men/Jackets/Hooded Jackets/12550261/12550261-3.jpg", # "Images/men/Shirts/Shirts/15614589/15614589-3.jpg", # "Images/women/Dresses/Day Dresses/10372515/10372515-0.jpg", # "Images/women/Dresses/Sundresses/18520992/18520992-1.jpg", # "Images/women/Skirts/Asymmetric & Draped Skirts/12404908/12404908-1.jpg", ] cloth_image_paths = [ os.path.join(data_root, path) for path in cloth_image_paths ] person_image_paths = [ os.path.join(data_root, path) for path in person_image_paths ] else: raise ValueError(f"Unknown dataset name: {dataset_name}") samples = [ { "folder": os.path.basename(os.path.dirname(cloth_image)), "cloth": cloth_image, "person": person_image, } for cloth_image, person_image in zip( cloth_image_paths, person_image_paths ) ] return samples def repaint_result(result, person_image, mask_image): result, person, mask = np.array(result), np.array(person_image), np.array(mask_image) # expand the mask to 3 channels & to 0~1 mask = np.expand_dims(mask, axis=2) mask = mask / 255.0 # mask for result, ~mask for person result_ = result * mask + person * (1 - mask) return Image.fromarray(result_.astype(np.uint8)) # 多通道 Sobel 算子处理 (用于获取模特图像的损失注意力图) def sobel(batch_image, mask=None, scale=4.0): """ 计算输入批量图像的Sobel梯度. batch_image: 输入的批量图像张量，大小为 [batch, channels, height, width] """ w, h = batch_image.size(3), batch_image.size(2) pool_kernel = (max(w, h) // 16) * 2 + 1 # 定义Sobel核 kernel_x = ( torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32) .view(1, 1, 3, 3) .to(batch_image.device) .repeat(1, batch_image.size(1), 1, 1) ) kernel_y = ( torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32) .view(1, 1, 3, 3) .to(batch_image.device) .repeat(1, batch_image.size(1), 1, 1) ) # 初始化梯度张量 grad_x = torch.zeros_like(batch_image) grad_y = torch.zeros_like(batch_image) # 边缘填充 batch_image = F.pad(batch_image, (1, 1, 1, 1), mode="reflect") # 应用Sobel算子 grad_x = F.conv2d(batch_image, kernel_x, padding=0) grad_y = F.conv2d(batch_image, kernel_y, padding=0) # 计算梯度的幅度 grad_magnitude = torch.sqrt(grad_x.pow(2) + grad_y.pow(2)) # Mask 处理 if mask is not None: grad_magnitude = grad_magnitude * mask # 剃度裁剪 grad_magnitude = torch.clamp(grad_magnitude, 0.2, 2.5) # 平均池化 grad_magnitude = F.avg_pool2d( grad_magnitude, kernel_size=pool_kernel, stride=1, padding=pool_kernel // 2 ) # 归一化 grad_magnitude = (grad_magnitude / grad_magnitude.max()) * scale return grad_magnitude # Sobel 加权平方误差, 增强边缘区域的损失（直接用于损失计算） def sobel_aug_squared_error(x, y, reference, mask=None, reduction="mean"): """ 计算x,y的逐元素平方误差，其中x和y是图像张量. 然后利用 x 的 sobel 结果作为权重，计算加权平方误差. x: Tensor, shape [batch, channels, height, width] y: Tensor, shape [batch, channels, height, width] """ ref_sobel = sobel(reference, mask=mask) # 计算 sobel 梯度作为损失权重 if ref_sobel.isnan().any(): print("Error: NaN Sobel Gradient") loss = F.mse_loss(x, y, reduction="mean") # 如果梯度为nan，则直接退化为MSE损失 else: squared_error = (x - y).pow(2) weighted_squared_error = squared_error * ref_sobel if reduction == "mean": loss = weighted_squared_error.mean() elif reduction == "sum": loss = weighted_squared_error.sum() elif reduction == "none": loss = weighted_squared_error # print("WSE Loss:", loss.mean(), loss.dtype) return loss # 准备图像（转换为 Batch 张量） def prepare_image(image): if isinstance(image, torch.Tensor): # Batch single image if image.ndim == 3: image = image.unsqueeze(0) image = image.to(dtype=torch.float32) else: # preprocess image if isinstance(image, (PIL.Image.Image, np.ndarray)): image = [image] if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): image = [np.array(i.convert("RGB"))[None, :] for i in image] image = np.concatenate(image, axis=0) elif isinstance(image, list) and isinstance(image[0], np.ndarray): image = np.concatenate([i[None, :] for i in image], axis=0) image = image.transpose(0, 3, 1, 2) image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 return image def prepare_mask_image(mask_image): if isinstance(mask_image, torch.Tensor): if mask_image.ndim == 2: # Batch and add channel dim for single mask mask_image = mask_image.unsqueeze(0).unsqueeze(0) elif mask_image.ndim == 3 and mask_image.shape[0] == 1: # Single mask, the 0'th dimension is considered to be # the existing batch size of 1 mask_image = mask_image.unsqueeze(0) elif mask_image.ndim == 3 and mask_image.shape[0] != 1: # Batch of mask, the 0'th dimension is considered to be # the batching dimension mask_image = mask_image.unsqueeze(1) # Binarize mask mask_image[mask_image < 0.5] = 0 mask_image[mask_image >= 0.5] = 1 else: # preprocess mask if isinstance(mask_image, (PIL.Image.Image, np.ndarray)): mask_image = [mask_image] if isinstance(mask_image, list) and isinstance(mask_image[0], PIL.Image.Image): mask_image = np.concatenate( [np.array(m.convert("L"))[None, None, :] for m in mask_image], axis=0 ) mask_image = mask_image.astype(np.float32) / 255.0 elif isinstance(mask_image, list) and isinstance(mask_image[0], np.ndarray): mask_image = np.concatenate([m[None, None, :] for m in mask_image], axis=0) mask_image[mask_image < 0.5] = 0 mask_image[mask_image >= 0.5] = 1 mask_image = torch.from_numpy(mask_image) return mask_image def numpy_to_pil(images): """ Convert a numpy image or a batch of images to a PIL image. """ 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 def load_eval_image_pairs(root, mode="logo"): # TODO 加载测试图像对，包括配对和非配对的图像对 test_name = "test" person_image_paths = [ os.path.join(root, test_name, "image", _) for _ in os.listdir(os.path.join(root, test_name, "image")) if _.endswith(".jpg") ] cloth_image_paths = [ person_image_path.replace("image", "cloth") for person_image_path in person_image_paths ] # 包含图案和文字的部分图像 if mode == "logo": filter_pairs = [ 6648, 6744, 6967, 6985, 14031, 12358, 4963, 4680, 499, 396, 345, 6648, 6744, 6967, 6985, 7510, 8205, 8254, 10545, 11485, 11632, 12354, 13144, 14112, 12570, 11766, ] filter_pairs.sort() filter_pairs = [f"{_:05d}_00.jpg" for _ in filter_pairs] cloth_image_paths = [ cloth_image_paths[i] for i in range(len(cloth_image_paths)) if os.path.basename(cloth_image_paths[i]) in filter_pairs ] person_image_paths = [ person_image_paths[i] for i in range(len(person_image_paths)) if os.path.basename(person_image_paths[i]) in filter_pairs ] return cloth_image_paths, person_image_paths def tensor_to_image(tensor: torch.Tensor): """ Converts a torch tensor to PIL Image. """ assert tensor.dim() == 3, "Input tensor should be 3-dimensional." assert tensor.dtype == torch.float32, "Input tensor should be float32." assert ( tensor.min() >= 0 and tensor.max() <= 1 ), "Input tensor should be in range [0, 1]." tensor = tensor.cpu() tensor = tensor * 255 tensor = tensor.permute(1, 2, 0) tensor = tensor.numpy().astype(np.uint8) image = Image.fromarray(tensor) return image def concat_images(images: List[Image.Image], divider: int = 4, cols: int = 4): """ Concatenates images horizontally and with """ widths = [image.size[0] for image in images] heights = [image.size[1] for image in images] total_width = cols * max(widths) total_width += divider * (cols - 1) # `col` images each row rows = math.ceil(len(images) / cols) total_height = max(heights) * rows # add divider between rows total_height += divider * (len(heights) // cols - 1) # all black image concat_image = Image.new("RGB", (total_width, total_height), (0, 0, 0)) x_offset = 0 y_offset = 0 for i, image in enumerate(images): concat_image.paste(image, (x_offset, y_offset)) x_offset += image.size[0] + divider if (i + 1) % cols == 0: x_offset = 0 y_offset += image.size[1] + divider return concat_image def read_prompt_file(prompt_file: str): if prompt_file is not None and os.path.isfile(prompt_file): with open(prompt_file, "r") as sample_prompt_file: sample_prompts = sample_prompt_file.readlines() sample_prompts = [sample_prompt.strip() for sample_prompt in sample_prompts] else: sample_prompts = [] return sample_prompts def save_tensors_to_npz(tensors: torch.Tensor, paths: List[str]): assert len(tensors) == len(paths), "Length of tensors and paths should be the same!" for tensor, path in zip(tensors, paths): np.savez_compressed(path, latent=tensor.cpu().numpy()) def deepspeed_zero_init_disabled_context_manager(): """ returns either a context list that includes one that will disable zero.Init or an empty context list """ deepspeed_plugin = ( AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None ) if deepspeed_plugin is None: return [] return [deepspeed_plugin.zero3_init_context_manager(enable=False)] def is_xformers_available(): try: import xformers xformers_version = version.parse(xformers.__version__) if xformers_version == version.parse("0.0.16"): print( "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, " "please update xFormers to at least 0.0.17. " "See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details." ) return True except ImportError: raise ValueError( "xformers is not available. Make sure it is installed correctly" ) def resize_and_crop(image, size): # Crop to size ratio w, h = image.size target_w, target_h = size if w / h < target_w / target_h: new_w = w new_h = w * target_h // target_w else: new_h = h new_w = h * target_w // target_h image = image.crop( ((w - new_w) // 2, (h - new_h) // 2, (w + new_w) // 2, (h + new_h) // 2) ) # resize image = image.resize(size, Image.LANCZOS) return image def resize_and_padding(image, size): # Padding to size ratio w, h = image.size target_w, target_h = size if w / h < target_w / target_h: new_h = target_h new_w = w * target_h // h else: new_w = target_w new_h = h * target_w // w image = image.resize((new_w, new_h), Image.LANCZOS) # padding padding = Image.new("RGB", size, (255, 255, 255)) padding.paste(image, ((target_w - new_w) // 2, (target_h - new_h) // 2)) return padding if __name__ == "__main__": import torch import torch.nn.functional as F from torchvision import transforms from PIL import Image, ImageFilter import numpy as np def vis_sobel_weight(image_path, mask_path) -> PIL.Image.Image: image = Image.open(image_path).convert("RGB") w, h = image.size l_w, l_h = w // 8, h // 8 image = image.resize((l_w, l_h)) mask = Image.open(mask_path).convert("L").resize((l_w, l_h)) image_pt = transforms.ToTensor()(image).unsqueeze(0).to("cuda") mask_pt = transforms.ToTensor()(mask).unsqueeze(0).to("cuda") sobel_pt = sobel(image_pt, mask_pt, scale=1.0) sobel_image = sobel_pt.squeeze().cpu().numpy() sobel_image = Image.fromarray((sobel_image * 255).astype(np.uint8)) sobel_image = sobel_image.resize((w, h), resample=Image.NEAREST) # 图像平滑 sobel_image = sobel_image.filter(ImageFilter.SMOOTH) from data.utils import grayscale_to_heatmap sobel_image = grayscale_to_heatmap(sobel_image) image = Image.open(image_path).convert("RGB").resize((w, h)) sobel_image = Image.blend(image, sobel_image, alpha=0.5) return sobel_image save_folder = "./sobel_vis-2.0" if not os.path.exists(save_folder): os.makedirs(save_folder) from data.utils import scan_files_in_dir for i in scan_files_in_dir( "/home/chongzheng/Projects/try-on-project/Datasets/VITONHD-1024/test/Images" ): image_path = i.path if i.path.endswith("-1.jpg"): result_path = os.path.join(save_folder, os.path.basename(image_path)) mask_path = image_path.replace("Images", "AgnosticMask").replace( "-1.jpg", "_mask-1.png" ) vis_sobel_weight(image_path, mask_path).save(result_path) pass