Update app.py

app.py

@@ -5,15 +5,14 @@ from copy import deepcopy
 import gradio as gr
 import numpy as np
 import PIL
-from PIL import Image, ImageFilter
 import spaces
 import torch
 import yaml
 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 diffusers import StableDiffusionPipeline
 from utils import extract_object, get_model_from_config, resize_and_center_crop
 
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
@@ -32,51 +31,20 @@ ASPECT_RATIOS = {
     str(1920 / 512): (1920, 512),
 }
 
-# Load relighting 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)
+
 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)
-
-# Load segmentation model
 birefnet = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True).cuda()
-
-# Load Stable Diffusion pipeline for background generation
-sd_pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16,
-    use_auth_token=huggingface_token,
-)
-sd_pipe.to("cuda")
-sd_pipe.enable_attention_slicing()
+image_size = (1024, 1024)
 
 @spaces.GPU
-def generate_background_image(bg_prompt: str):
-    if not bg_prompt or bg_prompt.strip() == "":
-        return None
-    with torch.inference_mode():
-        bg_img = sd_pipe(prompt=bg_prompt, height=1024, width=1024, num_inference_steps=20).images[0]
-    # Optional blur radius — tweak as you like or expose as a parameter
-    bg_img = bg_img.filter(ImageFilter.GaussianBlur(radius=5))
-    return bg_img
-
-@spaces.GPU
-def evaluate(
-    fg_image: PIL.Image.Image,
-    bg_image: PIL.Image.Image,
-    bg_prompt: str,
-    num_sampling_steps: int = 4,
-):
-    # Generate background if prompt is given
-    if bg_prompt and bg_prompt.strip() != "":
-        generated_bg = generate_background_image(bg_prompt)
-        if generated_bg is not None:
-            bg_image = generated_bg
-
+def evaluate(fg_image: PIL.Image.Image, bg_image: PIL.Image.Image, num_sampling_steps: int = 4):
     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))
@@ -114,11 +82,16 @@ button[aria-label="Download"] {
     margin: 0 !important;
     padding: 6px !important;
 }
-button[aria-label="Share"], button[aria-label="Copy link"], button[aria-label="Open in new tab"] {
+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.Row():
@@ -126,18 +99,24 @@ button[aria-label="Share"], button[aria-label="Copy link"], button[aria-label="O
             with gr.Row():
                 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)
-            bg_prompt = gr.Textbox(label="Sfondi (p.sh. 'në Milano')", placeholder="Shkruani një përshkrim për sfondin", lines=1)
 
             with gr.Row():
                 submit_button = gr.Button("Rindriço", variant="primary")
+            with gr.Row():
+                num_inference_steps = gr.Slider(minimum=1, maximum=4, value=4, step=1, visible=False)
 
-            num_inference_steps = gr.Slider(minimum=1, maximum=4, value=4, step=1, visible=False)
+            bg_gallery = gr.Gallery(object_fit="contain", visible=False)
 
         with gr.Column():
             output_slider = gr.ImageSlider(label="Para / Pas", type="numpy")
-            output_slider.upload(fn=evaluate, inputs=[fg_image, bg_image, bg_prompt, num_inference_steps], outputs=[output_slider])
+            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)
+
+    def bg_gallery_selected(gal, evt: gr.SelectData):
+        return gal[evt.index][0]
 
-    submit_button.click(evaluate, inputs=[fg_image, bg_image, bg_prompt, num_inference_steps], outputs=[output_slider], show_progress="full", show_api=False)
+    bg_gallery.select(bg_gallery_selected, inputs=bg_gallery, outputs=bg_image)
 
 if __name__ == "__main__":
     demo.queue().launch(show_api=False)