Inpainting using black-forest-labs/flux-fill-pro

Open In Colab Download .ipynb

!uv pip install gradio==5.49.1 replicate

Set Replicate API Token

import sys
import os
from dotenv import load_dotenv

if 'google.colab' in sys.modules:
  from google.colab import userdata # type:ignore
  os.environ['REPLICATE_API_TOKEN'] = userdata.get('REPLICATE_API_TOKEN')
  print("Replicate API Token set for Colab")
else:
  load_dotenv()
  print("Loaded env vars from .env")

Gradio Interface

import gradio as gr
import numpy as np
from PIL import Image
import replicate
import io

def inpaint(image_data, prompt):
    if image_data is None:
        return None
    
    # Get the original image from the background
    original_image = Image.fromarray(image_data['background'])
    
    # Get the mask from the layers
    if image_data['layers'] and len(image_data['layers']) > 0:
        mask_layer = image_data['layers'][0]
        mask_array = np.array(mask_layer)
        
        # Create binary mask: white where painted, black where not
        alpha_channel = mask_array[:, :, 3]
        binary_mask = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
        mask_image = Image.fromarray(binary_mask, mode='L')
    else:
        return None
    
    # Convert images to bytes for the replicate API
    image_bytes = io.BytesIO()
    original_image.save(image_bytes, format='PNG')
    image_bytes.seek(0)
    
    mask_bytes = io.BytesIO()
    mask_image.save(mask_bytes, format='PNG')
    mask_bytes.seek(0)
    
    # Call the Replicate API
    output = replicate.run(
        "black-forest-labs/flux-fill-pro",
        input={
            "image": image_bytes,
            "mask": mask_bytes,
            "prompt": prompt,
            "steps": 25,
            "guidance": 75,
            "outpaint": "None",
            "output_format": "jpg",
            "safety_tolerance": 2,
            "prompt_upsampling": False
        }
    )
    
    # Read the FileOutput and convert to PIL Image
    output_bytes = output.read()
    output_image = Image.open(io.BytesIO(output_bytes))
    
    return output_image

demo = gr.Interface(
    fn=inpaint,
    inputs=[
        gr.ImageEditor(
            label="Image (paint over areas to inpaint)",
            brush=gr.Brush(color_mode="fixed", colors=["#000000"]),
            layers=True
        ),
        gr.Textbox(label="Prompt", placeholder="Describe what should replace the masked area...")
    ],
    outputs=gr.Image(label="Output Image"),
    title="Inpainting using black-forest-labs/flux-fill-pro"
)
demo.launch(debug=True)
* Running on local URL:  http://127.0.0.1:7861
* To create a public link, set `share=True` in `launch()`.