3_ControlNet/3_inference.py

# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
import argparse
import datetime
import numpy as np
import time
import torch
import torch.backends.cudnn as cudnn
import json
from pathlib import Path
from timm.utils import NativeScaler, get_state_dict, ModelEma
from torchvision.utils import make_grid

import importlib
import utils
import logging
import torch.nn as nn

import sys
from typing import Iterable, Optional
import utils
import logging
import torch.distributed as dist
import os
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
from pycocotools.coco import COCO
from PIL import Image, ImageDraw, ImageFont
from cldm.model import create_model, load_state_dict
from cldm.ddim_hacked import DDIMSampler
from tqdm import tqdm
import matplotlib.pyplot as plt
import torchvision


def get_args_parser():
    parser = argparse.ArgumentParser('DeiT training and evaluation script', add_help=False)
    parser.add_argument('--batch-size', default=64, type=int)
    parser.add_argument('--epochs', default=300, type=int)
    parser.add_argument('--bce-loss', action='store_true')
    parser.add_argument('--unscale-lr', action='store_true')

    # Feature extractor parameters
    parser.add_argument('--model', default='deit_base_patch16_224', type=str)
    parser.add_argument('--target_model', default='deit_base_patch16_224', type=str)
    parser.add_argument('--input-size', default=224, type=int, help='images input size')

    parser.add_argument('--drop', type=float, default=0.0, metavar='PCT',
                        help='Dropout rate (default: 0.)')
    parser.add_argument('--drop-path', type=float, default=0.1, metavar='PCT',
                        help='Drop path rate (default: 0.1)')

    parser.add_argument('--model-ema', action='store_true')
    parser.add_argument('--no-model-ema', action='store_false', dest='model_ema')
    parser.set_defaults(model_ema=True)
    parser.add_argument('--model-ema-decay', type=float, default=0.99996, help='')
    parser.add_argument('--model-ema-force-cpu', action='store_true', default=False, help='')

    # CLDM parameters
    parser.add_argument('--total_train_steps', default=300000, type=int)
    parser.add_argument('--resume_path', default='', help='resume from checkpoint for controlnet')
    parser.add_argument('--cldm_learning_rate', default=1e-4, type=float, help='learning rate for controlnet')
    parser.add_argument('--sd_locked', default=True, help='whether to lock the sd of controlnet')
    parser.add_argument('--only_mid_control', default=False, help='only control the middle layers of controlnet')
    parser.add_argument('--cldm_yaml', default='./models/cldm_v15.yaml', help='yaml file for controlnet')
    parser.add_argument('--exp_dir', default='./exp', help='experiment directory')
    parser.add_argument('--image_floder', default='./data', help='training image floder')
    parser.add_argument('--txt_path', default='./data', help='training txt path')
    parser.add_argument("--log_every", default=20, type=int, help="log every n steps")
    parser.add_argument("--ckpt_every", default=1000, type=int, help="save checkpoint every n steps")
    parser.add_argument("--image_every", default=1000, type=int, help="log image every n steps")
    parser.add_argument('--controlnet_dir', default='', type=str, help='resume from checkpoint for controlnet')
    parser.add_argument('--global_step', default=0, type=int, help='global step')
    parser.add_argument('--save_img_path', default='./data', help='save image path')

    # Optimizer parameters
    parser.add_argument('--opt', default='adamw', type=str, metavar='OPTIMIZER',
                        help='Optimizer (default: "adamw"')
    parser.add_argument('--opt-eps', default=1e-8, type=float, metavar='EPSILON',
                        help='Optimizer Epsilon (default: 1e-8)')
    parser.add_argument('--opt-betas', default=None, type=float, nargs='+', metavar='BETA',
                        help='Optimizer Betas (default: None, use opt default)')
    parser.add_argument('--clip-grad', type=float, default=None, metavar='NORM',
                        help='Clip gradient norm (default: None, no clipping)')
    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
                        help='SGD momentum (default: 0.9)')
    parser.add_argument('--weight-decay', type=float, default=0.05,
                        help='weight decay (default: 0.05)')
    # Learning rate schedule parameters
    parser.add_argument('--sched', default='cosine', type=str, metavar='SCHEDULER',
                        help='LR scheduler (default: "cosine"')
    parser.add_argument('--lr', type=float, default=4e-4, metavar='LR',
                        help='learning rate (default: 5e-4)')
    parser.add_argument('--lr-noise', type=float, nargs='+', default=None, metavar='pct, pct',
                        help='learning rate noise on/off epoch percentages')
    parser.add_argument('--lr-noise-pct', type=float, default=0.67, metavar='PERCENT',
                        help='learning rate noise limit percent (default: 0.67)')
    parser.add_argument('--lr-noise-std', type=float, default=1.0, metavar='STDDEV',
                        help='learning rate noise std-dev (default: 1.0)')
    parser.add_argument('--warmup-lr', type=float, default=1e-6, metavar='LR',
                        help='warmup learning rate (default: 1e-6)')
    parser.add_argument('--min-lr', type=float, default=1e-5, metavar='LR',
                        help='lower lr bound for cyclic schedulers that hit 0 (1e-5)')

    parser.add_argument('--decay-epochs', type=float, default=30, metavar='N',
                        help='epoch interval to decay LR')
    parser.add_argument('--warmup-epochs', type=int, default=5, metavar='N',
                        help='epochs to warmup LR, if scheduler supports')
    parser.add_argument('--cooldown-epochs', type=int, default=10, metavar='N',
                        help='epochs to cooldown LR at min_lr, after cyclic schedule ends')
    parser.add_argument('--patience-epochs', type=int, default=10, metavar='N',
                        help='patience epochs for Plateau LR scheduler (default: 10')
    parser.add_argument('--decay-rate', '--dr', type=float, default=0.1, metavar='RATE',
                        help='LR decay rate (default: 0.1)')

    # Augmentation parameters
    parser.add_argument('--color-jitter', type=float, default=0.3, metavar='PCT',
                        help='Color jitter factor (default: 0.3)')
    parser.add_argument('--aa', type=str, default='rand-m9-mstd0.5-inc1', metavar='NAME',
                        help='Use AutoAugment policy. "v0" or "original". " + \
                             "(default: rand-m9-mstd0.5-inc1)'),
    parser.add_argument('--smoothing', type=float, default=0.1, help='Label smoothing (default: 0.1)')
    parser.add_argument('--train-interpolation', type=str, default='bicubic',
                        help='Training interpolation (random, bilinear, bicubic default: "bicubic")')

    parser.add_argument('--repeated-aug', action='store_true')
    parser.add_argument('--no-repeated-aug', action='store_false', dest='repeated_aug')
    parser.set_defaults(repeated_aug=True)
    
    parser.add_argument('--train-mode', action='store_true')
    parser.add_argument('--no-train-mode', action='store_false', dest='train_mode')
    parser.set_defaults(train_mode=True)
    
    parser.add_argument('--ThreeAugment', action='store_true') #3augment
    
    parser.add_argument('--src', action='store_true') #simple random crop
    
    # add dataset parameters
    parser.add_argument('--global_crops_size', '--img_size', default=224, type=int,
                        help="this should be equal to image size")
    parser.add_argument('--patch_size', default=16, type=int,
                        help="patch size for vit patch embedding")
    
    # add masking parameter
    parser.add_argument('--mask_ratio', default=(0.1, 0.5), type=float, nargs='+',
                        help="mask ratio can be either a value or a range")
    parser.add_argument('--mask_probability', default=0., type=float,
                        help="how many samples with be applied with masking")
    parser.add_argument('--mask_first_n', action='store_true',
                        help="mask the first n sample to avoid shuffling. Needed for MAE-style encoder")
    parser.add_argument('--clone_batch', default=1, type=int,
                        help="how many times to clone the batch for masking (default: 1, not cloning)")
    
    # * Random Erase params
    parser.add_argument('--reprob', type=float, default=0.25, metavar='PCT',
                        help='Random erase prob (default: 0.25)')
    parser.add_argument('--remode', type=str, default='pixel',
                        help='Random erase mode (default: "pixel")')
    parser.add_argument('--recount', type=int, default=1,
                        help='Random erase count (default: 1)')
    parser.add_argument('--resplit', action='store_true', default=False,
                        help='Do not random erase first (clean) augmentation split')

    # * Mixup params
    parser.add_argument('--mixup', type=float, default=0.8,
                        help='mixup alpha, mixup enabled if > 0. (default: 0.8)')
    parser.add_argument('--cutmix', type=float, default=1.0,
                        help='cutmix alpha, cutmix enabled if > 0. (default: 1.0)')
    parser.add_argument('--cutmix-minmax', type=float, nargs='+', default=None,
                        help='cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)')
    parser.add_argument('--mixup-prob', type=float, default=1.0,
                        help='Probability of performing mixup or cutmix when either/both is enabled')
    parser.add_argument('--mixup-switch-prob', type=float, default=0.5,
                        help='Probability of switching to cutmix when both mixup and cutmix enabled')
    parser.add_argument('--mixup-mode', type=str, default='batch',
                        help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')

    # Distillation parameters
    parser.add_argument('--teacher-model', default='base', type=str)
    parser.add_argument('--teacher-path', type=str, default='')
    parser.add_argument('--distillation-type', default='none', choices=['none', 'soft', 'hard'], type=str, help="")
    parser.add_argument('--distillation-alpha', default=0.5, type=float, help="")
    parser.add_argument('--distillation-tau', default=1.0, type=float, help="")
    parser.add_argument('--lambda_token', type=float, default=1.0)
    parser.add_argument('--lambda_fea', type=float, default=1.0)
    parser.add_argument('--lambda_patch', type=float, default=1.0)
    
    # * Cosub params
    parser.add_argument('--cosub', action='store_true') 
    
    # * Finetuning params
    parser.add_argument('--finetune', default='', help='finetune from checkpoint')
    parser.add_argument('--attn-only', action='store_true')
    parser.add_argument('--weight_inherit', default='') 
    
    # Dataset parameters
    parser.add_argument('--data-path', default='/datasets01/imagenet_full_size/061417/', type=str,
                        help='dataset path')
    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'IMNET_ibot', 'IMNET_ibot_aug', 'IMNET_ibot_fast_aug', 'INAT', 'INAT19', 'IMNET_L', 'IMNET_L_ibot'],
                        type=str, help='Image Net dataset path')
    parser.add_argument('--inat-category', default='name',
                        choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
                        type=str, help='semantic granularity')

    parser.add_argument('--output_dir', default='',
                        help='path where to save, empty for no saving')
    parser.add_argument('--log_dir', default='/data1/qiyp/Proteus-pytorch/pretrain/log/DINOv2_training/log', 
                        type=str, help='saving logging info every 20 iters')
    parser.add_argument('--device', default='cuda',
                        help='device to use for training / testing')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                        help='start epoch')
    parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
    parser.add_argument('--eval-crop-ratio', default=0.875, type=float, help="Crop ratio for evaluation")
    parser.add_argument('--dist-eval', action='store_true', default=False, help='Enabling distributed evaluation')
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--pin-mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem',
                        help='')
    parser.set_defaults(pin_mem=True)

    # distributed training parameters
    parser.add_argument('--distributed', action='store_true', default=False, help='Enabling distributed training')
    parser.add_argument('--world_size', default=1, type=int,
                        help='number of distributed processes')
    parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training')
    return parser

