Upload 3 files

app.py

@@ -0,0 +1,216 @@
+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)

requirements.txt

@@ -0,0 +1,14 @@
+accelerate>=1.2.1
+diffusers>=0.32.2
+torch==2.5.1
+torchvision>=0.20.0
+torchaudio>=2.5.0
+gradio==4.44.0
+gradio-imageslider
+lpips==0.1.4
+pydantic>=2.6.1
+transformers==4.42.3
+xformers
+einops==0.7.0
+timm==0.9.16
+kornia==0.8.0

utils.py

@@ -0,0 +1,177 @@
+import os
+from typing import List
+
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import Image
+from torchvision import transforms
+
+from lbm.models.embedders import (
+    ConditionerWrapper,
+    LatentsConcatEmbedder,
+    LatentsConcatEmbedderConfig,
+)
+from lbm.models.lbm import LBMConfig, LBMModel
+from lbm.models.unets import DiffusersUNet2DCondWrapper
+from lbm.models.vae import AutoencoderKLDiffusers, AutoencoderKLDiffusersConfig
+
+
+def get_model_from_config(
+    backbone_signature: str = "stabilityai/stable-diffusion-xl-base-1.0",
+    vae_num_channels: int = 4,
+    unet_input_channels: int = 4,
+    timestep_sampling: str = "log_normal",
+    selected_timesteps: List[float] = None,
+    prob: List[float] = None,
+    conditioning_images_keys: List[str] = [],
+    conditioning_masks_keys: List[str] = ["mask"],
+    source_key: str = "source_image",
+    target_key: str = "source_image_paste",
+    bridge_noise_sigma: float = 0.0,
+):
+
+    conditioners = []
+
+    denoiser = DiffusersUNet2DCondWrapper(
+        in_channels=unet_input_channels,  # Add downsampled_image
+        out_channels=vae_num_channels,
+        center_input_sample=False,
+        flip_sin_to_cos=True,
+        freq_shift=0,
+        down_block_types=[
+            "DownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+        ],
+        mid_block_type="UNetMidBlock2DCrossAttn",
+        up_block_types=["CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"],
+        only_cross_attention=False,
+        block_out_channels=[320, 640, 1280],
+        layers_per_block=2,
+        downsample_padding=1,
+        mid_block_scale_factor=1,
+        dropout=0.0,
+        act_fn="silu",
+        norm_num_groups=32,
+        norm_eps=1e-05,
+        cross_attention_dim=[320, 640, 1280],
+        transformer_layers_per_block=[1, 2, 10],
+        reverse_transformer_layers_per_block=None,
+        encoder_hid_dim=None,
+        encoder_hid_dim_type=None,
+        attention_head_dim=[5, 10, 20],
+        num_attention_heads=None,
+        dual_cross_attention=False,
+        use_linear_projection=True,
+        class_embed_type=None,
+        addition_embed_type=None,
+        addition_time_embed_dim=None,
+        num_class_embeds=None,
+        upcast_attention=None,
+        resnet_time_scale_shift="default",
+        resnet_skip_time_act=False,
+        resnet_out_scale_factor=1.0,
+        time_embedding_type="positional",
+        time_embedding_dim=None,
+        time_embedding_act_fn=None,
+        timestep_post_act=None,
+        time_cond_proj_dim=None,
+        conv_in_kernel=3,
+        conv_out_kernel=3,
+        projection_class_embeddings_input_dim=None,
+        attention_type="default",
+        class_embeddings_concat=False,
+        mid_block_only_cross_attention=None,
+        cross_attention_norm=None,
+        addition_embed_type_num_heads=64,
+    ).to(torch.bfloat16)
+
+    if conditioning_images_keys != [] or conditioning_masks_keys != []:
+
+        latents_concat_embedder_config = LatentsConcatEmbedderConfig(
+            image_keys=conditioning_images_keys,
+            mask_keys=conditioning_masks_keys,
+        )
+        latent_concat_embedder = LatentsConcatEmbedder(latents_concat_embedder_config)
+        latent_concat_embedder.freeze()
+        conditioners.append(latent_concat_embedder)
+
+        # Wrap conditioners and set to device
+    conditioner = ConditionerWrapper(
+        conditioners=conditioners,
+    )
+
+    ## VAE ##
+    # Get VAE model
+    vae_config = AutoencoderKLDiffusersConfig(
+        version=backbone_signature,
+        subfolder="vae",
+        tiling_size=(128, 128),
+    )
+    vae = AutoencoderKLDiffusers(vae_config).to(torch.bfloat16)
+    vae.freeze()
+    vae.to(torch.bfloat16)
+
+    ## Diffusion Model ##
+    # Get diffusion model
+    config = LBMConfig(
+        source_key=source_key,
+        target_key=target_key,
+        timestep_sampling=timestep_sampling,
+        selected_timesteps=selected_timesteps,
+        prob=prob,
+        bridge_noise_sigma=bridge_noise_sigma,
+    )
+
+    sampling_noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        backbone_signature,
+        subfolder="scheduler",
+    )
+
+    model = LBMModel(
+        config,
+        denoiser=denoiser,
+        sampling_noise_scheduler=sampling_noise_scheduler,
+        vae=vae,
+        conditioner=conditioner,
+    ).to(torch.bfloat16)
+
+    return model
+
+
+def extract_object(birefnet, img):
+    # Data settings
+    image_size = (1024, 1024)
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    image = img
+    input_images = transform_image(image).unsqueeze(0).cuda()
+
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image.size)
+    image = Image.composite(image, Image.new("RGB", image.size, (127, 127, 127)), mask)
+    return image, mask
+
+
+def resize_and_center_crop(image, target_width, target_height):
+    original_width, original_height = image.size
+    scale_factor = max(target_width / original_width, target_height / original_height)
+    resized_width = int(round(original_width * scale_factor))
+    resized_height = int(round(original_height * scale_factor))
+    resized_image = image.resize((resized_width, resized_height), Image.LANCZOS)
+    left = (resized_width - target_width) / 2
+    top = (resized_height - target_height) / 2
+    right = (resized_width + target_width) / 2
+    bottom = (resized_height + target_height) / 2
+    cropped_image = resized_image.crop((left, top, right, bottom))
+    return cropped_image