train_generator.py

import torch
import torch.nn as nn
from torch.nn import functional as F

import argparse
import os
import time
from cp_dataset import CPDataset, CPDataLoader
from cp_dataset_test import CPDatasetTest
from networks import ConditionGenerator, VGGLoss, load_checkpoint, save_checkpoint, make_grid, make_grid_3d
from network_generator import SPADEGenerator, MultiscaleDiscriminator, GANLoss, Projected_GANs_Loss, set_requires_grad

from sync_batchnorm import DataParallelWithCallback
from utils import create_network
import sys
from tqdm import tqdm

import numpy as np
from torch.utils.data import Subset
from torchvision.transforms import transforms
import eval_models as models
import torchgeometry as tgm

from pg_modules.discriminator import ProjectedDiscriminator
import cv2

def remove_overlap(seg_out, warped_cm):
    assert len(warped_cm.shape) == 4
    warped_cm = warped_cm - (torch.cat([seg_out[:, 1:3, :, :], seg_out[:, 5:, :, :]], dim=1)).sum(dim=1, keepdim=True) * warped_cm
    return warped_cm

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

    parser.add_argument('--name', type=str, required=True)
    parser.add_argument('--gpu_ids', type=str, default='0')
    parser.add_argument('-j', '--workers', type=int, default=4)
    parser.add_argument('-b', '--batch_size', type=int, default=8)
    parser.add_argument('--fp16', action='store_true', help='use amp')

    parser.add_argument("--dataroot", default="./data/")
    parser.add_argument("--datamode", default="train")
    parser.add_argument("--data_list", default="train_pairs.txt")
    parser.add_argument("--fine_width", type=int, default=768)
    parser.add_argument("--fine_height", type=int, default=1024)
    parser.add_argument("--radius", type=int, default=20)
    parser.add_argument("--grid_size", type=int, default=5)

    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints', help='save checkpoint infos')
    parser.add_argument('--tocg_checkpoint', type=str, help='condition generator checkpoint')
    parser.add_argument('--gen_checkpoint', type=str, default='', help='gen checkpoint')
    parser.add_argument('--dis_checkpoint', type=str, default='', help='dis checkpoint')

    parser.add_argument("--display_count", type=int, default=100)
    parser.add_argument("--save_count", type=int, default=1000)
    parser.add_argument("--load_step", type=int, default=0)
    parser.add_argument("--keep_step", type=int, default=100000)
    parser.add_argument("--decay_step", type=int, default=100000)
    parser.add_argument("--shuffle", action='store_true', help='shuffle input data')
    parser.add_argument('--resume', action='store_true', help='resume training from the last checkpoint')
    
    # test
    parser.add_argument("--lpips_count", type=int, default=1000)
    parser.add_argument("--test_datasetting", default="paired")
    parser.add_argument("--test_dataroot", default="./data/")
    parser.add_argument("--test_data_list", default="test_pairs.txt")

    # Hyper-parameters
    parser.add_argument('--G_lr', type=float, default=0.0001, help='initial learning rate for adam')
    parser.add_argument('--D_lr', type=float, default=0.0004, help='initial learning rate for adam')

    # SEAN-related hyper-parameters
    parser.add_argument('--GMM_const', type=float, default=None, help='constraint for GMM module')
    parser.add_argument('--semantic_nc', type=int, default=13, help='# of input label classes without unknown class')
    parser.add_argument('--gen_semantic_nc', type=int, default=7, help='# of input label classes without unknown class')
    parser.add_argument('--norm_G', type=str, default='spectralaliasinstance', help='instance normalization or batch normalization')
    parser.add_argument('--norm_D', type=str, default='spectralinstance', help='instance normalization or batch normalization')
    parser.add_argument('--ngf', type=int, default=64, help='# of gen filters in first conv layer')
    parser.add_argument('--ndf', type=int, default=64, help='# of discrim filters in first conv layer')
    parser.add_argument('--num_upsampling_layers', choices=['normal', 'more', 'most'], default='most',
                    help='If \'more\', add upsampling layer between the two middle resnet blocks. '
                            'If \'most\', also add one more (upsampling + resnet) layer at the end of the generator.')
    parser.add_argument('--init_type', type=str, default='xavier', help='network initialization [normal|xavier|kaiming|orthogonal]')
    parser.add_argument('--init_variance', type=float, default=0.02, help='variance of the initialization distribution')

    parser.add_argument('--no_ganFeat_loss', action='store_true', help='if specified, do *not* use discriminator feature matching loss')
    parser.add_argument('--lambda_l1', type=float, default=1.0, help='weight for image-level l1 loss')
    parser.add_argument('--lambda_feat', type=float, default=10.0, help='weight for feature matching loss')
    parser.add_argument('--lambda_vgg', type=float, default=10.0, help='weight for vgg loss')
    
    # D
    parser.add_argument('--n_layers_D', type=int, default=3, help='# layers in each discriminator')
    parser.add_argument('--netD_subarch', type=str, default='n_layer', help='architecture of each discriminator')
    parser.add_argument('--num_D', type=int, default=2, help='number of discriminators to be used in multiscale')
    
    # G & D arch-related
    parser.add_argument("--composition_mask", action='store_true', help='shuffle input data')

    # Training
    parser.add_argument('--occlusion', action='store_true')
    # tocg
    # network
    parser.add_argument('--cond_G_ngf', type=int, default=96)
    parser.add_argument("--cond_G_input_width", type=int, default=192)
    parser.add_argument("--cond_G_input_height", type=int, default=256)
    parser.add_argument('--cond_G_num_layers', type=int, default=5)
    parser.add_argument("--warp_feature", choices=['encoder', 'T1'], default="T1")
    parser.add_argument("--out_layer", choices=['relu', 'conv'], default="relu")

    # New arguments for selective layer freezing and last layer control
    parser.add_argument('--freeze_tocg_layers', type=int, default=0, help='number of layers to freeze in tocg from the start')
    parser.add_argument('--freeze_gen_layers', type=int, default=0, help='number of layers to freeze in generator from the start')
    parser.add_argument('--last_layer_mode', type=str, default='train', choices=['train', 'half', 'freeze'],
                        help='Mode for the last layer: train (full training), half (half parameters frozen), freeze (fully frozen)')

    opt = parser.parse_args()

    # set gpu ids
    str_ids = opt.gpu_ids.split(',')
    opt.gpu_ids = []
    for str_id in str_ids:
        id = int(str_id)
        if id >= 0:
            opt.gpu_ids.append(id)
    if len(opt.gpu_ids) > 0:
        torch.cuda.set_device(opt.gpu_ids[0])

    assert len(opt.gpu_ids) == 0 or opt.batch_size % len(opt.gpu_ids) == 0, \
        "Batch size %d is wrong. It must be a multiple of # GPUs %d." \
        % (opt.batch_size, len(opt.gpu_ids))

    return opt