def setup_logger(log_dir, rank=0):
    if rank != 0:
        return  # 只有主进程（rank 0）配置日志记录器
    log_formatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s]  %(message)s")
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)

    log_file_handler = logging.FileHandler(log_dir, encoding='utf-8')
    log_file_handler.setFormatter(log_formatter)
    root_logger.addHandler(log_file_handler)

    log_stream_handler = logging.StreamHandler(sys.stdout)
    log_stream_handler.setFormatter(log_formatter)
    root_logger.addHandler(log_stream_handler)

    logging.info('Logging file is %s' % log_dir)

class CustomCocoDataset(Dataset):
    def __init__(self, json_file, img_folder, common_transform=None):
        self.coco = COCO(json_file)
        self.img_folder = img_folder
        self.ids = list(self.coco.imgToAnns.keys())
        self.common_transform = common_transform
    
    def __len__(self):
        return len(self.ids)

    def __getitem__(self, index):
        img_id = self.ids[index]
        img_info = self.coco.loadImgs(img_id)[0]
        path = img_info['file_name']
        img_path = os.path.join(self.img_folder, path)
        image = Image.open(img_path).convert('RGB')

        # Perform a random crop
        i, j, h, w = transforms.RandomResizedCrop.get_params(
            image, scale=(0.9, 1.0), ratio=(1.0, 1.0))  # Ensure the same crop for both images
        
        cropped_image = transforms.functional.crop(image, i, j, h, w)

        # Resize to different resolutions
        jpg_image = transforms.functional.resize(cropped_image, 512, interpolation=transforms.InterpolationMode.BICUBIC)
        hint_image = transforms.functional.resize(cropped_image, 448, interpolation=transforms.InterpolationMode.BICUBIC)

        # Apply common transformations
        if self.common_transform is not None:
            jpg_image = self.common_transform(jpg_image)
            hint_image = self.common_transform(hint_image)

        ann_ids = self.coco.getAnnIds(imgIds=img_id)
        anns = self.coco.loadAnns(ann_ids)
        
        captions = [ann['caption'] for ann in anns]
        combined_caption = ' '.join(captions)
        
        return dict(jpg=jpg_image, txt=combined_caption, hint=hint_image)

