app.py

import gradio as gr
import os
import sys
import torch
import cv2
import einops
import numpy as np
import spaces
from omegaconf import OmegaConf
from huggingface_hub import snapshot_download
from PIL import Image

REPO_DIR = snapshot_download(repo_id="Hang2991/PICS")
os.chdir(REPO_DIR)
sys.path.insert(0, REPO_DIR)
sys.path.insert(0, os.path.join(REPO_DIR, "dinov2"))

from cldm.model import create_model, load_state_dict
from cldm.ddim_hacked import DDIMSampler
from datasets.data_utils import *
config = OmegaConf.load('configs/inference.yaml')
model = create_model(config.config_file).cpu()
model.load_state_dict(load_state_dict(config.pretrained_model, location='cpu'))
model.eval()

def get_input(batch, k):
    x = batch[k]
    if len(x.shape) == 3:
        x = x[None, ...]
    x = torch.tensor(x)
    x = einops.rearrange(x, 'b h w c -> b c h w')
    x = x.to(memory_format=torch.contiguous_format).float()
    return x

def get_unconditional_conditioning(N, obj_thr):
    x = [torch.zeros((1, 3, 224, 224)).to(model.device)] * N
    single_uc = model.get_learned_conditioning(x)
    uc = single_uc.unsqueeze(-1).repeat(1, 1, 1, obj_thr)
    return {"pch_code": uc}

def process_pairs_multiple(mask, tar_image, patch_dir, counter=0, max_ratio=0.8):
    view = cv2.imread(patch_dir)
    view = cv2.cvtColor(view, cv2.COLOR_BGR2RGB)
    view = pad_to_square(view, pad_value=255, random=False)
    view = cv2.resize(view.astype(np.uint8), (224, 224))
    view = view.astype(np.float32) / 255.0
    box_yyxx = get_bbox_from_mask(mask)
    H1, W1 = tar_image.shape[0], tar_image.shape[1]
    box_yyxx_crop = [0, H1, 0, W1]
    y1, y2, x1, x2 = box_in_box(box_yyxx, box_yyxx_crop)
    collage = tar_image.copy()
    source_collage = collage.copy()
    collage[y1:y2, x1:x2, :] = 0
    collage_mask = np.zeros_like(tar_image, dtype=np.float32)
    collage_mask[y1:y2, x1:x2, :] = 1.0
    tar_square = pad_to_square(tar_image, pad_value=0, random=False)
    collage_square = pad_to_square(collage, pad_value=0, random=False)
    mask_square = pad_to_square(collage_mask, pad_value=2, random=False)
    H2, W2 = collage_square.shape[0], collage_square.shape[1]
    tar_res = cv2.resize(tar_square, (512, 512)).astype(np.float32)
    col_res = cv2.resize(collage_square, (512, 512)).astype(np.float32)
    mask_res = cv2.resize(mask_square, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.float32)
    mask_res[mask_res == 2] = -1
    c_mask = np.where(mask_res[..., 0:1] == 1, 1.0, 0.0).astype(np.float32)
    tar_res = tar_res / 127.5 - 1.0
    col_res = col_res / 127.5 - 1.0
    hint_final = np.concatenate([col_res, mask_res[..., :1]], axis=-1)
    return {
        f'view{counter}': view, f'hint{counter}': hint_final, f'mask{counter}': c_mask,
        f'hint_sizes{counter}': np.array([y1, x1, y2, x2]), 'jpg': tar_res,
        'collage': source_collage, 'extra_sizes': np.array([H1, W1, H2, W2])
    }

def process_composition(item, obj_thr):
    collage = item['collage'].copy()
    collage_mask = np.zeros((collage.shape[0], collage.shape[1], 1), dtype=np.float32)
    for i in reversed(range(obj_thr)):
        y1, x1, y2, x2 = item['hint_sizes'+str(i)]
        collage[y1:y2, x1:x2, :] = 0
        collage_mask[y1:y2,x1:x2,:] = 1.0
    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.float32)
    collage = cv2.resize(collage.astype(np.uint8), (512, 512)).astype(np.float32) / 127.5 - 1.0 
    collage_mask = cv2.resize(collage_mask, (512, 512), interpolation=cv2.INTER_NEAREST).astype(np.float32)
    if len(collage_mask.shape) == 2: collage_mask = collage_mask[..., None]
    collage_mask[collage_mask == 2] = -1.0
    collage_final = np.concatenate([collage, collage_mask[:,:,:1]] , -1)
    item.update({'hint': collage_final.copy()})
    return item

def load_example_pil(path):
    return Image.open(path).convert("RGB")

