Inpaint / infer_inpaint.py
ZehanWang's picture
Upload ./infer_inpaint.py with huggingface_hub
7b88137 verified
import argparse
import logging
import os
import pdb
from peft import LoraConfig, get_peft_model
import torch
from safetensors.torch import load_model, save_model
from marigold.marigold_inpaint_pipeline import MarigoldInpaintPipeline
from marigold.duplicate_unet import DoubleUNet2DConditionModel
import json
from depth_anything_v2.dpt import DepthAnythingV2
from torchvision.transforms.functional import pil_to_tensor
from PIL import Image
import random
import numpy as np
from pycocotools import mask as coco_mask
from diffusers.schedulers import DDIMScheduler, PNDMScheduler
from torchvision.transforms import InterpolationMode, Resize, CenterCrop
import torchvision.transforms as transforms
model = MarigoldInpaintPipeline.from_pretrained('stabilityai/stable-diffusion-2')
unet_config_path = '/home/aiops/wangzh/.cache/huggingface/hub/models--stabilityai--stable-diffusion-2/snapshots/1e128c8891e52218b74cde8f26dbfc701cb99d79/unet/config.json'
# unet_checkpoint_path = '/home/aiops/wangzh/marigold/768_gen/diffusion_pytorch_model.safetensors'
model.unet = DoubleUNet2DConditionModel(**json.load(open(unet_config_path)))
# model.unet.load_state_dict(torch.load(unet_checkpoint_path, map_location='cpu'), strict=False)
model.unet.config["in_channels"] = 13
model.unet.duplicate_model()
model.unet.inpaint_rgb_conv_in()
model.unet.inpaint_depth_conv_in()
unet_lora_config = LoraConfig(
r=128,
lora_alpha=128,
init_lora_weights="gaussian",
target_modules=['to_k','to_q','to_v','to_out.0'],
)
model.unet = get_peft_model(model.unet, unet_lora_config)
sd2inpaint_ckpt = torch.load('/home/aiops/wangzh/marigold/output/512-inpaint-0.5-128-vitl-partition/checkpoint/latest/pytorch_model.bin', map_location='cpu')
model.unet.load_state_dict(sd2inpaint_ckpt)
model.to('cuda')
model_configs = {
'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536]}
}
model.rgb_scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2", subfolder="scheduler")
model.depth_scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2", subfolder="scheduler")
depth_model = DepthAnythingV2(**model_configs['vitl'])
depth_model.load_state_dict(
torch.load(f'/home/aiops/wangzh/Depth-Anything-V2/checkpoints/depth_anything_v2_vitl.pth', map_location='cpu'))
depth_model = depth_model.to('cuda').eval()
image_path = ['/dataset/~sa-1b/data/sa_001000/sa_10000335.jpg',
'/dataset/~sa-1b/data/sa_000357/sa_3572319.jpg',
'/dataset/~sa-1b/data/sa_000045/sa_457934.jpg']
prompt = ['A white car is parked in front of the factory',
'church with cemetery next to it',
'A house with a red brick roof']
imgs = [pil_to_tensor(Image.open(p)) for p in image_path]
depth_imgs = [depth_model(img.unsqueeze(0).cpu().numpy()) for img in imgs]
masks = []
for rgb_path in image_path:
anno = json.load(open(rgb_path.replace('.jpg', '.json')))['annotations']
random.shuffle(anno)
object_num = random.randint(5, 10)
mask = np.array(coco_mask.decode(anno[0]['segmentation']), dtype=np.uint8)
for single_anno in (anno[0:object_num] if len(anno)>object_num else anno):
mask += np.array(coco_mask.decode(single_anno['segmentation']), dtype=np.uint8)
mask = mask
mask = torch.stack([torch.tensor(mask) * 3], dim=0)
masks.append(mask)
# mask = torch.zeros((512,512))
# mask[100:300, 200:400] = 1
# masks.append(mask)
resize_transform = Resize(size=[512, 512], interpolation=InterpolationMode.NEAREST_EXACT)
imgs = [resize_transform(img) for img in imgs]
depth_imgs = [resize_transform(depth_img.unsqueeze(0)) for depth_img in depth_imgs]
masks = [resize_transform(mask.unsqueeze(0)) for mask in masks]
# for gs in [1,2,3,4,5]:
for i in range(len(imgs)):
output_image = model._rgbd_inpaint(imgs[i], depth_imgs[i].unsqueeze(0), masks[i], [prompt[i]], processing_res=512,
guidance_scale=3, mode='joint_inpaint' #'full_rgb_depth_inpaint', 'full_depth_rgb_inpaint', 'joint_inpaint'
)
output_image.save(f'./joint-{i}.jpg')