tdurzynski commited on
Commit
5661265
·
verified ·
1 Parent(s): d6c355a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -11
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: Determine new dimensions (multiples of 8) based on the foreground.
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 as needed.
55
- guidance_scale=7.5 # Adjust for prompt adherence.
56
  )
57
- # The generated background is in RGB mode; convert to RGBA for compositing.
58
  background = bg_output.images[0].convert("RGBA")
59
 
60
- # Step 4: If necessary, resize the foreground to match the background.
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 foreground over the new background.
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" # You may change the model if desired.
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.inputs.Image(type="pil", label="Upload Your Image"),
102
- gr.inputs.Textbox(lines=2, placeholder="Describe the new background...", label="Background Prompt")
103
  ],
104
- outputs=gr.outputs.Image(type="pil", label="Output Image"),
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"