Upload train_generator.py with huggingface_hub

train_generator.py

@@ -0,0 +1,635 @@
+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()