File size: 3,789 Bytes
5fbf40f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import spaces
import gradio as gr
import os
import torch
import uuid

from PIL import Image
from enhance_utils import enhance_image


DEFAULT_SRC_PROMPT = "a woman"
DEFAULT_EDIT_PROMPT = "a woman, with red lips, 8k, high quality"

device = "cuda" if torch.cuda.is_available() else "cpu"

def create_demo() -> gr.Blocks:
    from inversion_run_adapter import run as adapter_run

    @spaces.GPU(duration=15)
    def image_to_image(
        input_image_path: str,
        input_image_prompt: str,
        edit_prompt: str,
        seed: int,
        w1: float,
        num_steps: int,
        start_step: int,
        guidance_scale: float,
        generate_size: int,
        lineart_scale: float,
        canny_scale: float,
        lineart_detect: float,
        canny_detect: float,
    ):
        w2 = 1.0

        input_image = Image.open(input_image_path)

        w2 = 1.0
        run_model = adapter_run
        generated_image = run_model(
            input_image,
            input_image_prompt,
            edit_prompt,
            generate_size,
            seed,
            w1,
            w2,
            num_steps,
            start_step,
            guidance_scale,
            lineart_scale,
            canny_scale,
            lineart_detect,
            canny_detect,
        )
        enhanced_image = enhance_image(generated_image, False)

        tmpPrefix = "/tmp/gradio/"

        extension = 'png'
        if enhanced_image.mode == 'RGBA':
            extension = 'png'
        else:
            extension = 'jpg'

        targetDir = f"{tmpPrefix}output/"
        if not os.path.exists(targetDir):
            os.makedirs(targetDir)

        enhanced_path = f"{targetDir}{uuid.uuid4()}.{extension}"
        enhanced_image.save(enhanced_path, quality=100)

        return enhanced_path

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                input_image_path = gr.File(label="Input Image")
            with gr.Column():
                generated_image_path = gr.File(label="Download the segment image", interactive=False)
        with gr.Row():
            with gr.Column():
                input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
                with gr.Accordion("Advanced Options", open=False):
                    guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale")
                    seed = gr.Number(label="Seed", value=8)
                    generate_size = gr.Number(label="Generate Size", value=1024)
                    lineart_scale = gr.Slider(minimum=0, maximum=5, value=0.8, step=0.1, label="Lineart Weights", visible=True)
                    canny_scale = gr.Slider(minimum=0, maximum=5, value=0.4, step=0.1, label="Canny Weights", visible=True)
                    lineart_detect = gr.Number(label="Lineart Detect", value=0.375, visible=True)
                    canny_detect = gr.Number(label="Canny Detect", value=0.375, visible=True)
            with gr.Column():
                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
                start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step")
                w1 = gr.Number(label="W1", value=2)
                g_btn = gr.Button("Edit Image")
                
        
        g_btn.click(
            fn=image_to_image,
            inputs=[input_image_path, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, lineart_scale, canny_scale, lineart_detect, canny_detect],
            outputs=[generated_image_path],
        )

    return demo