File size: 3,002 Bytes
0d89801
03bb3c4
0d89801
 
 
 
 
83cae6c
 
 
0d89801
 
f0e8d1f
0d89801
 
f0e8d1f
0d89801
 
3e075bb
c2f8847
0d89801
 
 
 
 
c2f8847
 
 
 
0d89801
 
 
 
 
 
c2f8847
0d89801
 
867296e
 
0d89801
 
 
 
867296e
 
bc5af2e
0d89801
 
 
 
 
 
 
 
 
 
 
 
 
3e075bb
 
 
 
 
 
 
 
bc5af2e
 
 
 
 
 
 
c2f8847
 
0d89801
 
 
 
 
 
 
 
 
 
c2f8847
bc5af2e
 
0d89801
 
 
 
 
 
bc5af2e
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
import torch
import spaces
import gradio as gr
from diffusers import FluxInpaintPipeline

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.
"""

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)


@spaces.GPU()
def process(input_image_editor, uploaded_mask, input_text, strength, progress=gr.Progress(track_tqdm=True)):
    if not input_text:
        gr.Info("Please enter a text prompt.")
        return None

    image = input_image_editor['background']
    if uploaded_mask is None:
        mask_image = input_image_editor['layers'][0]
    else:
        mask_image = uploaded_mask

    if not image:
        gr.Info("Please upload an image.")
        return None

    if not mask_image:
        gr.Info("Please draw or upload a mask on the image.")
        return None

    width, height = image.size

    return pipe(
        prompt=input_text,
        image=image,
        mask_image=mask_image,
        width=width,
        height=height,
        strength=strength
    ).images[0]


with gr.Blocks() as demo:
    gr.Markdown(MARKDOWN)
    with gr.Row():
        with gr.Column():
            input_image_editor_component = gr.ImageEditor(
                label='Image',
                type='pil',
                sources=["upload", "webcam"],
                image_mode='RGB',
                layers=False,
                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"))
            input_text_component = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            strength_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.7,
                step=0.01,
                label="Strength"
            )
            with gr.Accordion("Upload a mask", open = False):
                uploaded_mask_component = gr.Image(label = "Already made mask (black pixels will be preserved, white pixels will be redrawn)", sources = ["upload"], type = "pil")
            submit_button_component = gr.Button(
                value='Submit', variant='primary')
        with gr.Column():
            output_image_component = gr.Image(
                type='pil', image_mode='RGB', label='Generated image')

    submit_button_component.click(
        fn=process,
        inputs=[
            input_image_editor_component,
            uploaded_mask_component,
            input_text_component,
            strength_slider
        ],
        outputs=[
            output_image_component
        ]
    )

demo.launch(debug=False, show_error=True)