from functools import partial import cv2 import random from typing import Tuple, Optional import gradio as gr import numpy as np import requests import spaces import torch from PIL import Image, ImageFilter from diffusers import FluxInpaintPipeline from gradio_client import Client, handle_file MARKDOWN = """ # FLUX.1 Inpainting 🔥 Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) team for creating this amazing model, and a big thanks to [Gothos](https://github.com/Gothos) for taking it to the next level by enabling inpainting with the FLUX. """ MAX_SEED = np.iinfo(np.int32).max IMAGE_SIZE = 1024 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" PIPE = FluxInpaintPipeline.from_pretrained( "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE) CLIENT = Client("SkalskiP/florence-sam-masking") EXAMPLES = [ [ { "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-image.png", stream=True).raw), "layers": [Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-mask-2-removebg.png", stream=True).raw)], "composite": Image.open(requests.get("https://media.roboflow.com/spaces/doge-2-composite-2.png", stream=True).raw), }, "little lion", "", 5, 5, 42, False, 0.85, 20 ], [ { "background": Image.open(requests.get("https://media.roboflow.com/spaces/doge-5.jpeg", stream=True).raw), "layers": None, "composite": None }, "big blue eyes", "eyes", 10, 5, 42, False, 0.9, 20 ] ] def calculate_image_dimensions_for_flux( original_resolution_wh: Tuple[int, int], maximum_dimension: int = IMAGE_SIZE ) -> Tuple[int, int]: width, height = original_resolution_wh if width > height: scaling_factor = maximum_dimension / width else: scaling_factor = maximum_dimension / height new_width = int(width * scaling_factor) new_height = int(height * scaling_factor) new_width = new_width - (new_width % 32) new_height = new_height - (new_height % 32) return new_width, new_height def is_mask_empty(image: Image.Image) -> bool: gray_img = image.convert("L") pixels = list(gray_img.getdata()) return all(pixel == 0 for pixel in pixels) def process_mask( mask: Image.Image, mask_inflation: Optional[int] = None, mask_blur: Optional[int] = None ) -> Image.Image: """ Inflates and blurs the white regions of a mask. Args: mask (Image.Image): The input mask image. mask_inflation (Optional[int]): The number of pixels to inflate the mask by. mask_blur (Optional[int]): The radius of the Gaussian blur to apply. Returns: Image.Image: The processed mask with inflated and/or blurred regions. """ if mask_inflation and mask_inflation > 0: mask_array = np.array(mask) kernel = np.ones((mask_inflation, mask_inflation), np.uint8) mask_array = cv2.dilate(mask_array, kernel, iterations=1) mask = Image.fromarray(mask_array) if mask_blur and mask_blur > 0: mask = mask.filter(ImageFilter.GaussianBlur(radius=mask_blur)) return mask def set_client_for_session(request: gr.Request): try: x_ip_token = request.headers['x-ip-token'] return Client("SkalskiP/florence-sam-masking", headers={"X-IP-Token": x_ip_token}) except: return CLIENT @spaces.GPU(duration=50) def run_flux( image: Image.Image, mask: Image.Image, prompt: str, seed_slicer: int, randomize_seed_checkbox: bool, strength_slider: float, num_inference_steps_slider: int, resolution_wh: Tuple[int, int], ) -> Image.Image: print("Running FLUX...") width, height = resolution_wh if randomize_seed_checkbox: seed_slicer = random.randint(0, MAX_SEED) generator = torch.Generator().manual_seed(seed_slicer) return PIPE( prompt=prompt, image=image, mask_image=mask, width=width, height=height, strength=strength_slider, generator=generator, num_inference_steps=num_inference_steps_slider ).images[0] def process( client, input_image_editor: dict, inpainting_prompt_text: str, masking_prompt_text: str, mask_inflation_slider: int, mask_blur_slider: int, seed_slicer: int, randomize_seed_checkbox: bool, strength_slider: float, num_inference_steps_slider: int ): if not inpainting_prompt_text: gr.Info("Please enter inpainting text prompt.") return None, None image_path = input_image_editor['background'] mask_path = input_image_editor['layers'][0] image = Image.open(image_path) mask = Image.open(mask_path) if not image: gr.Info("Please upload an image.") return None, None if is_mask_empty(mask) and not masking_prompt_text: gr.Info("Please draw a mask or enter a masking prompt.") return None, None if not is_mask_empty(mask) and masking_prompt_text: gr.Info("Both mask and masking prompt are provided. Please provide only one.") return None, None if is_mask_empty(mask): print("Generating mask...") mask = client.predict( image_input=handle_file(image_path), text_input=masking_prompt_text, api_name="/process_image") mask = Image.open(mask) print("Mask generated.") width, height = calculate_image_dimensions_for_flux(original_resolution_wh=image.size) image = image.resize((width, height), Image.LANCZOS) mask = mask.resize((width, height), Image.LANCZOS) mask = process_mask(mask, mask_inflation=mask_inflation_slider, mask_blur=mask_blur_slider) image = run_flux( image=image, mask=mask, prompt=inpainting_prompt_text, seed_slicer=seed_slicer, randomize_seed_checkbox=randomize_seed_checkbox, strength_slider=strength_slider, num_inference_steps_slider=num_inference_steps_slider, resolution_wh=(width, height) ) return image, mask process_example = partial(process, client=CLIENT) with gr.Blocks() as demo: client_component = gr.State() gr.Markdown(MARKDOWN) with gr.Row(): with gr.Column(): input_image_editor_component = gr.ImageEditor( label='Image', type='filepath', sources=["upload", "webcam"], image_mode='RGB', layers=False, brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed")) with gr.Row(): inpainting_prompt_text_component = gr.Text( label="Inpainting prompt", show_label=False, max_lines=1, placeholder="Enter text to generate inpainting", container=False, ) submit_button_component = gr.Button( value='Submit', variant='primary', scale=0) with gr.Accordion("Advanced Settings", open=False): masking_prompt_text_component = gr.Text( label="Masking prompt", show_label=False, max_lines=1, placeholder="Enter text to generate masking", container=False, ) with gr.Row(): mask_inflation_slider_component = gr.Slider( label="Mask inflation", info="Adjusts the amount of mask edge expansion before " "inpainting.", minimum=0, maximum=20, step=1, value=5, ) mask_blur_slider_component = gr.Slider( label="Mask blur", info="Controls the intensity of the Gaussian blur applied to " "the mask edges.", minimum=0, maximum=20, step=1, value=5, ) seed_slicer_component = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, ) randomize_seed_checkbox_component = gr.Checkbox( label="Randomize seed", value=True) with gr.Row(): strength_slider_component = gr.Slider( label="Strength", info="Indicates extent to transform the reference `image`. " "Must be between 0 and 1. `image` is used as a starting " "point and more noise is added the higher the `strength`.", minimum=0, maximum=1, step=0.01, value=0.85, ) num_inference_steps_slider_component = gr.Slider( label="Number of inference steps", info="The number of denoising steps. More denoising steps " "usually lead to a higher quality image at the", minimum=1, maximum=50, step=1, value=20, ) with gr.Column(): output_image_component = gr.Image( type='pil', image_mode='RGB', label='Generated image', format="png") with gr.Accordion("Debug", open=False): output_mask_component = gr.Image( type='pil', image_mode='RGB', label='Input mask', format="png") gr.Examples( fn=process_example, examples=EXAMPLES, inputs=[ input_image_editor_component, inpainting_prompt_text_component, masking_prompt_text_component, mask_inflation_slider_component, mask_blur_slider_component, seed_slicer_component, randomize_seed_checkbox_component, strength_slider_component, num_inference_steps_slider_component ], outputs=[ output_image_component, output_mask_component ], run_on_click=False ) submit_button_component.click( fn=process, inputs=[ client_component, input_image_editor_component, inpainting_prompt_text_component, masking_prompt_text_component, mask_inflation_slider_component, mask_blur_slider_component, seed_slicer_component, randomize_seed_checkbox_component, strength_slider_component, num_inference_steps_slider_component ], outputs=[ output_image_component, output_mask_component ] ) demo.load(set_client_for_session, None, client_component) demo.launch(debug=False, show_error=True)