Spaces:
Runtime error
Runtime error
File size: 5,216 Bytes
0c1bd79 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 |
from typing import Tuple, Dict
import requests
import random
import numpy as np
import gradio as gr
import torch
from PIL import Image
from diffusers import StableDiffusionInpaintPipeline
INFO = """
# FLUX-Based Inpainting 🎨
This interface utilizes a FLUX model variant for precise inpainting. Special thanks to the [Black Forest Labs](https://huggingface.co/black-forest-labs) team
and [Gothos](https://github.com/Gothos) for contributing to this advanced solution.
"""
# Constants
MAX_SEED_VALUE = np.iinfo(np.int32).max
TARGET_DIM = 1024
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# Function to clear background
def clear_background(image: Image.Image, threshold: int = 50) -> Image.Image:
image = image.convert("RGBA")
pixels = image.getdata()
processed_data = [
(0, 0, 0, 0) if sum(pixel[:3]) / 3 < threshold else pixel for pixel in pixels
]
image.putdata(processed_data)
return image
# Sample data examples
EXAMPLES = [
[
{
"background": Image.open(requests.get("https://example.com/doge-1.png", stream=True).raw),
"layers": [clear_background(Image.open(requests.get("https://example.com/mask-1.png", stream=True).raw))],
"composite": Image.open(requests.get("https://example.com/composite-1.png", stream=True).raw),
},
"desert mirage",
42,
False,
0.75,
25
],
[
{
"background": Image.open(requests.get("https://example.com/doge-2.png", stream=True).raw),
"layers": [clear_background(Image.open(requests.get("https://example.com/mask-2.png", stream=True).raw))],
"composite": Image.open(requests.get("https://example.com/composite-2.png", stream=True).raw),
},
"neon city",
100,
True,
0.9,
35
]
]
# Load model
inpainting_pipeline = StableDiffusionInpaintPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
# Utility to adjust image size
def get_scaled_dimensions(
original_size: Tuple[int, int], max_dim: int = TARGET_DIM
) -> Tuple[int, int]:
width, height = original_size
scaling_factor = max_dim / max(width, height)
return (int(width * scaling_factor) // 32 * 32, int(height * scaling_factor) // 32 * 32)
@spaces.GPU(duration=100)
def generate_inpainting(
input_data: Dict,
prompt_text: str,
chosen_seed: int,
use_random_seed: bool,
inpainting_strength: float,
steps: int,
progress=gr.Progress(track_tqdm=True)
):
if not prompt_text:
return gr.Info("Provide a prompt to proceed."), None
background = input_data.get("background")
mask_layer = input_data.get("layers")[0]
if not background:
return gr.Info("Background image is missing."), None
if not mask_layer:
return gr.Info("Mask layer is missing."), None
new_width, new_height = get_scaled_dimensions(background.size)
resized_background = background.resize((new_width, new_height), Image.LANCZOS)
resized_mask = mask_layer.resize((new_width, new_height), Image.LANCZOS)
if use_random_seed:
chosen_seed = random.randint(0, MAX_SEED_VALUE)
torch.manual_seed(chosen_seed)
generated_image = inpainting_pipeline(
prompt=prompt_text,
image=resized_background,
mask_image=resized_mask,
strength=inpainting_strength,
num_inference_steps=steps,
).images[0]
return generated_image, resized_mask
# Build the Gradio interface
with gr.Blocks() as flux_app:
gr.Markdown(INFO)
with gr.Row():
with gr.Column():
image_editor = gr.ImageEditor(
label="Edit Image",
type="pil",
sources=["upload", "webcam"],
brush=gr.Brush(colors=["#FFF"], color_mode="fixed")
)
prompt_box = gr.Text(
label="Inpainting Prompt", placeholder="Describe the change you'd like."
)
run_button = gr.Button(value="Run Inpainting")
with gr.Accordion("Settings"):
seed_slider = gr.Slider(0, MAX_SEED_VALUE, step=1, value=42, label="Seed")
random_seed_toggle = gr.Checkbox(label="Randomize Seed", value=True)
inpainting_strength_slider = gr.Slider(0.0, 1.0, step=0.01, value=0.85, label="Inpainting Strength")
steps_slider = gr.Slider(1, 50, step=1, value=25, label="Inference Steps")
with gr.Column():
output_image = gr.Image(label="Output Image")
output_mask = gr.Image(label="Processed Mask")
run_button.click(
generate_inpainting,
inputs=[image_editor, prompt_box, seed_slider, random_seed_toggle, inpainting_strength_slider, steps_slider],
outputs=[output_image, output_mask]
)
gr.Examples(
examples=EXAMPLES,
fn=generate_inpainting,
inputs=[image_editor, prompt_box, seed_slider, random_seed_toggle, inpainting_strength_slider, steps_slider],
outputs=[output_image, output_mask],
run_on_click=True,
)
flux_app.launch(debug=False, show_error=True)
|