import os
import sys
#sys.path.append('.')
import cv2
import einops
import numpy as np
import torch
import random
import gradio as gr
import albumentations as A
from PIL import Image
import torchvision.transforms as T
from mydatasets.data_utils import *
from cldm.model import create_model, load_state_dict
from cldm.ddim_hacked import DDIMSampler
from omegaconf import OmegaConf
from cldm.hack import disable_verbosity, enable_sliced_attention
from huggingface_hub import snapshot_download
snapshot_download(repo_id="xichenhku/AnyDoor_models", local_dir="./AnyDoor_models")
snapshot_download(repo_id="xichenhku/mask_refine", local_dir="./mask_refine")
cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)
save_memory = False
disable_verbosity()
if save_memory:
enable_sliced_attention()
config = OmegaConf.load('./configs/demo.yaml')
model_ckpt = config.pretrained_model
model_config = config.config_file
use_interactive_seg = config.config_file
model = create_model(model_config ).cpu()
model.load_state_dict(load_state_dict(model_ckpt, location='cuda'))
model = model.cuda()
ddim_sampler = DDIMSampler(model)
if use_interactive_seg:
from iseg.coarse_mask_refine_util import BaselineModel
model_path = './mask_refine/coarse_mask_refine.pth'
iseg_model = BaselineModel().eval()
weights = torch.load(model_path , map_location='cpu')['state_dict']
iseg_model.load_state_dict(weights, strict= True)
def crop_back( pred, tar_image, extra_sizes, tar_box_yyxx_crop):
H1, W1, H2, W2 = extra_sizes
y1,y2,x1,x2 = tar_box_yyxx_crop
pred = cv2.resize(pred, (W2, H2))
m = 3 # maigin_pixel
if W1 == H1:
tar_image[y1+m :y2-m, x1+m:x2-m, :] = pred[m:-m, m:-m]
return tar_image
if W1 < W2:
pad1 = int((W2 - W1) / 2)
pad2 = W2 - W1 - pad1
pred = pred[:,pad1: -pad2, :]
else:
pad1 = int((H2 - H1) / 2)
pad2 = H2 - H1 - pad1
pred = pred[pad1: -pad2, :, :]
tar_image[y1+m :y2-m, x1+m:x2-m, :] = pred[m:-m, m:-m]
return tar_image
def inference_single_image(ref_image,
ref_mask,
tar_image,
tar_mask,
strength,
ddim_steps,
scale,
seed,
enable_shape_control
):
raw_background = tar_image.copy()
item = process_pairs(ref_image, ref_mask, tar_image, tar_mask, enable_shape_control = enable_shape_control)
ref = item['ref']
hint = item['hint']
num_samples = 1
control = torch.from_numpy(hint.copy()).float().cuda()
control = torch.stack([control for _ in range(num_samples)], dim=0)
control = einops.rearrange(control, 'b h w c -> b c h w').clone()
clip_input = torch.from_numpy(ref.copy()).float().cuda()
clip_input = torch.stack([clip_input for _ in range(num_samples)], dim=0)
clip_input = einops.rearrange(clip_input, 'b h w c -> b c h w').clone()
H,W = 512,512
cond = {"c_concat": [control], "c_crossattn": [model.get_learned_conditioning( clip_input )]}
un_cond = {"c_concat": [control],
"c_crossattn": [model.get_learned_conditioning([torch.zeros((1,3,224,224))] * num_samples)]}
shape = (4, H // 8, W // 8)
if save_memory:
model.low_vram_shift(is_diffusing=True)
model.control_scales = ([strength] * 13)
samples, _ = ddim_sampler.sample(ddim_steps, num_samples,
shape, cond, verbose=False, eta=0,
unconditional_guidance_scale=scale,
unconditional_conditioning=un_cond)
if save_memory:
model.low_vram_shift(is_diffusing=False)
x_samples = model.decode_first_stage(samples)
x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy()
result = x_samples[0][:,:,::-1]
result = np.clip(result,0,255)
pred = x_samples[0]
pred = np.clip(pred,0,255)[1:,:,:]
sizes = item['extra_sizes']
tar_box_yyxx_crop = item['tar_box_yyxx_crop']
tar_image = crop_back(pred, tar_image, sizes, tar_box_yyxx_crop)
# keep background unchanged
y1,y2,x1,x2 = item['tar_box_yyxx']
raw_background[y1:y2, x1:x2, :] = tar_image[y1:y2, x1:x2, :]
return raw_background
def process_pairs(ref_image, ref_mask, tar_image, tar_mask, max_ratio = 0.8, enable_shape_control = False):
# ========= Reference ===========
# ref expand
ref_box_yyxx = get_bbox_from_mask(ref_mask)
# ref filter mask
ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)
masked_ref_image = ref_image * ref_mask_3 + np.ones_like(ref_image) * 255 * (1-ref_mask_3)
y1,y2,x1,x2 = ref_box_yyxx
masked_ref_image = masked_ref_image[y1:y2,x1:x2,:]
ref_mask = ref_mask[y1:y2,x1:x2]
ratio = np.random.randint(11, 15) / 10 #11,13
masked_ref_image, ref_mask = expand_image_mask(masked_ref_image, ref_mask, ratio=ratio)
ref_mask_3 = np.stack([ref_mask,ref_mask,ref_mask],-1)
# to square and resize
masked_ref_image = pad_to_square(masked_ref_image, pad_value = 255, random = False)
masked_ref_image = cv2.resize(masked_ref_image.astype(np.uint8), (224,224) ).astype(np.uint8)
ref_mask_3 = pad_to_square(ref_mask_3 * 255, pad_value = 0, random = False)
ref_mask_3 = cv2.resize(ref_mask_3.astype(np.uint8), (224,224) ).astype(np.uint8)
ref_mask = ref_mask_3[:,:,0]
# collage aug
masked_ref_image_compose, ref_mask_compose = masked_ref_image, ref_mask
ref_mask_3 = np.stack([ref_mask_compose,ref_mask_compose,ref_mask_compose],-1)
ref_image_collage = sobel(masked_ref_image_compose, ref_mask_compose/255)
# ========= Target ===========
tar_box_yyxx = get_bbox_from_mask(tar_mask)
tar_box_yyxx = expand_bbox(tar_mask, tar_box_yyxx, ratio=[1.1,1.2]) #1.1 1.3
tar_box_yyxx_full = tar_box_yyxx
# crop
tar_box_yyxx_crop = expand_bbox(tar_image, tar_box_yyxx, ratio=[1.3, 3.0])
tar_box_yyxx_crop = box2squre(tar_image, tar_box_yyxx_crop) # crop box
y1,y2,x1,x2 = tar_box_yyxx_crop
cropped_target_image = tar_image[y1:y2,x1:x2,:]
cropped_tar_mask = tar_mask[y1:y2,x1:x2]
tar_box_yyxx = box_in_box(tar_box_yyxx, tar_box_yyxx_crop)
y1,y2,x1,x2 = tar_box_yyxx
# collage
ref_image_collage = cv2.resize(ref_image_collage.astype(np.uint8), (x2-x1, y2-y1))
ref_mask_compose = cv2.resize(ref_mask_compose.astype(np.uint8), (x2-x1, y2-y1))
ref_mask_compose = (ref_mask_compose > 128).astype(np.uint8)
collage = cropped_target_image.copy()
collage[y1:y2,x1:x2,:] = ref_image_collage
collage_mask = cropped_target_image.copy() * 0.0
collage_mask[y1:y2,x1:x2,:] = 1.0
if enable_shape_control:
collage_mask = np.stack([cropped_tar_mask,cropped_tar_mask,cropped_tar_mask],-1)
# the size before pad
H1, W1 = collage.shape[0], collage.shape[1]
cropped_target_image = pad_to_square(cropped_target_image, pad_value = 0, random = False).astype(np.uint8)
collage = pad_to_square(collage, pad_value = 0, random = False).astype(np.uint8)
collage_mask = pad_to_square(collage_mask, pad_value = 2, random = False).astype(np.uint8)
# the size after pad
H2, W2 = collage.shape[0], collage.shape[1]
cropped_target_image = cv2.resize(cropped_target_image.astype(np.uint8), (512,512)).astype(np.float32)
collage = cv2.resize(collage.astype(np.uint8), (512,512)).astype(np.float32)
collage_mask = cv2.resize(collage_mask.astype(np.uint8), (512,512), interpolation = cv2.INTER_NEAREST).astype(np.float32)
collage_mask[collage_mask == 2] = -1
masked_ref_image = masked_ref_image / 255
cropped_target_image = cropped_target_image / 127.5 - 1.0
collage = collage / 127.5 - 1.0
collage = np.concatenate([collage, collage_mask[:,:,:1] ] , -1)
item = dict(ref=masked_ref_image.copy(), jpg=cropped_target_image.copy(), hint=collage.copy(),
extra_sizes=np.array([H1, W1, H2, W2]),
tar_box_yyxx_crop=np.array( tar_box_yyxx_crop ),
tar_box_yyxx=np.array(tar_box_yyxx_full),
)
return item
ref_dir='./examples/Gradio/FG'
image_dir='./examples/Gradio/BG'
ref_list=[os.path.join(ref_dir,file) for file in os.listdir(ref_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file ]
ref_list.sort()
image_list=[os.path.join(image_dir,file) for file in os.listdir(image_dir) if '.jpg' in file or '.png' in file or '.jpeg' in file]
image_list.sort()
def mask_image(image, mask):
blanc = np.ones_like(image) * 255
mask = np.stack([mask,mask,mask],-1) / 255
masked_image = mask * ( 0.5 * blanc + 0.5 * image) + (1-mask) * image
return masked_image.astype(np.uint8)
def run_local(base,
ref,
*args):
image = base["image"].convert("RGB")
mask = base["mask"].convert("L")
ref_image = ref["image"].convert("RGB")
ref_mask = ref["mask"].convert("L")
image = np.asarray(image)
mask = np.asarray(mask)
mask = np.where(mask > 128, 1, 0).astype(np.uint8)
ref_image = np.asarray(ref_image)
ref_mask = np.asarray(ref_mask)
ref_mask = np.where(ref_mask > 128, 1, 0).astype(np.uint8)
synthesis = inference_single_image(ref_image.copy(), ref_mask.copy(), image.copy(), mask.copy(), *args)
synthesis = torch.from_numpy(synthesis).permute(2, 0, 1)
synthesis = synthesis.permute(1, 2, 0).numpy()
return [synthesis]
with gr.Blocks() as demo:
with gr.Column():
gr.Markdown("# Play with AnyDoor to Teleport your Target Objects! ")
with gr.Row():
baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", columns=1, height=768)
with gr.Accordion("Advanced Option", open=True):
num_samples = 1
strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=30, step=1)
scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=4.5, step=0.1)
seed = gr.Slider(label="Seed", minimum=-1, maximum=999999999, step=1, value=-1)
enable_shape_control = gr.Checkbox(label='Enable Shape Control', value=False)
gr.Markdown("### Guidelines")
gr.Markdown(" Higher guidance-scale makes higher fidelity, while lower one makes more harmonized blending.")
gr.Markdown(" Enable shape control means the generation results would consider user-drawn masks; otherwise it \
considers the location and size to adjust automatically.")
gr.Markdown(" Users should annotate the mask of the target object, too coarse mask would lead to bad generation.")
gr.Markdown("# Upload / Select Images for the Background (left) and Reference Object (right)")
gr.Markdown("### You could draw coarse masks on the background to indicate the desired location and shape.")
gr.Markdown("### <u>Do not forget</u> to annotate the target object on the reference image.")
with gr.Row():
base = gr.Image(label="Background", source="upload", tool="sketch", type="pil", height=512, brush_color='#FFFFFF', mask_opacity=0.5)
ref = gr.Image(label="Reference", source="upload", tool="sketch", type="pil", height=512, brush_color='#FFFFFF', mask_opacity=0.5)
run_local_button = gr.Button(label="Generate", value="Run")
with gr.Row():
with gr.Column():
gr.Examples(image_list, inputs=[base],label="Examples - Background Image",examples_per_page=16)
with gr.Column():
gr.Examples(ref_list, inputs=[ref],label="Examples - Reference Object",examples_per_page=16)
run_local_button.click(fn=run_local,
inputs=[base,
ref,
strength,
ddim_steps,
scale,
seed,
enable_shape_control,
],
outputs=[baseline_gallery]
)
demo.launch()