File size: 3,235 Bytes
3911a99
1e1a292
 
3911a99
 
 
1e1a292
3911a99
 
1e1a292
3911a99
1e1a292
3911a99
 
 
1e1a292
 
 
 
 
 
 
3911a99
da59864
 
3911a99
 
 
 
 
1e1a292
3911a99
 
 
 
1e1a292
 
3911a99
1e1a292
 
 
3911a99
 
 
 
 
 
a8ce530
 
 
 
1e1a292
 
 
 
 
f58ced3
a8ce530
 
1e1a292
 
 
a8ce530
 
 
1e1a292
a8ce530
 
 
1e1a292
 
 
a8ce530
 
 
 
1e1a292
 
 
 
 
a8ce530
 
 
1e1a292
 
 
 
 
a8ce530
 
1e1a292
a8ce530
 
 
1e1a292
a8ce530
 
 
 
 
 
 
 
 
 
1e1a292
3911a99
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
import gradio as gr
import torch
from diffusers import DDIMScheduler, DiffusionPipeline

stable_inpiant_model_list = [
    "stabilityai/stable-diffusion-2-inpainting",
    "runwayml/stable-diffusion-inpainting",
]

stable_prompt_list = ["a photo of a man.", "a photo of a girl."]

stable_negative_prompt_list = ["bad, ugly", "deformed"]


def stable_diffusion_inpaint(
    dict: str,
    model_path: str,
    prompt: str,
    negative_prompt: str,
    guidance_scale: int,
    num_inference_step: int,
):

    image = dict["image"].convert("RGB").resize((512, 512))
    mask_image = dict["mask"].convert("RGB").resize((512, 512))
    pipe = DiffusionPipeline.from_pretrained(
        model_path,
        revision="fp16",
        torch_dtype=torch.float16,
    )
    pipe.to("cuda")
    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
    pipe.enable_xformers_memory_efficient_attention()

    output = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask_image,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_step,
        guidance_scale=guidance_scale,
    ).images

    return output[0]


def stable_diffusion_inpaint_app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                inpaint_image_file = gr.Image(
                    source="upload",
                    tool="sketch",
                    elem_id="image_upload",
                    type="pil",
                    label="Upload",
                )

                inpaint_model_id = gr.Dropdown(
                    choices=stable_inpiant_model_list,
                    value=stable_inpiant_model_list[0],
                    label="Inpaint Model Id",
                )

                inpaint_prompt = gr.Textbox(
                    lines=1, value=stable_prompt_list[0], label="Prompt"
                )

                inpaint_negative_prompt = gr.Textbox(
                    lines=1,
                    value=stable_negative_prompt_list[0],
                    label="Negative Prompt",
                )

                with gr.Accordion("Advanced Options", open=False):
                    inpaint_guidance_scale = gr.Slider(
                        minimum=0.1,
                        maximum=15,
                        step=0.1,
                        value=7.5,
                        label="Guidance Scale",
                    )

                    inpaint_num_inference_step = gr.Slider(
                        minimum=1,
                        maximum=100,
                        step=1,
                        value=50,
                        label="Num Inference Step",
                    )

                inpaint_predict = gr.Button(value="Generator")

            with gr.Column():
                output_image = gr.Gallery(label="Outputs")

        inpaint_predict.click(
            fn=stable_diffusion_inpaint,
            inputs=[
                inpaint_image_file,
                inpaint_model_id,
                inpaint_prompt,
                inpaint_negative_prompt,
                inpaint_guidance_scale,
                inpaint_num_inference_step,
            ],
            outputs=output_image,
        )