ysharma's picture
ysharma HF staff
Update app.py
aa440b5 verified
from typing import Tuple
import uuid
import random
import os
import numpy as np
import gradio as gr
import spaces
import torch
from PIL import Image
from diffusers import FluxInpaintPipeline
from gradio_client import Client, handle_file
# Set an environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)
MARKDOWN = """
# FLUX.1 Inpainting with Text guided Mask🔥
Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) team for FLUX!
Special thanks to [Piotr Skalski](https://huggingface.co/SkalskiP) and [Gothos](https://github.com/Gothos)
for their work on enabling and [showcasing inpainting](https://huggingface.co/spaces/SkalskiP/FLUX.1-inpaint) with the FLUX.
<br>We have used Gradio clients to access [EVF-SAM Spaces demo](https://huggingface.co/spaces/wondervictor/evf-sam) for text-guided segmentation.
Visit [Gradio docs](https://www.gradio.app) to start building!
"""
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Using Gradio Python Client to query EVF-SAM demo, hosted on SPaces, as an endpoint
client = Client("ysharma/evf-sam", hf_token=HF_TOKEN)
pipe = FluxInpaintPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
def resize_image_dimensions(
original_resolution_wh: Tuple[int, int],
maximum_dimension: int = 2048
) -> Tuple[int, int]:
width, height = original_resolution_wh
if width <= maximum_dimension and height <= maximum_dimension:
width = width - (width % 32)
height = height - (height % 32)
return width, height
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 evf_sam_mask(image, prompt):
images = client.predict(
image_np=handle_file(image),
prompt=prompt,
api_name="/predict")
# Open the mask image
pil_image = Image.open(images[1])
return pil_image
@spaces.GPU(duration=150)
def process(
input_image: dict,
input_text: str,
inpaint_text: str,
seed_slicer: int,
randomize_seed_checkbox: bool,
strength_slider: float,
num_inference_steps_slider: int,
progress=gr.Progress(track_tqdm=True)
):
if not input_text:
gr.Info("Please enter a text prompt.")
return None
mask = evf_sam_mask(input_image, input_text)
if not input_image:
gr.Info("Please upload an image.")
return None
else:
input_image = Image.open(input_image)
if not mask:
gr.Info("Please draw a mask on the image.")
return None
width, height = resize_image_dimensions(original_resolution_wh=input_image.size)
resized_image = input_image.resize((width, height), Image.LANCZOS)
resized_mask = mask.resize((width, height), Image.NEAREST)
if randomize_seed_checkbox:
seed_slicer = random.randint(0, MAX_SEED)
generator = torch.Generator().manual_seed(seed_slicer)
result = pipe(
prompt=inpaint_text,
image=resized_image,
mask_image=resized_mask,
width=width,
height=height,
strength=strength_slider,
generator=generator,
num_inference_steps=num_inference_steps_slider
).images[0]
print('INFERENCE DONE')
return result, resized_mask
with gr.Blocks() as demo:
gr.Markdown(MARKDOWN)
with gr.Row():
with gr.Column():
input_image = gr.Image(
label='Image',
type='filepath',
sources=["upload", "webcam", "clipboard"],
image_mode='RGB',
)
with gr.Row():
with gr.Column():
input_text_component = gr.Text(
label="Text-guided segmentation",
show_label=True,
max_lines=1,
placeholder="Enter text for generating the segmentation mask",
container=False,
)
inpaint_text_component = gr.Text(
label="Text-guided Inpainting",
show_label=True,
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):
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=False)
with gr.Row():
strength_slider_component = gr.Slider(
label="Strength",
minimum=0,
maximum=1,
step=0.01,
value=0.75,
)
num_inference_steps_slider_component = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=50,
step=1,
value=20,
)
with gr.Column():
output_image_component = gr.Image(
type='pil', image_mode='RGB', label='Generated image')
with gr.Accordion("Generated Mask", open=False):
output_mask_component = gr.Image(
type='pil', image_mode='RGB', label='Input mask')
submit_button_component.click(
fn=process,
inputs=[
input_image, #input_image_editor_component,
input_text_component,
inpaint_text_component,
seed_slicer_component,
randomize_seed_checkbox_component,
strength_slider_component,
num_inference_steps_slider_component
],
outputs=[
output_image_component,
output_mask_component,
]
)
demo.launch(debug=True)