Update app.py
Browse files
app.py
CHANGED
@@ -40,7 +40,7 @@ def process_image(input_image: Image.Image, bg_prompt: str) -> Image.Image:
|
|
40 |
foreground = remove(input_image)
|
41 |
foreground = foreground.convert("RGBA")
|
42 |
|
43 |
-
# Step 2:
|
44 |
orig_w, orig_h = foreground.size
|
45 |
gen_w, gen_h = adjust_size(orig_w, orig_h)
|
46 |
print(f"Original size: {orig_w}x{orig_h} | Adjusted size: {gen_w}x{gen_h}")
|
@@ -51,18 +51,18 @@ def process_image(input_image: Image.Image, bg_prompt: str) -> Image.Image:
|
|
51 |
bg_prompt,
|
52 |
height=gen_h,
|
53 |
width=gen_w,
|
54 |
-
num_inference_steps=50, # Adjust
|
55 |
-
guidance_scale=7.5 # Adjust for prompt adherence.
|
56 |
)
|
57 |
-
#
|
58 |
background = bg_output.images[0].convert("RGBA")
|
59 |
|
60 |
-
# Step 4:
|
61 |
if foreground.size != background.size:
|
62 |
print("Resizing foreground to match background dimensions...")
|
63 |
foreground = foreground.resize(background.size, Image.ANTIALIAS)
|
64 |
|
65 |
-
# Step 5: Composite the
|
66 |
print("Compositing images...")
|
67 |
final_image = Image.alpha_composite(background, foreground)
|
68 |
|
@@ -71,7 +71,7 @@ def process_image(input_image: Image.Image, bg_prompt: str) -> Image.Image:
|
|
71 |
# -----------------------------------------------------------------------------
|
72 |
# Load the Stable Diffusion pipeline from Hugging Face.
|
73 |
# -----------------------------------------------------------------------------
|
74 |
-
MODEL_ID = "stabilityai/stable-diffusion-2" #
|
75 |
|
76 |
# Use half precision if GPU is available.
|
77 |
if torch.cuda.is_available():
|
@@ -86,7 +86,7 @@ if torch.cuda.is_available():
|
|
86 |
print("Stable Diffusion pipeline loaded.")
|
87 |
|
88 |
# -----------------------------------------------------------------------------
|
89 |
-
# Create the Gradio Interface.
|
90 |
# -----------------------------------------------------------------------------
|
91 |
title = "Background Removal & Replacement"
|
92 |
description = (
|
@@ -98,10 +98,10 @@ description = (
|
|
98 |
iface = gr.Interface(
|
99 |
fn=process_image,
|
100 |
inputs=[
|
101 |
-
gr.
|
102 |
-
gr.
|
103 |
],
|
104 |
-
outputs=gr.
|
105 |
title=title,
|
106 |
description=description,
|
107 |
allow_flagging="never"
|
|
|
40 |
foreground = remove(input_image)
|
41 |
foreground = foreground.convert("RGBA")
|
42 |
|
43 |
+
# Step 2: Adjust dimensions for background generation.
|
44 |
orig_w, orig_h = foreground.size
|
45 |
gen_w, gen_h = adjust_size(orig_w, orig_h)
|
46 |
print(f"Original size: {orig_w}x{orig_h} | Adjusted size: {gen_w}x{gen_h}")
|
|
|
51 |
bg_prompt,
|
52 |
height=gen_h,
|
53 |
width=gen_w,
|
54 |
+
num_inference_steps=50, # Adjust if needed.
|
55 |
+
guidance_scale=7.5 # Adjust for more/less prompt adherence.
|
56 |
)
|
57 |
+
# Convert the generated background to RGBA.
|
58 |
background = bg_output.images[0].convert("RGBA")
|
59 |
|
60 |
+
# Step 4: Ensure the foreground matches the background dimensions.
|
61 |
if foreground.size != background.size:
|
62 |
print("Resizing foreground to match background dimensions...")
|
63 |
foreground = foreground.resize(background.size, Image.ANTIALIAS)
|
64 |
|
65 |
+
# Step 5: Composite the images.
|
66 |
print("Compositing images...")
|
67 |
final_image = Image.alpha_composite(background, foreground)
|
68 |
|
|
|
71 |
# -----------------------------------------------------------------------------
|
72 |
# Load the Stable Diffusion pipeline from Hugging Face.
|
73 |
# -----------------------------------------------------------------------------
|
74 |
+
MODEL_ID = "stabilityai/stable-diffusion-2" # Change the model if desired.
|
75 |
|
76 |
# Use half precision if GPU is available.
|
77 |
if torch.cuda.is_available():
|
|
|
86 |
print("Stable Diffusion pipeline loaded.")
|
87 |
|
88 |
# -----------------------------------------------------------------------------
|
89 |
+
# Create the Gradio Interface using the updated API.
|
90 |
# -----------------------------------------------------------------------------
|
91 |
title = "Background Removal & Replacement"
|
92 |
description = (
|
|
|
98 |
iface = gr.Interface(
|
99 |
fn=process_image,
|
100 |
inputs=[
|
101 |
+
gr.Image(type="pil", label="Upload Your Image"),
|
102 |
+
gr.Textbox(lines=2, placeholder="Describe the new background...", label="Background Prompt")
|
103 |
],
|
104 |
+
outputs=gr.Image(label="Output Image"),
|
105 |
title=title,
|
106 |
description=description,
|
107 |
allow_flagging="never"
|