File size: 2,289 Bytes
0d89801
 
 
 
 
 
031c42b
83cae6c
 
 
0d89801
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import torch
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 = torch.device("cuda")
# DEVICE = torch.device("cpu")

pipe = FluxInpaintPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16)
pipe.to(DEVICE)


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

    image = input_image_editor['background']
    mask_image = input_image_editor['layers'][0]

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

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

    generator = torch.Generator().manual_seed(42)
    return pipe(
        prompt=input_text,
        image=image,
        mask_image=mask_image,
        width=1024,
        height=1024,
        strength=0.9
    ).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=["#000000"], color_mode="fixed"))
            input_text_component = gr.Textbox(
                label='Text prompt',
                placeholder='Cartoon cactus',)
            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,
            input_text_component
        ],
        outputs=[
            output_image_component
        ]
    )

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