def show_params(model, show_learnable=True):
    params = set()
    for name, param in model.named_parameters():
        if param.requires_grad == show_learnable:
            params.add(name)

    if show_learnable:
        logging.info("Parameters to be updated: ")
    else:
        logging.info("Parameters that are not being updated: ")

    for each in params:
        logging.info('\t{}'.format(str(each)))
    logging.info('\n')

def log_txt_as_img(wh, xc):
    # wh a tuple of (width, height)
    # xc a list of captions to plot
    b = len(xc)
    txts = list()
    for bi in range(b):
        txt = Image.new("RGB", wh, color="white")
        draw = ImageDraw.Draw(txt)
        # font = ImageFont.truetype('font/DejaVuSans.ttf', size=size)
        font = ImageFont.load_default()
        nc = int(40 * (wh[0] / 256))
        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))

        try:
            draw.text((0, 0), lines, fill="black", font=font)
        except UnicodeEncodeError:
            print("Cant encode string for logging. Skipping.")

        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
        txts.append(txt)
    txts = np.stack(txts)
    txts = torch.tensor(txts)
    return txts

def visualize_features(features):
    # Assuming features are of shape (batch_size, num_features, height, width)
    batch_size, num_features, height, width = features.shape
    
    # Normalize the feature maps to the range [0, 1]
    vis = features.mean(dim=1, keepdim=True)
    vis = vis - vis.min()
    vis = vis / vis.max()
    
    # Squeeze the channel dimension
    vis = vis.squeeze(1).cpu().detach().numpy()
    
    # Apply a colormap (e.g., viridis) to convert it to RGB
    vis_colored = np.zeros((batch_size, height, width, 3))
    for i in range(batch_size):
        vis_colored[i] = plt.cm.viridis(vis[i])[:, :, :3]  # Drop the alpha channel
    
    # Convert vis_colored to a tensor and save using torchvision
    vis_colored = torch.tensor(vis_colored).permute(0, 3, 1, 2)  # Convert to (batch, channels, height, width)

    return vis_colored

