test.py

import argparse, os, sys, glob
import cv2
import torch
import numpy as np
from omegaconf import OmegaConf
from PIL import Image
from torch.utils.data import DataLoader
from torchvision import transforms
from tqdm import tqdm, trange
from itertools import islice
from einops import rearrange
from torchvision.utils import make_grid
import time
from pytorch_lightning import seed_everything
from torch import autocast
from contextlib import contextmanager, nullcontext
import torchvision

from ldm.data.cp_dataset import CPDataset
from ldm.resizer import Resizer
from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler
from ldm.data.deepfashions import DFPairDataset

import clip
from torchvision.transforms import Resize


def chunk(it, size):
    it = iter(it)
    return iter(lambda: tuple(islice(it, size)), ())


def get_tensor_clip(normalize=True, toTensor=True):
    transform_list = []
    if toTensor:
        transform_list += [torchvision.transforms.ToTensor()]

    if normalize:
        transform_list += [torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
                                                            (0.26862954, 0.26130258, 0.27577711))]
    return torchvision.transforms.Compose(transform_list)


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")
    pil_images = [Image.fromarray(image) for image in images]

    return pil_images


def load_model_from_config(config, ckpt, verbose=False):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu")
    if "global_step" in pl_sd:
        print(f"Global Step: {pl_sd['global_step']}")
    sd = pl_sd["state_dict"]
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    if len(m) > 0 and verbose:
        print("missing keys:")
        print(m)
    if len(u) > 0 and verbose:
        print("unexpected keys:")
        print(u)

    model.cuda()
    model.eval()
    return model


def put_watermark(img, wm_encoder=None):
    if wm_encoder is not None:
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        img = wm_encoder.encode(img, 'dwtDct')
        img = Image.fromarray(img[:, :, ::-1])
    return img


def load_replacement(x):
    try:
        hwc = x.shape
        y = Image.open("assets/rick.jpeg").convert("RGB").resize((hwc[1], hwc[0]))
        y = (np.array(y) / 255.0).astype(x.dtype)
        assert y.shape == x.shape
        return y
    except Exception:
        return x


def get_tensor(normalize=True, toTensor=True):
    transform_list = []
    if toTensor:
        transform_list += [torchvision.transforms.ToTensor()]

    if normalize:
        transform_list += [torchvision.transforms.Normalize((0.5, 0.5, 0.5),
                                                            (0.5, 0.5, 0.5))]
    return torchvision.transforms.Compose(transform_list)


