Spaces:
Runtime error
Runtime error
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 | |
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) | |