def main(args):
    utils.init_distributed_mode(args)

    print(args)

    device = torch.device(args.device)
     # 获取当前进程的 rank
    rank = dist.get_rank() if dist.is_initialized() else 0
    # set up logger
    os.makedirs(args.log_dir, exist_ok=True)
    setup_logger(args.log_dir + '/' + time.strftime('%Y%m%d_%H%M%S') + '.log', rank)
    logging.info('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
    logging.info("{}".format(args).replace(', ', ',\n') + '\n')

    # fix the seed for reproducibility
    seed = args.seed + utils.get_rank()
    torch.manual_seed(seed)
    np.random.seed(seed)
    # random.seed(seed)

    cudnn.benchmark = True
    save_img_path = args.save_img_path
    os.makedirs(save_img_path, exist_ok=True)

    # Setup an experiment folder
    exp_dir = args.exp_dir
    os.makedirs(exp_dir, exist_ok=True)
    ckpt_dir = os.path.join(exp_dir, "checkpoints")
    os.makedirs(ckpt_dir, exist_ok=True)
    logging.info(f"Experiment directory created at {exp_dir}")
    # dist.barrier() # wait for all processes to create the experiment directory

    # Create Dataset
    common_transforms = transforms.Compose([
        # transforms.RandomHorizontalFlip(),
        transforms.ToTensor()
    ])
    txt_path = args.txt_path
    img_folder = args.image_floder
    dataset = CustomCocoDataset(txt_path, img_folder, common_transform=common_transforms)
    logging.info(f"Loaded train dataset with {len(dataset)} samples")

    # Distributed Sampler
    if args.distributed:
        num_tasks = utils.get_world_size()
        global_rank = utils.get_rank()
        sampler_train = torch.utils.data.DistributedSampler(
            dataset, num_replicas=num_tasks, rank=global_rank, shuffle=True
        )
    else:
        sampler_train = torch.utils.data.RandomSampler(dataset)
    logging.info("Sampler_train = %s" % str(sampler_train))

    # Dataloader
    dataloader = DataLoader(
        dataset, sampler=sampler_train, batch_size=1, num_workers=args.num_workers,
        pin_memory=args.pin_mem, drop_last=True)
    logging.info('Dataloader created')


    # create IF model
    logging.info(f"# ===== Creating Feature Extractor: {args.model} ===== #")
    # print(f"Creating model: {args.model}")  # models_proteus_dinov2
    meta_arch_module = importlib.import_module(args.model)
    MetaArch = meta_arch_module.MetaArch

    model = MetaArch(args)
    logging.info("Model = %s" % str(model))

    if args.finetune:
        checkpoint = torch.load(args.finetune, map_location='cpu')

        if 'state_dict' in checkpoint:
            pretrained_dict = checkpoint['state_dict']
        elif 'model' in checkpoint:
            pretrained_dict = checkpoint['model']
        else:
            pretrained_dict = checkpoint

        missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, False)
        # print('missing_keys: ', missing_keys)
        # print('unexpected_keys: ', unexpected_keys)
        logging.info('Finetuning from %s' % args.finetune)
        logging.info('missing_keys: %s' % str(missing_keys))
        logging.info('unexpected_keys: %s' % str(unexpected_keys))
            
    model.to(device)
    model.eval()
    # create cldm
    logging.info(f"# ===== Creating ControlNet: ===== #")
    cldm = create_model(args.cldm_yaml).cpu()
    cldm.load_state_dict(load_state_dict(args.resume_path, location='cpu'))
    cldm.learning_rate = args.cldm_learning_rate
    cldm.sd_locked = args.sd_locked
    cldm.only_mid_control = args.only_mid_control
    cldm = cldm.to(device)

    if args.controlnet_dir:
        logging.info(f"Resuming controlnet from {args.controlnet_dir}")
        checkpoint = torch.load(args.controlnet_dir, map_location='cpu')
        cldm_state_dict = checkpoint['model_state_dict']
        cldm.control_model.load_state_dict(cldm_state_dict)

    if args.distributed:
        cldm = torch.nn.parallel.DistributedDataParallel(cldm, device_ids=[args.gpu], find_unused_parameters=True)
        cldm_without_ddp = cldm.module
    else:
        cldm_without_ddp = cldm

    ddim_sampler = DDIMSampler(cldm_without_ddp)
    
    logging.info("# ========== Building model done! ========== #")
    # Create SummaryWriter only in the main process
    if rank == 0:
        writer = SummaryWriter(exp_dir)
    epoch = 0
    ddim_steps = 50
    eta = 0
    cldm_without_ddp.eval()
    pbar = tqdm(iterable=dataloader, unit="batch")
    for i, batch in enumerate(dataloader):
        if args.distributed:
            sampler_train.set_epoch(epoch)

        gt = batch["jpg"].to(device)
        hint = batch["hint"].to(device)
        prompt = batch["txt"]
        _, _, H, W = gt.shape
        _, _, h, w = hint.shape
        with torch.no_grad():
            z_0 = cldm_without_ddp.get_latent(gt)
            features_dict = model.student.backbone(hint, is_training=True)
            features = features_dict['x_norm_patchtokens']
            features, _ = model.info_bottleneck(features, is_training=False)
            features = features.view(-1, h//14, w//14, features.shape[2])  # [B, h, w, c]
            features = features.permute(0, 3, 1, 2)
            # features = (features - features.mean()) / features.std()
            # features = torch.clamp(features, -5, 5)
            cond = cldm_without_ddp.prepare_condition(prompt, features) # cond["c_crossattn"]: txt embed, cond["c_concat"]: features

        # t = torch.randint(0, cldm_without_ddp.num_timesteps, (z_0.shape[0],), device=device).long()
        shape = (4, H // 8, W // 8)
        samples, intermediates = ddim_sampler.sample(ddim_steps, 1, shape, cond, verbose=False, eta=eta)
        x_samples = cldm_without_ddp.decode_first_stage(samples)

        vis_colored = visualize_features(features)
        out_feat_path = os.path.join(args.save_img_path, f"features_{i}.png")
        torchvision.utils.save_image(vis_colored, out_feat_path, normalize=True)
        # Normalize the x_samples from [-1, 1] to [0, 1]
        x_samples = (x_samples + 1) / 2
        out_img_path = os.path.join(args.save_img_path, f"output_{i}.png")
        # Save the images using torchvision.save_image
        torchvision.utils.save_image(x_samples, out_img_path, normalize=True)
        orig_img_path = os.path.join(args.save_img_path, f"original_{i}.png")
        torchvision.utils.save_image(gt, orig_img_path, normalize=True)

        hint_img_path = os.path.join(args.save_img_path, f"hint_{i}.png")
        torchvision.utils.save_image(hint, hint_img_path, normalize=True)
        
    pbar.close()
    logging.info("done!")
    if rank == 0:
        writer.close()

    # Ensure all processes wait before ending
    dist.barrier()

if __name__ == '__main__':
    parser = argparse.ArgumentParser('DeiT training and evaluation script', parents=[get_args_parser()])
    args = parser.parse_args()
    if args.output_dir:
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    main(args)