Update app.py

app.py

@@ -31,7 +31,6 @@ ASPECT_RATIOS = {
     str(1920 / 512): (1920, 512),
 }
 
-# download the config and model
 MODEL_PATH = hf_hub_download("jasperai/LBM_relighting", "model.safetensors", token=huggingface_token)
 CONFIG_PATH = hf_hub_download("jasperai/LBM_relighting", "config.yaml", token=huggingface_token)
 
@@ -41,20 +40,11 @@ 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()
+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,
-):
-    gr.Info("Relighting Image...", duration=3)
-
+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))
@@ -67,159 +57,69 @@ def evaluate(
     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),
-    }
+    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 = 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(css="""
+body::before {
+    content: "";
+    display: block;
+    height: 640px;
+    background-color: #0f1117;
+}
+button[aria-label="Download"] {
+    transform: scale(2);
+    transform-origin: top right;
+    margin: 0 !important;
+    padding: 6px !important;
+}
+button[aria-label="Share"] {
+    display: none;
+}
+button[aria-label="Copy link"] {
+    display: none;
+}
+button[aria-label="Open in new tab"] {
+    display: none;
+}
+""", title="LBM Object Relighting") as demo:
+    gr.Markdown("# 🌄 Rindriçim i Objektit me Sfondin e Zgjedhur")
 
-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/2503.07535) *by Jasper Research*. This demo is based on the [LBM relighting checkpoint](https://huggingface.co/jasperai/LBM_relighting).
-    """
-    )
-    gr.Markdown(
-        """
-        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,
-                )
+                fg_image = gr.Image(type="pil", label="Imazhi Kryesor", image_mode="RGB", height=360)
+                bg_image = gr.Image(type="pil", label="Sfondi i Ri", image_mode="RGB", height=360)
 
             with gr.Row():
-                submit_button = gr.Button("Relight", variant="primary")
+                submit_button = gr.Button("Rindriço", variant="primary")
             with gr.Row():
-                num_inference_steps = gr.Slider(
-                    minimum=1,
-                    maximum=4,
-                    value=1,
-                    step=1,
-                    label="Number of Inference Steps",
-                )
+                num_inference_steps = gr.Slider(minimum=1, maximum=4, value=1, step=1, label="Hapat e Inferencës")
 
             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,
-            )
+                label="Sfondet", value=[path for path in glob.glob("examples/backgrounds/*.jpg")],
+                columns=5, allow_preview=False)
 
         with gr.Column():
-            output_slider = gr.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="full",
-        show_api=False,
-    )
+            output_slider = gr.ImageSlider(label="Para / Pas", type="numpy", show_share_button=False)
+            output_slider.upload(fn=evaluate, inputs=[fg_image, bg_image, num_inference_steps], outputs=[output_slider])
 
-    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."
-    )
-    gr.Markdown("**Note:** Some backgrounds example are taken from [IC-Light repo](https://github.com/lllyasviel/IC-Light)")
+    submit_button.click(evaluate, inputs=[fg_image, bg_image, num_inference_steps], outputs=[output_slider], show_progress="full", show_api=False)
 
     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(show_api=False)