def apply_layer_freezing(model, num_layers_to_freeze, last_layer_mode):
    """Apply selective layer freezing and handle the last layer based on mode."""
    children = list(model.named_children())
    total_layers = len(children)
    
    # Freeze specified layers from the start
    for i, (name, module) in enumerate(children):
        if i < num_layers_to_freeze:
            for param in module.parameters():
                param.requires_grad = False
    
    # Handle the last layer based on mode
    if total_layers > 0 and last_layer_mode != 'train':
        last_name, last_module = children[-1]
        if last_layer_mode == 'freeze':
            for param in last_module.parameters():
                param.requires_grad = False
        elif last_layer_mode == 'half':
            # Freeze half of the parameters in the last layer
            params = list(last_module.parameters())
            half_idx = len(params) // 2
            for param in params[:half_idx]:
                param.requires_grad = False
            for param in params[half_idx:]:
                param.requires_grad = True

def train(opt, train_loader, test_loader, tocg, generator, discriminator, model):
    """
        Train Generator and Condition Generator
    """
    # Model
    tocg.cuda()
    tocg.train()  # Enable training for tocg
    generator.train()
    discriminator.train()
    if not opt.composition_mask:
        discriminator.feature_network.requires_grad_(False)
    discriminator.cuda()
    model.eval()

    # Apply layer freezing
    apply_layer_freezing(tocg, opt.freeze_tocg_layers, opt.last_layer_mode)
    apply_layer_freezing(generator, opt.freeze_gen_layers, opt.last_layer_mode)

    # criterion
    criterionGAN = None
    if opt.fp16:
        if opt.composition_mask:
            criterionGAN = GANLoss('hinge', tensor=torch.cuda.HalfTensor)
        else:
            criterionGAN = Projected_GANs_Loss(tensor=torch.cuda.HalfTensor)
    else:
        if opt.composition_mask:
            criterionGAN = GANLoss('hinge', tensor=torch.cuda.FloatTensor)
        else:
            criterionGAN = Projected_GANs_Loss(tensor=torch.cuda.FloatTensor)
    
    criterionL1 = nn.L1Loss()
    criterionFeat = nn.L1Loss()
    criterionVGG = VGGLoss()

    # optimizer
    optimizer_gen = torch.optim.Adam(
        list(generator.parameters()) + list(tocg.parameters()),  # Include tocg parameters
        lr=opt.G_lr, betas=(0.0, 0.9)
    )
    scheduler_gen = torch.optim.lr_scheduler.LambdaLR(optimizer_gen, lr_lambda=lambda step: 1.0 -
            max(0, step * 1000 + opt.load_step - opt.keep_step) / float(opt.decay_step + 1))
    optimizer_dis = torch.optim.Adam(discriminator.parameters(), lr=opt.D_lr, betas=(0.0, 0.9))
    scheduler_dis = torch.optim.lr_scheduler.LambdaLR(optimizer_dis, lr_lambda=lambda step: 1.0 -
            max(0, step * 1000 + opt.load_step - opt.keep_step) / float(opt.decay_step + 1))

    if opt.fp16:
        from apex import amp
        [tocg, generator, discriminator], [optimizer_gen, optimizer_dis] = amp.initialize(
            [tocg, generator, discriminator], [optimizer_gen, optimizer_dis], opt_level='O1', num_losses=2)
        
    if len(opt.gpu_ids) > 0:
        tocg = DataParallelWithCallback(tocg, device_ids=opt.gpu_ids)
        generator = DataParallelWithCallback(generator, device_ids=opt.gpu_ids)
        discriminator = DataParallelWithCallback(discriminator, device_ids=opt.gpu_ids)
        criterionGAN = DataParallelWithCallback(criterionGAN, device_ids=opt.gpu_ids)
        criterionFeat = DataParallelWithCallback(criterionFeat, device_ids=opt.gpu_ids)
        criterionVGG = DataParallelWithCallback(criterionVGG, device_ids=opt.gpu_ids)
        criterionL1 = DataParallelWithCallback(criterionL1, device_ids=opt.gpu_ids)
        
    upsample = torch.nn.Upsample(scale_factor=4, mode='bilinear')
    gauss = tgm.image.GaussianBlur((15, 15), (3, 3))
    gauss = gauss.cuda()

    checkpoint_path = os.path.join(opt.checkpoint_dir, opt.name, 'checkpoint.pth')
    if opt.resume:
        if os.path.exists(checkpoint_path):
            print(f"Resuming from checkpoint: {checkpoint_path}")
            checkpoint = torch.load(checkpoint_path)
            opt.load_step = checkpoint['step']
            generator.load_state_dict(checkpoint['generator_state_dict'])
            discriminator.load_state_dict(checkpoint['discriminator_state_dict'])
            tocg.load_state_dict(checkpoint['tocg_state_dict'])  # Load tocg state
            optimizer_gen.load_state_dict(checkpoint['optimizer_gen_state_dict'])
            optimizer_dis.load_state_dict(checkpoint['optimizer_dis_state_dict'])
            scheduler_gen.load_state_dict(checkpoint['scheduler_gen_state_dict'])
            scheduler_dis.load_state_dict(checkpoint['scheduler_dis_state_dict'])
        else:
            print(f"Checkpoint not found at {checkpoint_path}, starting from scratch")

    for step in tqdm(range(opt.load_step, opt.keep_step + opt.decay_step)):
        iter_start_time = time.time()
        inputs = train_loader.next_batch()

        # input
        agnostic = inputs['agnostic'].cuda()
        parse_GT = inputs['parse'].cuda()
        pose = inputs['densepose'].cuda()
        parse_cloth = inputs['parse_cloth'].cuda()
        parse_agnostic = inputs['parse_agnostic'].cuda()
        pcm = inputs['pcm'].cuda()
        cm = inputs['cloth_mask']['paired'].cuda()
        c_paired = inputs['cloth']['paired'].cuda()
        
        # target
        im = inputs['image'].cuda()

        # Warping Cloth (tocg is now trainable)
        pre_clothes_mask_down = F.interpolate(cm, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
        input_parse_agnostic_down = F.interpolate(parse_agnostic, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
        clothes_down = F.interpolate(c_paired, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
        densepose_down = F.interpolate(pose, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
        
        input1 = torch.cat([clothes_down, pre_clothes_mask_down], 1)
        input2 = torch.cat([input_parse_agnostic_down, densepose_down], 1)

        flow_list_taco, fake_segmap, warped_cloth_paired_taco, warped_clothmask_paired_taco, flow_list_tvob, warped_cloth_paired_tvob, warped_clothmask_paired_tvob = tocg(input1, input2)
        
        warped_clothmask_paired_taco_onehot = torch.FloatTensor((warped_clothmask_paired_taco.detach().cpu().numpy() > 0.5).astype(float)).cuda()
        
        cloth_mask = torch.ones_like(fake_segmap)
        cloth_mask[:,3:4, :, :] = warped_clothmask_paired_taco
        fake_segmap = fake_segmap * cloth_mask
                
        N, _, iH, iW = c_paired.shape
        N, flow_iH, flow_iW, _ = flow_list_tvob[-1].shape

        flow_tvob = F.interpolate(flow_list_tvob[-1].permute(0, 3, 1, 2), size=(iH, iW), mode='bilinear').permute(0, 2, 3, 1)
        flow_tvob_norm = torch.cat([flow_tvob[:, :, :, 0:1] / ((flow_iW - 1.0) / 2.0), flow_tvob[:, :, :, 1:2] / ((flow_iH - 1.0) / 2.0)], 3)

        grid = make_grid(N, iH, iW)
        grid_3d = make_grid_3d(N, iH, iW)

        warped_grid_tvob = grid + flow_tvob_norm
        warped_cloth_tvob = F.grid_sample(c_paired, warped_grid_tvob, padding_mode='border')
        warped_clothmask_tvob = F.grid_sample(cm, warped_grid_tvob, padding_mode='border')

        flow_taco = F.interpolate(flow_list_taco[-1].permute(0, 4, 1, 2, 3), size=(2,iH,iW), mode='trilinear').permute(0, 2, 3, 4, 1)
        flow_taco_norm = torch.cat([flow_taco[:, :, :, :, 0:1] / ((flow_iW - 1.0) / 2.0), flow_taco[:, :, :, :, 1:2] / ((flow_iH - 1.0) / 2.0), flow_taco[:, :, :, :, 2:3]], 4)
        warped_cloth_tvob = warped_cloth_tvob.unsqueeze(2)
        warped_cloth_paired_taco = F.grid_sample(torch.cat((warped_cloth_tvob, torch.zeros_like(warped_cloth_tvob).cuda()), dim=2), flow_taco_norm + grid_3d, padding_mode='border')
        warped_cloth_paired_taco = warped_cloth_paired_taco[:,:,0,:,:]

        warped_clothmask_tvob = warped_clothmask_tvob.unsqueeze(2)
        warped_clothmask_taco = F.grid_sample(torch.cat((warped_clothmask_tvob, torch.zeros_like(warped_clothmask_tvob).cuda()), dim=2), flow_taco_norm + grid_3d, padding_mode='border')
        warped_clothmask_taco = warped_clothmask_taco[:,:,0,:,:]

        fake_parse_gauss = gauss(F.interpolate(fake_segmap, size=(iH, iW), mode='bilinear'))
        fake_parse = fake_parse_gauss.argmax(dim=1)[:, None]

        if opt.occlusion:
            warped_clothmask_taco = remove_overlap(F.softmax(fake_parse_gauss, dim=1), warped_clothmask_taco)
            warped_cloth_paired_taco = warped_cloth_paired_taco * warped_clothmask_taco + torch.ones_like(warped_cloth_paired_taco) * (1-warped_clothmask_taco)
            warped_cloth_paired_taco = warped_cloth_paired_taco.detach()
            
        old_parse = torch.FloatTensor(fake_parse.size(0), 13, opt.fine_height, opt.fine_width).zero_().cuda()
        old_parse.scatter_(1, fake_parse, 1.0)

        labels = {
            0:  ['background',  [0]],
            1:  ['paste',       [2, 4, 7, 8, 9, 10, 11]],
            2:  ['upper',       [3]],
            3:  ['hair',        [1]],
            4:  ['left_arm',    [5]],
            5:  ['right_arm',   [6]],
            6:  ['noise',       [12]]
        }
        parse = torch.FloatTensor(fake_parse.size(0), 7, opt.fine_height, opt.fine_width).zero_().cuda()
        for i in range(len(labels)):
            for label in labels[i][1]:
                parse[:, i] += old_parse[:, label]
                
        parse = parse.detach()

        # Train the generator and tocg
        G_losses = {}
        if opt.composition_mask:
            output_paired_rendered, output_paired_comp = generator(torch.cat((agnostic, pose, warped_cloth_paired_taco), dim=1), parse)
            output_paired_comp1 = output_paired_comp * warped_clothmask_taco
            output_paired_comp = parse[:,2:3,:,:] * output_paired_comp1
            output_paired = warped_cloth_paired_taco * output_paired_comp + output_paired_rendered * (1 - output_paired_comp)

            fake_concat = torch.cat((parse, output_paired_rendered), dim=1)
            real_concat = torch.cat((parse, im), dim=1)
            pred = discriminator(torch.cat((fake_concat, real_concat), dim=0))

            pred_fake = []
            pred_real = []
            for p in pred:
                pred_fake.append([tensor[:tensor.size(0) // 2] for tensor in p])
                pred_real.append([tensor[tensor.size(0) // 2:] for tensor in p])
            
            G_losses['GAN'] = criterionGAN(pred_fake, True, for_discriminator=False)

            num_D = len(pred_fake)
            GAN_Feat_loss = torch.cuda.FloatTensor(len(opt.gpu_ids)).zero_()
            for i in range(num_D):
                num_intermediate_outputs = len(pred_fake[i]) - 1
                for j in range(num_intermediate_outputs):
                    unweighted_loss = criterionFeat(pred_fake[i][j], pred_real[i][j].detach())
                    GAN_Feat_loss += unweighted_loss * opt.lambda_feat / num_D
            G_losses['GAN_Feat'] = GAN_Feat_loss

            G_losses['VGG'] = criterionVGG(output_paired, im) * opt.lambda_vgg + criterionVGG(output_paired_rendered, im) * opt.lambda_vgg
            G_losses['L1'] = criterionL1(output_paired_rendered, im) * opt.lambda_l1 + criterionL1(output_paired, im) * opt.lambda_l1
            G_losses['Composition_Mask'] = torch.mean(torch.abs(1 - output_paired_comp))

            loss_gen = sum(G_losses.values()).mean()

        else:
            set_requires_grad(discriminator, False)
            output_paired = generator(torch.cat((agnostic, pose, warped_cloth_paired_taco), dim=1), parse)

            pred_fake, feats_fake = discriminator(output_paired)
            pred_real, feats_real = discriminator(im)

            G_losses['GAN'] = criterionGAN(pred_fake, True, for_discriminator=False) * 0.5

            num_D = len(feats_fake)
            GAN_Feat_loss = torch.cuda.FloatTensor(len(opt.gpu_ids)).zero_()
            for i in range(num_D):
                num_intermediate_outputs = len(feats_fake[i])
                for j in range(num_intermediate_outputs):
                    unweighted_loss = criterionFeat(feats_fake[i][j], feats_real[i][j].detach())
                    GAN_Feat_loss += unweighted_loss * opt.lambda_feat / num_D
            G_losses['GAN_Feat'] = GAN_Feat_loss

            G_losses['VGG'] = criterionVGG(output_paired, im) * opt.lambda_vgg
            G_losses['L1'] = criterionL1(output_paired, im) * opt.lambda_l1

            loss_gen = sum(G_losses.values()).mean()

        optimizer_gen.zero_grad()
        if opt.fp16:
            with amp.scale_loss(loss_gen, optimizer_gen, loss_id=0) as loss_gen_scaled:
                loss_gen_scaled.backward()
        else:
            loss_gen.backward()
        optimizer_gen.step()

        # Train the discriminator
        D_losses = {}
        if opt.composition_mask:
            with torch.no_grad():            
                output_paired_rendered, output_comp = generator(torch.cat((agnostic, pose, warped_cloth_paired_taco), dim=1), parse)
                output_comp1 = output_comp * warped_clothmask_taco
                output_comp = parse[:,2:3,:,:] * output_comp1
                output = warped_cloth_paired_taco * output_comp + output_paired_rendered * (1 - output_comp)
                output_comp = output_comp.detach()
                output = output.detach()
                output_comp.requires_grad_()
                output.requires_grad_()

            fake_concat = torch.cat((parse, output_paired_rendered), dim=1)
            real_concat = torch.cat((parse, im), dim=1)
            pred = discriminator(torch.cat((fake_concat, real_concat), dim=0))

            pred_fake = []
            pred_real = []
            for p in pred:
                pred_fake.append([tensor[:tensor.size(0) // 2] for tensor in p])
                pred_real.append([tensor[tensor.size(0) // 2:] for tensor in p])
            
            D_losses['D_Fake'] = criterionGAN(pred_fake, False, for_discriminator=True)
            D_losses['D_Real'] = criterionGAN(pred_real, True, for_discriminator=True)
            
            loss_dis = sum(D_losses.values()).mean()

        else:
            set_requires_grad(discriminator, True)
            discriminator.module.feature_network.requires_grad_(False)

            with torch.no_grad():
                output = generator(torch.cat((agnostic, pose, warped_cloth_paired_taco), dim=1), parse)
                output = output.detach()
                output.requires_grad_()

            pred_fake, _ = discriminator(output)
            pred_real, _ = discriminator(im)

            D_losses['D_Fake'] = criterionGAN(pred_fake, False, for_discriminator=True)
            D_losses['D_Real'] = criterionGAN(pred_real, True, for_discriminator=True)

            loss_dis = sum(D_losses.values()).mean()

        optimizer_dis.zero_grad()
        if opt.fp16:
            with amp.scale_loss(loss_dis, optimizer_dis, loss_id=1) as loss_dis_scaled:
                loss_dis_scaled.backward()
        else:
            loss_dis.backward()
        optimizer_dis.step()
        
        if not opt.composition_mask:
            set_requires_grad(discriminator, False)

        if (step+1) % 100 == 0:
            a_0 = im.cuda()[0]
            b_0 = output.cuda()[0]
            c_0 = warped_cloth_paired_taco.cuda()[0]
            combine = torch.cat((a_0, b_0, c_0), dim=2)
            cv_img=(combine.permute(1,2,0).detach().cpu().numpy()+1)/2
            rgb=(cv_img*255).astype(np.uint8)
            bgr=cv2.cvtColor(rgb,cv2.COLOR_RGB2BGR)
            cv2.imwrite('sample_fs_toig/'+str(step)+'.jpg',bgr)

        # Evaluate the generator
        if (step + 1) % opt.lpips_count == 0:
            generator.eval()
            tocg.eval()
            T2 = transforms.Compose([transforms.Resize((128, 128))])
            lpips_list = []
            avg_distance = 0.0
            
            with torch.no_grad():
                print("LPIPS")
                for i in tqdm(range(500)):
                    inputs = test_loader.next_batch()
                    agnostic = inputs['agnostic'].cuda()
                    parse_GT = inputs['parse'].cuda()
                    pose = inputs['densepose'].cuda()
                    parse_cloth = inputs['parse_cloth'].cuda()
                    parse_agnostic = inputs['parse_agnostic'].cuda()
                    pcm = inputs['pcm'].cuda()
                    cm = inputs['cloth_mask']['paired'].cuda()
                    c_paired = inputs['cloth']['paired'].cuda()
                    im = inputs['image'].cuda()
                                
                    pre_clothes_mask_down = F.interpolate(cm, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
                    input_parse_agnostic_down = F.interpolate(parse_agnostic, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='nearest')
                    clothes_down = F.interpolate(c_paired, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
                    densepose_down = F.interpolate(pose, size=(opt.cond_G_input_height, opt.cond_G_input_width), mode='bilinear')
                    
                    input1 = torch.cat([clothes_down, pre_clothes_mask_down], 1)
                    input2 = torch.cat([input_parse_agnostic_down, densepose_down], 1)

                    flow_list_taco, fake_segmap, warped_cloth_paired_taco, warped_clothmask_paired_taco, flow_list_tvob, warped_cloth_paired_tvob, warped_clothmask_paired_tvob = tocg(input1, input2)
                    
                    warped_clothmask_paired_taco_onehot = torch.FloatTensor((warped_clothmask_paired_taco.detach().cpu().numpy() > 0.5).astype(float)).cuda()
                    
                    cloth_mask = torch.ones_like(fake_segmap)
                    cloth_mask[:,3:4, :, :] = warped_clothmask_paired_taco
                    fake_segmap = fake_segmap * cloth_mask
                            
                    N, _, iH, iW = c_paired.shape
                    N, flow_iH, flow_iW, _ = flow_list_tvob[-1].shape

                    flow_tvob = F.interpolate(flow_list_tvob[-1].permute(0, 3, 1, 2), size=(iH, iW), mode='bilinear').permute(0, 2, 3, 1)
                    flow_tvob_norm = torch.cat([flow_tvob[:, :, :, 0:1] / ((flow_iW - 1.0) / 2.0), flow_tvob[:, :, :, 1:2] / ((flow_iH - 1.0) / 2.0)], 3)

                    grid = make_grid(N, iH, iW)
                    grid_3d = make_grid_3d(N, iH, iW)

                    warped_grid_tvob = grid + flow_tvob_norm
                    warped_cloth_tvob = F.grid_sample(c_paired, warped_grid_tvob, padding_mode='border')
                    warped_clothmask_tvob = F.grid_sample(cm, warped_grid_tvob, padding_mode='border')

                    flow_taco = F.interpolate(flow_list_taco[-1].permute(0, 4, 1, 2, 3), size=(2, iH, iW), mode='trilinear').permute(0, 2, 3, 4, 1)
                    flow_taco_norm = torch.cat([flow_taco[:, :, :, :, 0:1] / ((flow_iW - 1.0) / 2.0), flow_taco[:, :, :, :, 1:2] / ((flow_iH - 1.0) / 2.0), flow_taco[:, :, :, :, 2:3]], 4)
                    warped_cloth_tvob = warped_cloth_tvob.unsqueeze(2)
                    warped_cloth_paired_taco = F.grid_sample(torch.cat((warped_cloth_tvob, torch.zeros_like(warped_cloth_tvob).cuda()), dim=2), flow_taco_norm + grid_3d, padding_mode='border')
                    warped_cloth_paired_taco = warped_cloth_paired_taco[:,:,0,:,:]

                    warped_clothmask_tvob = warped_clothmask_tvob.unsqueeze(2)
                    warped_clothmask_taco = F.grid_sample(torch.cat((warped_clothmask_tvob, torch.zeros_like(warped_clothmask_tvob).cuda()), dim=2), flow_taco_norm + grid_3d, padding_mode='border')
                    warped_clothmask_taco = warped_clothmask_taco[:,:,0,:,:]

                    fake_parse_gauss = gauss(F.interpolate(fake_segmap, size=(iH, iW), mode='bilinear'))
                    fake_parse = fake_parse_gauss.argmax(dim=1)[:, None]

                    if opt.occlusion:
                        warped_clothmask_taco = remove_overlap(F.softmax(fake_parse_gauss, dim=1), warped_clothmask_taco)
                        warped_cloth_paired_taco = warped_cloth_paired_taco * warped_clothmask_taco + torch.ones_like(warped_cloth_paired_taco) * (1-warped_clothmask_taco)
                        warped_cloth_paired_taco = warped_cloth_paired_taco.detach()
                        
                    old_parse = torch.FloatTensor(fake_parse.size(0), 13, opt.fine_height, opt.fine_width).zero_().cuda()
                    old_parse.scatter_(1, fake_parse, 1.0)

                    labels = {
                        0:  ['background',  [0]],
                        1:  ['paste',       [2, 4, 7, 8, 9, 10, 11]],
                        2:  ['upper',       [3]],
                        3:  ['hair',        [1]],
                        4:  ['left_arm',    [5]],
                        5:  ['right_arm',   [6]],
                        6:  ['noise',       [12]]
                    }
                    parse = torch.FloatTensor(fake_parse.size(0), 7, opt.fine_height, opt.fine_width).zero_().cuda()
                    for i in range(len(labels)):
                        for label in labels[i][1]:
                            parse[:, i] += old_parse[:, label]
                            
                    parse = parse.detach()
                
                    if opt.composition_mask:
                        output_paired_rendered, output_paired_comp = generator(torch.cat((agnostic, pose, warped_cloth_paired_taco), dim=1), parse)
                        output_paired_comp1 = output_paired_comp * warped_clothmask_taco
                        output_paired_comp = parse[:,2:3,:,:] * output_paired_comp1
                        output_paired = warped_cloth_paired_taco * output_paired_comp + output_paired_rendered * (1 - output_paired_comp)
                    else:
                        output_paired = generator(torch.cat((agnostic, pose, warped_cloth_paired_taco), dim=1), parse)

                    avg_distance += model.forward(T2(im), T2(output_paired))
                    
            avg_distance = avg_distance / 500
            print(f"LPIPS: {avg_distance}")
            generator.train()
            tocg.train()

        if (step + 1) % opt.display_count == 0:
            t = time.time() - iter_start_time
            print("step: %8d, time: %.3f, G_loss: %.4f, G_adv_loss: %.4f, D_loss: %.4f, D_fake_loss: %.4f, D_real_loss: %.4f"
                  % (step + 1, t, loss_gen.item(), G_losses['GAN'].mean().item(), loss_dis.item(),
                     D_losses['D_Fake'].mean().item(), D_losses['D_Real'].mean().item()), flush=True)

        if (step + 1) % opt.save_count == 0:
            checkpoint = {
                'step': step + 1,
                'generator_state_dict': generator.state_dict(),
                'discriminator_state_dict': discriminator.state_dict(),
                'tocg_state_dict': tocg.state_dict(),  # Save tocg state
                'optimizer_gen_state_dict': optimizer_gen.state_dict(),
                'optimizer_dis_state_dict': optimizer_dis.state_dict(),
                'scheduler_gen_state_dict': scheduler_gen.state_dict(),
                'scheduler_dis_state_dict': scheduler_dis.state_dict(),
            }
            torch.save(checkpoint, checkpoint_path)

        if (step + 1) % 1000 == 0:
            scheduler_gen.step()
            scheduler_dis.step()

def main():
    opt = get_opt()
    print(opt)
    print("Start to train %s!" % opt.name)

    os.makedirs('sample_fs_toig', exist_ok=True)
    os.makedirs(os.path.join(opt.checkpoint_dir, opt.name), exist_ok=True)

    train_dataset = CPDataset(opt)
    train_loader = CPDataLoader(opt, train_dataset)
    
    opt.batch_size = 1
    opt.dataroot = opt.test_dataroot
    opt.datamode = 'test'
    opt.data_list = opt.test_data_list
    test_dataset = CPDatasetTest(opt)
    test_dataset = Subset(test_dataset, np.arange(500))
    test_loader = CPDataLoader(opt, test_dataset)
    
    input1_nc = 4
    input2_nc = opt.semantic_nc + 3
    tocg = ConditionGenerator(opt, input1_nc=input1_nc, input2_nc=input2_nc, output_nc=13, ngf=opt.cond_G_ngf, norm_layer=nn.BatchNorm2d, num_layers=opt.cond_G_num_layers)
    load_checkpoint(tocg, opt.tocg_checkpoint)

    generator = SPADEGenerator(opt, 3+3+3)
    generator.print_network()
    if len(opt.gpu_ids) > 0:
        assert(torch.cuda.is_available())
        generator.cuda()
    generator.init_weights(opt.init_type, opt.init_variance)

    discriminator = None
    if opt.composition_mask:
        discriminator = create_network(MultiscaleDiscriminator, opt)
    else:
        discriminator = ProjectedDiscriminator(interp224=False)

    model = models.PerceptualLoss(model='net-lin',net='alex',use_gpu=True)

    if opt.gen_checkpoint and os.path.exists(opt.gen_checkpoint):
        load_checkpoint(generator, opt.gen_checkpoint)
    if opt.dis_checkpoint and os.path.exists(opt.dis_checkpoint):
        load_checkpoint(discriminator, opt.dis_checkpoint)

    train(opt, train_loader, test_loader, tocg, generator, discriminator, model)

    save_checkpoint(generator, os.path.join(opt.checkpoint_dir, opt.name, 'gen_model_final.pth'))
    save_checkpoint(discriminator, os.path.join(opt.checkpoint_dir, opt.name, 'dis_model_final.pth'))
    save_checkpoint(tocg, os.path.join(opt.checkpoint_dir, opt.name, 'tocg_model_final.pth'))

    print("Finished training %s!" % opt.name)

if __name__ == "__main__":
    main()