Replacing Only the Background of an Image with AI Generation Using the Stable Diffusion Web API

Introduction

This guide demonstrates how to replace the background of an image using Python code only, without relying on image editing software like Photoshop. The goal is to keep the subject intact while swapping in an AI-generated background.

While this approach may not be revolutionary, it addresses a common need, so I hope it will be helpful for those with similar requirements.

Input and Output Images

Let’s start with the results.

The following output image was generated from the input image shown below.

Input Image

Output Image

Libraries

Install requests to handle the API calls.

$ pip install requests

I verified the version as follows:

$ pip list | grep -e requests
requests                  2.31.0

API Key

For background generation, we’ll use Stability AI’s Web API.

To access this API, you’ll need to obtain an API Key from their Developer Platform. For pricing, refer to the Pricing page.

To keep your key secure, save it as an environment variable rather than hardcoding it in your code.

In my environment, I use the zshrc settings file.

$ open ~/.zshrc

I saved the key under the name STABILITY_API_KEY.

export STABILITY_API_KEY=your_api_key_here

Code

Here, we use the Remove Background API to isolate the subject. We then pass the extracted image to the Inpaint API to create the new background.

The prompt used is "Large glass windows with a view of the metropolis behind"

import os
import requests

# File paths
input_path = './input.png'  # Original image
mask_path = './mask.png'  # Mask image (temporarily generated)
output_path = './output.png'  # Output image

# Check for API Key
api_key = os.getenv("STABILITY_API_KEY")
if api_key is None:
    raise Exception("Missing Stability API key.")

headers = {
    "Accept": "image/*",
    "Authorization": f"Bearer {api_key}"
}

# Call Remove Background API
response = requests.post(
    f"https://api.stability.ai/v2beta/stable-image/edit/remove-background",
    headers=headers,
    files={
        "image": open(input_path, "rb")
    },
    data={
        "output_format": "png"
    },
)

# Save mask image
if response.status_code == 200:
    with open(mask_path, 'wb') as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

# Call Inpaint API
response = requests.post(
    "https://api.stability.ai/v2beta/stable-image/edit/inpaint",
    headers=headers,
    files={
        "image": open(mask_path, "rb"),
    },
    data={
        "prompt": "Large glass windows with a view of the metropolis behind",
        "output_format": "png",
        "grow_mask": 0,  # Disable blurring around the mask
    },
)

# Delete mask image
os.remove(mask_path)

# Save output image
if response.status_code == 200:
    with open(output_path, "wb") as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

Using rembg

Another approach for background removal is to use rembg. This method requires only one API call, making it more cost-effective, though it may result in differences in extraction accuracy.

First, install rembg.

$ pip install rembg

I verified the version as follows:

$ pip list | grep -e rembg
rembg                     2.0.59

Here’s the code for this approach:

from rembg import remove
import os
import requests

# File paths
input_path = './input.png'  # Input image path
mask_path = './mask.png'  # Mask image path (temporarily generated)
output_path = './output.png'  # Output image path

# Generate mask image with background removed
with open(input_path, 'rb') as i:
    with open(mask_path, 'wb') as o:
        input_image = i.read()
        mask_image = remove(input_image)
        o.write(mask_image)

# Check for API Key
api_key = os.getenv("STABILITY_API_KEY")
if api_key is None:
    raise Exception("Missing Stability API key.")

# Call Inpaint API
response = requests.post(
    "https://api.stability.ai/v2beta/stable-image/edit/inpaint",
    headers={
        "Accept": "image/*",
        "Authorization": f"Bearer {api_key}"
    },
    files={
        "image": open(mask_path, "rb"),
    },
    data={
        "prompt": "Large glass windows with a view of the metropolis behind",
        "output_format": "png",
        "grow_mask": 0,
    },
)

# Delete mask image
os.remove(mask_path)

# Save output image
if response.status_code == 200:
    with open(output_path, "wb") as file:
        file.write(response.content)
else:
    raise Exception(str(response.json()))

Here’s the output image. In this case, the accuracy of the extraction seems satisfactory.

If you set up a local Stable Diffusion environment, you can eliminate API call costs, so feel free to explore that option if it suits your needs.

Conclusion

Being able to achieve this through code alone is highly convenient.

It’s exciting to witness the ongoing improvements in workflow efficiency.

Original Japanese Article

Stable DiffusionのWeb APIを用いて画像上の人物はそのままに背景だけをAI生成で入れ替えてみた

Blog