@spaces.GPU(duration=120)
def pics_pairwise_inference(background, img_a, mask_a, img_b, mask_b):
    device = "cuda"
    model.to(device)
    ddim_sampler = DDIMSampler(model)
    back_image = np.array(background)
    
    item_with_collage = {}
    objs = [(img_a, mask_a), (img_b, mask_b)]
    for j, (img, mask) in enumerate(objs):
        temp_patch = f"temp_obj_{j}.png"
        cv2.imwrite(temp_patch, cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
        tar_mask = (np.array(mask)[:, :, 0] > 128).astype(np.uint8)
        item_with_collage.update(process_pairs_multiple(tar_mask, back_image, temp_patch, counter=j))

    item_with_collage = process_composition(item_with_collage, obj_thr=2)

    obj_thr = 2
    num_samples = 1
    H, W = 512, 512
    guidance_scale = 5.0
    
    xc = []
    xc_mask = []
    for i in range(obj_thr):
        xc.append(get_input(item_with_collage, f"view{i}").to(device))
        xc_mask.append(get_input(item_with_collage, f"mask{i}"))

    c_list = [model.get_learned_conditioning(xc_i) for xc_i in xc]
    c_tensor = torch.stack(c_list).permute(1, 2, 3, 0)
    cond_cross = {"pch_code": c_tensor}
    c_mask = torch.stack(xc_mask).permute(1, 2, 3, 4, 0).to(device)

    hint = item_with_collage['hint']
    control = torch.from_numpy(hint.copy()).float().to(device)
    control = torch.stack([control] * num_samples, dim=0)
    control = einops.rearrange(control, 'b h w c -> b c h w').clone()

    cond = {"c_concat": [control], "c_crossattn": [cond_cross], "c_mask": [c_mask]}
    uc_pch = get_unconditional_conditioning(num_samples, obj_thr)
    un_cond = {"c_concat": [control], "c_crossattn": [uc_pch], "c_mask": [c_mask]}

    shape = (4, H // 8, W // 8)
    model.control_scales = [1.0] * 13
    
    samples, _ = ddim_sampler.sample(50, num_samples, shape, cond, verbose=False, eta=0.0,
                                     unconditional_guidance_scale=guidance_scale, unconditional_conditioning=un_cond)

    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()
    pred = np.clip(x_samples[0], 0, 255).astype(np.uint8)
    
    side = max(back_image.shape[0], back_image.shape[1])
    pred_res = cv2.resize(pred, (side, side))

    final_image = crop_back(pred_res, back_image, item_with_collage['extra_sizes'], 
                            item_with_collage['hint_sizes0'], item_with_collage['hint_sizes1'], is_masked=True)
    
    return final_image

with gr.Blocks(title="PICS: Pairwise Spatial Compositing with Spatial Interactions") as demo:
    gr.Markdown("# 🚀 PICS: Pairwise Image Compositing with Spatial Interactions")
    gr.Markdown("Submit **Background**, **Two Objects**, and their **Two Masks** to reason about spatial interactions.")

    with gr.Row():
        with gr.Column(scale=2):
            bg_input = gr.Image(label="1. Scene Background", type="pil")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Object A")
                    obj_a_img = gr.Image(label="Image A", type="pil")
                    obj_a_mask = gr.Image(label="Mask A", type="pil")
                
                with gr.Column():
                    gr.Markdown("### Object B")
                    obj_b_img = gr.Image(label="Image B", type="pil")
                    obj_b_mask = gr.Image(label="Mask B", type="pil")
            
            run_btn = gr.Button("Execute PICS Inference ✨", variant="primary")

        with gr.Column(scale=1):
            output_img = gr.Image(label="PICS Composite Result")
            gr.Markdown("""
            ### 🎨 PICS Reasoning Logic
            * **Pairwise Interaction**: Model reasons about spatial relations between Object A and B.
            * **Composition**: It intelligently composites objects into the provided scene background.
            """)

    gr.Markdown("### 💡 Quick Examples")
    gr.Examples(
        examples=[
            [
                load_example_pil("sample/bread_basket/image.jpg"), 
                load_example_pil("sample/bread_basket/object_0.png"), 
                load_example_pil("sample/bread_basket/object_0_mask.png"), 
                load_example_pil("sample/bread_basket/object_1.png"), 
                load_example_pil("sample/bread_basket/object_1_mask.png")
            ], 
            [
                load_example_pil("sample/pen_penholder/image.jpg"), 
                load_example_pil("sample/pen_penholder/object_0.png"), 
                load_example_pil("sample/pen_penholder/object_0_mask.png"), 
                load_example_pil("sample/pen_penholder/object_1.png"), 
                load_example_pil("sample/pen_penholder/object_1_mask.png")
            ]
        ],
        inputs=[bg_input, obj_a_img, obj_a_mask, obj_b_img, obj_b_mask],
        cache_examples=False, 
    )

    run_btn.click(
        fn=pics_pairwise_inference,
        inputs=[bg_input, obj_a_img, obj_a_mask, obj_b_img, obj_b_mask],
        outputs=output_img
    )

if __name__ == "__main__":
    demo.launch(allowed_paths=[REPO_DIR])