def get_tensor_clip(normalize=True, toTensor=True):
    transform_list = []
    if toTensor:
        transform_list += [torchvision.transforms.ToTensor()]

    if normalize:
        transform_list += [torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
                                                            (0.26862954, 0.26130258, 0.27577711))]
    return torchvision.transforms.Compose(transform_list)


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--outdir",
        type=str,
        nargs="?",
        help="dir to write results to",
        default="outputs/txt2img-samples"
    )
    parser.add_argument(
        "--skip_grid",
        action='store_true',
        help="do not save a grid, only individual samples. Helpful when evaluating lots of samples",
    )
    parser.add_argument(
        "--skip_save",
        action='store_true',
        help="do not save individual samples. For speed measurements.",
    )
    parser.add_argument(
        "--gpu_id",
        type=int,
        default=0,
        help="which gpu to use",
    )
    parser.add_argument(
        "--ddim_steps",
        type=int,
        default=30,
        help="number of ddim sampling steps",
    )
    parser.add_argument(
        "--plms",
        action='store_true',
        help="use plms sampling",
    )
    parser.add_argument(
        "--fixed_code",
        action='store_true',
        help="if enabled, uses the same starting code across samples ",
    )
    parser.add_argument(
        "--ddim_eta",
        type=float,
        default=0.0,
        help="ddim eta (eta=0.0 corresponds to deterministic sampling",
    )
    parser.add_argument(
        "--n_iter",
        type=int,
        default=2,
        help="sample this often",
    )
    parser.add_argument(
        "--H",
        type=int,
        default=512,
        help="image height, in pixel space",
    )
    parser.add_argument(
        "--W",
        type=int,
        default=512,
        help="image width, in pixel space",
    )
    parser.add_argument(
        "--n_imgs",
        type=int,
        default=100,
        help="image width, in pixel space",
    )
    parser.add_argument(
        "--C",
        type=int,
        default=4,
        help="latent channels",
    )
    parser.add_argument(
        "--f",
        type=int,
        default=8,
        help="downsampling factor",
    )
    parser.add_argument(
        "--n_samples",
        type=int,
        default=1,
        help="how many samples to produce for each given reference image. A.k.a. batch size",
    )
    parser.add_argument(
        "--n_rows",
        type=int,
        default=0,
        help="rows in the grid (default: n_samples)",
    )
    parser.add_argument(
        "--scale",
        type=float,
        default=1,
        help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
    )
    parser.add_argument(
        "--config",
        type=str,
        default="",
        help="path to config which constructs model",
    )
    parser.add_argument(
        "--ckpt",
        type=str,
        default="",
        help="path to checkpoint of model",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="the seed (for reproducible sampling)",
    )
    parser.add_argument(
        "--precision",
        type=str,
        help="evaluate at this precision",
        choices=["full", "autocast"],
        default="autocast"
    )
    parser.add_argument(
        "--unpaired",
        action='store_true',
        help="if enabled, uses the same starting code across samples "
    )
    parser.add_argument(
        "--dataroot",
        type=str,
        help="path to dataroot of the dataset",
        default=""
    )

    opt = parser.parse_args()

    seed_everything(opt.seed)

    device = torch.device("cuda:{}".format(opt.gpu_id)) if torch.cuda.is_available() else torch.device("cpu")
    torch.cuda.set_device(device)

    config = OmegaConf.load(f"{opt.config}")
    version = opt.config.split('/')[-1].split('.')[0]
    model = load_model_from_config(config, f"{opt.ckpt}")

    # model = model.to(device)
    dataset = CPDataset(opt.dataroot, opt.H, mode='test', unpaired=opt.unpaired)
    loader = DataLoader(dataset, batch_size=opt.n_samples, shuffle=False, num_workers=4, pin_memory=True)
    if opt.plms:
        sampler = PLMSSampler(model)
    else:
        sampler = DDIMSampler(model)

    os.makedirs(opt.outdir, exist_ok=True)
    outpath = opt.outdir

    result_path = os.path.join(outpath, "upper_body")
    os.makedirs(result_path, exist_ok=True)

    start_code = None
    if opt.fixed_code:
        start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)

    iterator = tqdm(loader, desc='Test Dataset', total=len(loader))
    precision_scope = autocast if opt.precision == "autocast" else nullcontext
    with torch.no_grad():
        with precision_scope("cuda"):
            with model.ema_scope():
                for data in iterator:
                    mask_tensor = data['inpaint_mask']
                    inpaint_image = data['inpaint_image']
                    ref_tensor_f = data['ref_imgs_f']
                    ref_tensor_b = data['ref_imgs_b']
                    skeleton_cf = data['skeleton_cf']
                    skeleton_cb = data['skeleton_cb']
                    skeleton_p = data['skeleton_p']
                    order = data['order']
                    feat_tensor = data['warp_feat']
                    image_tensor = data['GT']

                    controlnet_cond_f = data['controlnet_cond_f']
                    controlnet_cond_b = data['controlnet_cond_b']

                    ref_tensor = ref_tensor_f
                    for i in range(len(order)):
                        if order[i] == "1" or order[i] == "2":
                            continue
                        elif order[i] == "3":
                            ref_tensor[i] = ref_tensor_b[i]
                        else:
                            raise ValueError("Invalid order")

                    # filename = data['file_name']

                    test_model_kwargs = {}
                    test_model_kwargs['inpaint_mask'] = mask_tensor.to(device)
                    test_model_kwargs['inpaint_image'] = inpaint_image.to(device)
                    feat_tensor = feat_tensor.to(device)
                    ref_tensor = ref_tensor.to(device)

                    controlnet_cond_f = controlnet_cond_f.to(device)
                    controlnet_cond_b = controlnet_cond_b.to(device)
                    skeleton_cf = skeleton_cf.to(device)
                    skeleton_cb = skeleton_cb.to(device)
                    skeleton_p = skeleton_p.to(device)

                    uc = None
                    if opt.scale != 1.0:
                        uc = model.learnable_vector
                        uc = uc.repeat(ref_tensor.size(0), 1, 1)
                    c = model.get_learned_conditioning(ref_tensor.to(torch.float16))
                    c = model.proj_out(c)

                    # z_gt = model.encode_first_stage(image_tensor.to(device))
                    # z_gt = model.get_first_stage_encoding(z_gt).detach()

                    z_inpaint = model.encode_first_stage(test_model_kwargs['inpaint_image'])
                    z_inpaint = model.get_first_stage_encoding(z_inpaint).detach()
                    test_model_kwargs['inpaint_image'] = z_inpaint
                    test_model_kwargs['inpaint_mask'] = Resize([z_inpaint.shape[-2], z_inpaint.shape[-1]])(
                        test_model_kwargs['inpaint_mask'])

                    warp_feat = model.encode_first_stage(feat_tensor)
                    warp_feat = model.get_first_stage_encoding(warp_feat).detach()

                    ts = torch.full((1,), 999, device=device, dtype=torch.long)
                    start_code = model.q_sample(warp_feat, ts)

                    # local_controlnet
                    ehs_cf = model.pose_model(skeleton_cf)
                    ehs_cb = model.pose_model(skeleton_cb)
                    ehs_p = model.pose_model(skeleton_p)
                    ehs_text = torch.zeros((c.shape[0], 1, 768)).to("cuda")
                    # controlnet_cond = torch.cat((controlnet_cond_f, controlnet_cond_b, ehs_cf, ehs_cb, ehs_p), dim=1)
                    x_noisy = torch.cat(
                        (start_code, test_model_kwargs['inpaint_image'], test_model_kwargs['inpaint_mask']), dim=1)

                    down_samples_f, mid_samples_f = model.local_controlnet(x_noisy, ts,
                                                                           encoder_hidden_states=ehs_text.to("cuda"), controlnet_cond=controlnet_cond_f, ehs_c=ehs_cf, ehs_p=ehs_p)
                    down_samples_b, mid_samples_b = model.local_controlnet(x_noisy, ts,
                                                                           encoder_hidden_states=ehs_text.to("cuda"), controlnet_cond=controlnet_cond_b, ehs_c=ehs_cb, ehs_p=ehs_p)

                    # print(torch.max(down_samples_f[0]))
                    # print(torch.min(down_samples_f[0]))

                    # normalized_tensor = (down_samples_f[0] + 1) / 2

                    # # 将张量值范围从[0,1]转换到[0,255]
                    # scaled_tensor = normalized_tensor * 255

                    # # 将张量转换为NumPy数组
                    # numpy_array = scaled_tensor.squeeze().cpu().numpy().astype(np.uint8)

                    # # 将NumPy数组转换为PIL图像
                    # image = Image.fromarray(numpy_array)

                    # # 保存图像
                    # image.save("down_samples_f.jpg")

                    # normalized_tensor = (down_samples_b[0] + 1) / 2

                    # # 将张量值范围从[0,1]转换到[0,255]
                    # scaled_tensor = normalized_tensor * 255

                    # # 将张量转换为NumPy数组
                    # numpy_array = scaled_tensor.squeeze().cpu().numpy().astype(np.uint8)

                    # # 将NumPy数组转换为PIL图像
                    # image = Image.fromarray(numpy_array)

                    # # 保存图像
                    # image.save("down_samples_b.jpg")                    

                    mid_samples = mid_samples_f + mid_samples_b
                    down_samples = ()
                    for ds in range(len(down_samples_f)):
                        tmp = torch.cat((down_samples_f[ds], down_samples_b[ds]), dim=1)
                        down_samples = down_samples + (tmp,)

                    shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
                    samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
                                                     conditioning=c,
                                                     batch_size=opt.n_samples,
                                                     shape=shape,
                                                     verbose=False,
                                                     unconditional_guidance_scale=opt.scale,
                                                     unconditional_conditioning=uc,
                                                     eta=opt.ddim_eta,
                                                     x_T=start_code,
                                                     down_samples=down_samples,
                                                     test_model_kwargs=test_model_kwargs)

                    x_samples_ddim = model.decode_first_stage(samples_ddim)
                    x_sample_result = x_samples_ddim
                    x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
                    x_samples_ddim = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy()

                    x_checked_image = x_samples_ddim
                    x_checked_image_torch = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)
                    x_source = torch.clamp((image_tensor + 1.0) / 2.0, min=0.0, max=1.0)
                    x_result = x_checked_image_torch * (1 - mask_tensor) + mask_tensor * x_source
                    # x_result = x_checked_image_torch

                    resize = transforms.Resize((opt.H, int(opt.H / 256 * 192)))

                    if not opt.skip_save:

                        def un_norm(x):
                            return (x + 1.0) / 2.0

                        for i, x_sample in enumerate(x_result):
                            filename = data['file_name'][i]
                            # filename = data['file_name']
                            save_x = resize(x_sample)
                            save_x = 255. * rearrange(save_x.cpu().numpy(), 'c h w -> h w c')
                            img = Image.fromarray(save_x.astype(np.uint8))
                            img.save(os.path.join(result_path, filename[:-4] + ".png"))

    print(f"Your samples are ready and waiting for you here: \n{outpath} \n"
          f" \nEnjoy.")


if __name__ == "__main__":
    main()