vipainting.py

from functools import partial
import os
import argparse
import yaml
from omegaconf import OmegaConf
from ldm.util import instantiate_from_config, get_obj_from_str
import torch
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from utils.logger import get_logger
from utils.mask_generator import mask_generator
from utils.helper import encoder_kl, clean_directory, to_img, encoder_vq, load_file
from ldm.guided_diffusion.h_posterior import HPosterior
from PIL import Image
import numpy as np 
from torchvision.transforms.functional import pil_to_tensor

def load_yaml(file_path: str) -> dict:
    with open(file_path) as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    return config

def save_segmentation(s, img_path, name):
  s = s.detach().cpu().numpy().transpose(0,2,3,1)[0,:,:,None,:]
  colorize = np.random.RandomState(1).randn(1,1,s.shape[-1],3)
  colorize = colorize / colorize.sum(axis=2, keepdims=True)
  s = s@colorize
  s = s[...,0,:]
  s = ((s+1.0)*127.5).clip(0,255).astype(np.uint8)
  s = Image.fromarray(s)
  s.save(os.path.join(img_path, name))

def vipaint(num, mask_web, image_queue, sampling_queue):
    parser = argparse.ArgumentParser()
    parser.add_argument('--inpaint_config', type=str, default='configs/inpainting/lands_config_mountain.yaml') #lsun_config, imagenet_config
    parser.add_argument('--working_directory', type=str, default='results/')
    parser.add_argument('--gpu', type=int, default=0)
    parser.add_argument('--id', type=int, default=0)
    parser.add_argument('--k_steps', type=int, default=2)
    parser.add_argument('--case', type=str, default="random_all")
    args = parser.parse_args()
   
    
    # Device setting
    print("================= Device setting")
    device_str = f"cuda:{args.gpu}" if torch.cuda.is_available() else 'cpu'
    device = torch.device(device_str)  
    
    # Load configurations
    print("================= Load config") 
    inpaint_config = load_yaml(args.inpaint_config)
    working_directory = args.working_directory

    # Load model
    print("================= Load model") 
    config = OmegaConf.load(inpaint_config['diffusion'])
    vae_config = OmegaConf.load(inpaint_config['autoencoder'])

    diff = instantiate_from_config(config.model)
    diff.load_state_dict(torch.load(inpaint_config['diffusion_model'],
                                     map_location='cpu')["state_dict"], strict=False)    
    diff = diff.to(device)
    diff.model.eval()
    diff.first_stage_model.eval()
    diff.eval()

    # Load pre-trained autoencoder loss config
    print("================= Load pre-trained") 
    loss_config = vae_config['model']['params']['lossconfig']
    vae_loss = get_obj_from_str(inpaint_config['name'],
                                 reload=False)(**loss_config.get("params", dict()))
    
    # Load test data
    print("================= Load test data") 
    if os.path.exists(inpaint_config['data']['file_name']):
        dataset = np.load(inpaint_config['data']['file_name'])
    loader = torch.utils.data.DataLoader(dataset= dataset, batch_size=1)

    # Working directory
    print("================= working directory") 
    out_path = working_directory 
    os.makedirs(out_path, exist_ok=True)
        

    #mask = torch.tensor(np.load("masks/mask_" + str(args.id) + ".npy")).to(device)
    posterior = inpaint_config['posterior']
    if args.k_steps == 1: 
        posterior = "gauss"
        t_steps_hierarchy = [400]
    else : 
        posterior = "hierarchical"
        if args.k_steps == 2: t_steps_hierarchy = [inpaint_config[posterior]['t_steps_hierarchy'][0], 
                                                   inpaint_config[posterior]['t_steps_hierarchy'][-1]]
        elif args.k_steps == 4: t_steps_hierarchy = inpaint_config[posterior]['t_steps_hierarchy'] # [550, 500, 450, 400]
        elif args.k_steps == 6: t_steps_hierarchy = [650, 600, 550, 500, 450, 400]

        
    # Prepare VI method
    print("=================== Prepare VI method")
    h_inpainter = HPosterior(diff, vae_loss, 
                             eta = inpaint_config[posterior]["eta"],
                             z0_size = inpaint_config["data"]["latent_size"],
                             img_size = inpaint_config["data"]["image_size"],
                             latent_channels = inpaint_config["data"]["latent_channels"],
                             first_stage=inpaint_config[posterior]["first_stage"],
                             t_steps_hierarchy=t_steps_hierarchy, #inpaint_config[posterior]['t_steps_hierarchy'],
                             posterior = inpaint_config['posterior'], image_queue = image_queue,
                             sampling_queue = sampling_queue) 
    
    h_inpainter.descretize(inpaint_config[posterior]['rho']) 

    x_size = inpaint_config['mask_opt']['image_size']
    channels = inpaint_config['data']['channels']
    
    # Do Inference
    print("=================== Do Inference")
    imgs = [num]
    for i, random_num in enumerate(imgs):
        img_path = os.path.join(out_path, str(random_num) )  # +str(args.k_steps) + "_h" #"Loss-ablation"
        for img_dir in ['progress', 'params', 'mus']:
            sub_dir = os.path.join(img_path, img_dir)
            os.makedirs(sub_dir, exist_ok=True)

        bs = inpaint_config[posterior]["batch_size"]

        batch_size = bs
        channels = 182 
        # For conditional models
        segmentation = loader.dataset["segmentation"][random_num]
        if inpaint_config["conditional_model"] : 
            segment_c = torch.tensor(segmentation.transpose(2,0,1)[None]).to(dtype=torch.float32, device=diff.device)
            segment_c = segment_c.repeat(batch_size, 1, 1, 1)  
            uc = diff.get_learned_conditioning(
                        {diff.cond_stage_key: segment_c.to(diff.device)}['segmentation']
                        ).detach()

        #Get Image/Labels
        print("==================== get image/labels")
        #Get Image/Labels
        if len(loader.dataset) ==2: 
            ref_img = loader.dataset["images"][random_num] #512, 512, 3
            ref_img = torch.tensor(ref_img[None]).to(dtype=torch.float32, device=diff.device)
            print(f"ref_img {ref_img.shape}") #1, 512, 512, 3
            ref_img = ref_img/127.5 - 1

            label = torch.tensor(segmentation.transpose(2,0,1)[None]).to(dtype=torch.float32, device=diff.device)
            save_segmentation(label, img_path, 'input.png')
            label = label.repeat(batch_size, 1, 1, 1)  # Now shape is [batch_size, 182, 128, 128]
            xc = torch.tensor(label)
            c = diff.get_learned_conditioning({diff.cond_stage_key: xc}['segmentation']).detach()
        else:
            ref_img = loader.dataset[random_num].reshape(1,x_size,x_size,channels)
            c = None
            uc = None
        
        ref_img = torch.tensor(ref_img).to(device)

        # #Get mask
        mask_tensor = torch.tensor(mask_web).to(device)
        mask_tensor = mask_tensor.float() / 255.0  # Convert to float and normalize to [0, 1]
        ref_img = torch.permute(ref_img, (0,3,1,2)) 
        y = torch.Tensor.repeat(mask_tensor*ref_img, [bs,1,1,1]).float()
        
        if inpaint_config[posterior]["first_stage"] == "kl":
            y_encoded = encoder_kl(diff, y)[0]
        else:
            y_encoded = encoder_vq(diff, y)

        # print(f"shape {ref_img.shape} {mask.shape}")
        plt.imsave(os.path.join(img_path, 'true.png'), to_img(ref_img).astype(np.uint8)[0])
        plt.imsave(os.path.join(img_path, 'observed.png'), to_img(y).astype(np.uint8)[0])
        
        lambda_ = h_inpainter.init(y_encoded, inpaint_config["init"]["var_scale"], 
                                inpaint_config[posterior]["mean_scale"], inpaint_config["init"]["prior_scale"],
                                inpaint_config[posterior]["mean_scale_top"])
        # Fit posterior once
        print("============ fit posterior once")
        torch.cuda.empty_cache()  
        h_inpainter.fit(lambda_ = lambda_, cond=c, shape = (bs, *y_encoded.shape[1:]),
                quantize_denoised=False, mask_pixel = mask_tensor, y =y,
                log_every_t=25, iterations = inpaint_config[posterior]['iterations'],
                unconditional_guidance_scale= inpaint_config[posterior]["unconditional_guidance_scale"] ,
                unconditional_conditioning=uc, kl_weight_1=inpaint_config[posterior]["beta_1"],
                 kl_weight_2 = inpaint_config[posterior]["beta_2"],
                debug=True, wdb = False,
                dir_name = img_path,
                batch_size = bs, 
                lr_init_gamma = inpaint_config[posterior]["lr_init_gamma"],
                recon_weight = inpaint_config[posterior]["recon"], 
                )  
        
        # Load parameters and sample
        print("============= load parameters and sample") 
        params_path = os.path.join(img_path, 'params', f'{inpaint_config[posterior]["iterations"]}.pt') #, j+1
        [mu, logvar, gamma] = torch.load(params_path)
        
        h_inpainter.sample(inpaint_config["sampling"]["scale"], inpaint_config[posterior]["eta"],
                            mu.cuda(), logvar.cuda(), gamma.cuda(), mask_tensor,  y,
                            n_samples=inpaint_config["sampling"]["n_samples"], 
                            batch_size = bs, dir_name= img_path, cond=c, 
                            unconditional_conditioning=uc, 
                            unconditional_guidance_scale=inpaint_config["sampling"]["unconditional_guidance_scale"], 
                            samples_iteration=inpaint_config[posterior]["iterations"])