app.py

import glob
from copy import deepcopy

import gradio as gr
import numpy as np
import PIL
import spaces
import torch
import yaml
from gradio_imageslider import ImageSlider
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file
from torchvision.transforms import ToPILImage, ToTensor
from transformers import AutoModelForImageSegmentation
from utils import extract_object, get_model_from_config, resize_and_center_crop

ASPECT_RATIOS = {
    str(512 / 2048): (512, 2048),
    str(1024 / 1024): (1024, 1024),
    str(2048 / 512): (2048, 512),
    str(896 / 1152): (896, 1152),
    str(1152 / 896): (1152, 896),
    str(512 / 1920): (512, 1920),
    str(640 / 1536): (640, 1536),
    str(768 / 1280): (768, 1280),
    str(1280 / 768): (1280, 768),
    str(1536 / 640): (1536, 640),
    str(1920 / 512): (1920, 512),
}

# download the config and model
MODEL_PATH = hf_hub_download("jasperai/LBM_relighting", "relight.safetensors")
CONFIG_PATH = hf_hub_download("jasperai/LBM_relighting", "relight.yaml")

with open(CONFIG_PATH, "r") as f:
    config = yaml.safe_load(f)
model = get_model_from_config(**config)
sd = load_file(MODEL_PATH)
model.load_state_dict(sd, strict=True)
model.to("cuda").to(torch.bfloat16)
birefnet = AutoModelForImageSegmentation.from_pretrained(
    "ZhengPeng7/BiRefNet", trust_remote_code=True
).cuda()
image_size = (1024, 1024)


@spaces.GPU
def evaluate(
    fg_image: PIL.Image.Image,
    bg_image: PIL.Image.Image,
    num_sampling_steps: int = 1,
):

    ori_h_bg, ori_w_bg = fg_image.size
    ar_bg = ori_h_bg / ori_w_bg
    closest_ar_bg = min(ASPECT_RATIOS, key=lambda x: abs(float(x) - ar_bg))
    dimensions_bg = ASPECT_RATIOS[closest_ar_bg]

    _, fg_mask = extract_object(birefnet, deepcopy(fg_image))

    fg_image = resize_and_center_crop(fg_image, dimensions_bg[0], dimensions_bg[1])
    fg_mask = resize_and_center_crop(fg_mask, dimensions_bg[0], dimensions_bg[1])
    bg_image = resize_and_center_crop(bg_image, dimensions_bg[0], dimensions_bg[1])

    img_pasted = Image.composite(fg_image, bg_image, fg_mask)

    img_pasted_tensor = ToTensor()(img_pasted).unsqueeze(0) * 2 - 1
    batch = {
        "source_image": img_pasted_tensor.cuda().to(torch.bfloat16),
    }

    z_source = model.vae.encode(batch[model.source_key])

    output_image = model.sample(
        z=z_source,
        num_steps=num_sampling_steps,
        conditioner_inputs=batch,
        max_samples=1,
    ).clamp(-1, 1)

    output_image = (output_image[0].float().cpu() + 1) / 2
    output_image = ToPILImage()(output_image)

    # paste the output image on the background image
    output_image = Image.composite(output_image, bg_image, fg_mask)

    output_image.resize((ori_h_bg, ori_w_bg))
    print(output_image.size, img_pasted.size)

    return (np.array(img_pasted), np.array(output_image))


with gr.Blocks(title="LBM Object Relighting") as demo:
    gr.Markdown(
        f"""
        # Object Relighting with Latent Bridge Matching
        This is an interactive demo of [LBM: Latent Bridge Matching for Fast Image-to-Image Translation](https://arxiv.org/abs/2403.03025) *by Jasper Research*. We are internally exploring the possibility of releasing the model. If you enjoy the space, please also promote *open-source* by giving a ⭐ to the <a href='https://github.com/gojasper/LBM' target='_blank'>Github Repo</a>.
    """
    )
    gr.Markdown(
        "💡 *Hint:* To better appreciate the low latency of our method, run the demo locally !"
    )
    with gr.Row():
        with gr.Column():
            with gr.Row():
                fg_image = gr.Image(
                    type="pil",
                    label="Input Image",
                    image_mode="RGB",
                    height=360,
                    # width=360,
                )
                bg_image = gr.Image(
                    type="pil",
                    label="Target Background",
                    image_mode="RGB",
                    height=360,
                    # width=360,
                )

            with gr.Row():
                submit_button = gr.Button("Relight", variant="primary")
            with gr.Row():
                num_inference_steps = gr.Slider(
                    minimum=1,
                    maximum=4,
                    value=1,
                    step=1,
                    label="Number of Inference Steps",
                )

            bg_gallery = gr.Gallery(
                # height=450,
                object_fit="contain",
                label="Background List",
                value=[path for path in glob.glob("examples/backgrounds/*.jpg")],
                columns=5,
                allow_preview=False,
            )

        with gr.Column():
            output_slider = ImageSlider(label="Composite vs LBM", type="numpy")
            output_slider.upload(
                fn=evaluate,
                inputs=[fg_image, bg_image, num_inference_steps],
                outputs=[output_slider],
            )

    submit_button.click(
        evaluate,
        inputs=[fg_image, bg_image, num_inference_steps],
        outputs=[output_slider],
        show_progress=False,
        show_api=False,
    )

    with gr.Row():
        gr.Examples(
            fn=evaluate,
            examples=[
                [
                    "examples/foregrounds/2.jpg",
                    "examples/backgrounds/14.jpg",
                    1,
                ],
                [
                    "examples/foregrounds/10.jpg",
                    "examples/backgrounds/4.jpg",
                    1,
                ],
                [
                    "examples/foregrounds/11.jpg",
                    "examples/backgrounds/24.jpg",
                    1,
                ],
                [
                    "examples/foregrounds/19.jpg",
                    "examples/backgrounds/3.jpg",
                    1,
                ],
                [
                    "examples/foregrounds/4.jpg",
                    "examples/backgrounds/6.jpg",
                    1,
                ],
                [
                    "examples/foregrounds/14.jpg",
                    "examples/backgrounds/22.jpg",
                    1,
                ],
                [
                    "examples/foregrounds/12.jpg",
                    "examples/backgrounds/1.jpg",
                    1,
                ],
            ],
            inputs=[fg_image, bg_image, num_inference_steps],
            outputs=[output_slider],
            run_on_click=True,
        )

    gr.Markdown("**Disclaimer:**")
    gr.Markdown(
        "This demo is only for research purpose. Jasper cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. Jasper provides the tools, but the responsibility for their use lies with the individual user."
    )

    def bg_gallery_selected(gal, evt: gr.SelectData):
        print(gal, evt.index)
        return gal[evt.index][0]

    bg_gallery.select(bg_gallery_selected, inputs=bg_gallery, outputs=bg_image)

if __name__ == "__main__":

    demo.queue().launch(share=True, show_api=False)