Update app.py

app.py

@@ -5,14 +5,15 @@ 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")
@@ -31,20 +32,52 @@ 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()
-image_size = (1024, 1024)
+
+# Load Stable Diffusion pipeline for background generation
+sd_pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+    revision="fp16",
+    use_auth_token=huggingface_token,
+)
+sd_pipe.to("cuda")
+sd_pipe.enable_attention_slicing()
 
 @spaces.GPU
-def evaluate(fg_image: PIL.Image.Image, bg_image: PIL.Image.Image, num_sampling_steps: int = 4):
+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
+
     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))
@@ -82,41 +115,30 @@ button[aria-label="Download"] {
     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"] {
+button[aria-label="Share"], button[aria-label="Copy link"], 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")
+
+    gr.Markdown("# Rindriçim i Objektit me Sfondin e Zgjedhur")
 
     with gr.Row():
         with gr.Column():
             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)
 
-            bg_gallery = gr.Gallery(object_fit="contain", visible=False)
+            num_inference_steps = gr.Slider(minimum=1, maximum=4, value=4, step=1, visible=False)
 
         with gr.Column():
             output_slider = gr.ImageSlider(label="Para / Pas", 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)
-
-    def bg_gallery_selected(gal, evt: gr.SelectData):
-        return gal[evt.index][0]
+            output_slider.upload(fn=evaluate, inputs=[fg_image, bg_image, bg_prompt, num_inference_steps], outputs=[output_slider])
 
-    bg_gallery.select(bg_gallery_selected, inputs=bg_gallery, outputs=bg_image)
+    submit_button.click(evaluate, inputs=[fg_image, bg_image, bg_prompt, num_inference_steps], outputs=[output_slider], show_progress="full", show_api=False)
 
 if __name__ == "__main__":
     demo.queue().launch(show_api=False)