File size: 7,919 Bytes
107d5d9
61c2d32
 
 
80ccb59
 
7aefd5a
61c2d32
7aefd5a
61c2d32
80ccb59
7aefd5a
 
61c2d32
80ccb59
 
7aefd5a
 
61c2d32
 
 
 
7aefd5a
61c2d32
 
 
 
 
983b029
 
61c2d32
 
4405af8
61c2d32
7aefd5a
 
 
 
61c2d32
 
 
 
80ccb59
983b029
80ccb59
 
 
 
 
983b029
80ccb59
 
 
 
7aefd5a
 
80ccb59
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7aefd5a
80ccb59
 
 
 
7aefd5a
80ccb59
7aefd5a
80ccb59
 
7aefd5a
80ccb59
61c2d32
80ccb59
 
 
 
7aefd5a
107d5d9
7aefd5a
 
80ccb59
 
61c2d32
 
 
80ccb59
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7aefd5a
 
 
 
 
 
 
 
 
 
80ccb59
 
 
7aefd5a
 
 
 
 
 
80ccb59
 
7aefd5a
80ccb59
 
 
 
 
 
 
 
 
 
61c2d32
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
80ccb59
61c2d32
 
 
80ccb59
61c2d32
80ccb59
61c2d32
 
 
80ccb59
61c2d32
983b029
 
61c2d32
 
 
80ccb59
4405af8
80ccb59
 
61c2d32
80ccb59
61c2d32
 
4405af8
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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
# import spaces
import os
import sys
import time
from glob import glob
from os.path import join as opj
from pathlib import Path

import apply_net
import gradio as gr
import torch
from omegaconf import OmegaConf
from PIL import Image

from cldm.model import create_model
from cldm.plms_hacked import PLMSSampler
from detectron2.data.detection_utils import _apply_exif_orientation, convert_PIL_to_numpy
from utils_stableviton import get_batch, get_mask_location, tensor2img

PROJECT_ROOT = Path(__file__).absolute().parents[1].absolute()
sys.path.insert(0, str(PROJECT_ROOT))

# from detectron2.projects.DensePose.apply_net_gradio import DensePose4Gradio
from preprocess.humanparsing.run_parsing import Parsing
from preprocess.openpose.run_openpose import OpenPose

os.environ['GRADIO_TEMP_DIR'] = './tmp'  # TODO: turn off when final upload

IMG_H = 512
IMG_W = 384

openpose_model_hd = OpenPose(0)
openpose_model_hd.preprocessor.body_estimation.model.to('cuda')
parsing_model_hd = Parsing(0)
# densepose_model_hd = DensePose4Gradio(
#     cfg='preprocess/detectron2/projects/DensePose/configs/densepose_rcnn_R_50_FPN_s1x.yaml',
#     model='https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl',
# )

category_dict = ['upperbody', 'lowerbody', 'dress']
category_dict_utils = ['upper_body', 'lower_body', 'dresses']

# #### model init >>>>
config = OmegaConf.load("./configs/VITON.yaml")
config.model.params.img_H = IMG_H
config.model.params.img_W = IMG_W
params = config.model.params

model = create_model(config_path=None, config=config)
model.load_state_dict(torch.load("./checkpoints/VITONHD.ckpt", map_location="cpu")["state_dict"])
model = model.cuda()
model.eval()
sampler = PLMSSampler(model)
# #### model init <<<<


def stable_viton_model_hd(
        batch,
        n_steps,
):
    z, cond = model.get_input(batch, params.first_stage_key)
    bs = z.shape[0]
    c_crossattn = cond["c_crossattn"][0][:bs]
    if c_crossattn.ndim == 4:
        c_crossattn = model.get_learned_conditioning(c_crossattn)
        cond["c_crossattn"] = [c_crossattn]
    uc_cross = model.get_unconditional_conditioning(bs)
    uc_full = {"c_concat": cond["c_concat"], "c_crossattn": [uc_cross]}
    uc_full["first_stage_cond"] = cond["first_stage_cond"]
    for k, v in batch.items():
        if isinstance(v, torch.Tensor):
            batch[k] = v.cuda()
    sampler.model.batch = batch

    ts = torch.full((1,), 999, device=z.device, dtype=torch.long)
    start_code = model.q_sample(z, ts)

    output, _, _ = sampler.sample(
        n_steps,
        bs,
        (4, IMG_H // 8, IMG_W // 8),
        cond,
        x_T=start_code,
        verbose=False,
        eta=0.0,
        unconditional_conditioning=uc_full,
    )

    output = model.decode_first_stage(output)
    output = tensor2img(output)
    pil_output = Image.fromarray(output)
    return pil_output

# @spaces.GPU  # TODO: turn on when final upload


@torch.no_grad()
def process_hd(vton_img, garm_img, n_steps):
    model_type = 'hd'
    category = 0  # 0:upperbody; 1:lowerbody; 2:dress

    stt = time.time()
    print('load images... ', end='')
    garm_img = Image.open(garm_img).resize((IMG_W, IMG_H))
    vton_img = Image.open(vton_img).resize((IMG_W, IMG_H))
    print('%.2fs' % (time.time() - stt))

    stt = time.time()
    print('get agnostic map... ', end='')
    keypoints = openpose_model_hd(vton_img.resize((IMG_W, IMG_H)))
    model_parse, _ = parsing_model_hd(vton_img.resize((IMG_W, IMG_H)))
    mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints)
    mask = mask.resize((IMG_W, IMG_H), Image.NEAREST)
    mask_gray = mask_gray.resize((IMG_W, IMG_H), Image.NEAREST)
    masked_vton_img = Image.composite(mask_gray, vton_img, mask)  # agnostic map
    print('%.2fs' % (time.time() - stt))

    stt = time.time()
    print('get densepose... ', end='')
    vton_img = vton_img.resize((IMG_W, IMG_H))  # size for densepose
    # densepose = densepose_model_hd.execute(vton_img)  # densepose

    human_img_arg = _apply_exif_orientation(vton_img.resize((IMG_W, IMG_H)))
    human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")
    args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
    # verbosity = getattr(args, "verbosity", None)
    pose_img = args.func(args, human_img_arg)
    pose_img = pose_img[:, :, ::-1]
    pose_img = Image.fromarray(pose_img).resize((IMG_W, IMG_H))

    print('%.2fs' % (time.time() - stt))

    batch = get_batch(
        vton_img,
        garm_img,
        densepose,
        masked_vton_img,
        mask,
        IMG_H,
        IMG_W
    )

    sample = stable_viton_model_hd(
        batch,
        n_steps
    )
    return sample


example_path = opj(os.path.dirname(__file__), 'examples')
example_model_ps = sorted(glob(opj(example_path, "model/*")))
example_garment_ps = sorted(glob(opj(example_path, "garment/*")))

with gr.Blocks(css='style.css') as demo:
    gr.HTML(
        """
        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
            <div>
                <h1>StableVITON Demo πŸ‘•πŸ‘”πŸ‘—</h1>
                <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
                    <a href='https://arxiv.org/abs/2312.01725'>
                        <img src="https://img.shields.io/badge/arXiv-2312.01725-red">
                    </a>
                    &nbsp;
                    <a href='https://rlawjdghek.github.io/StableVITON/'>
                        <img src='https://img.shields.io/badge/page-github.io-blue.svg'>
                    </a>
                    &nbsp;
                    <a href='https://github.com/rlawjdghek/StableVITON'>
                        <img src='https://img.shields.io/github/stars/rlawjdghek/StableVITON'>
                    </a>
                    &nbsp;
                    <a href='https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode'>
                        <img src='https://img.shields.io/badge/license-CC_BY--NC--SA_4.0-lightgrey'>
                    </a>
                </div>
            </div>
        </div>
        """
    )
    with gr.Row():
        gr.Markdown("## Experience virtual try-on with your own images!")
    with gr.Row():
        with gr.Column():
            vton_img = gr.Image(label="Model", type="filepath", height=384, value=example_model_ps[0])
            example = gr.Examples(
                inputs=vton_img,
                examples_per_page=14,
                examples=example_model_ps)
        with gr.Column():
            garm_img = gr.Image(label="Garment", type="filepath", height=384, value=example_garment_ps[0])
            example = gr.Examples(
                inputs=garm_img,
                examples_per_page=14,
                examples=example_garment_ps)
        with gr.Column():
            result_gallery = gr.Image(label='Output', show_label=False, preview=True, scale=1)
            # result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True, scale=1)
    with gr.Column():
        run_button = gr.Button(value="Run")
        # TODO: change default values (important!)
        # n_samples = gr.Slider(label="Images", minimum=1, maximum=4, value=1, step=1)
        n_steps = gr.Slider(label="Steps", minimum=20, maximum=70, value=25, step=1)
        # guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=5.0, value=2.0, step=0.1)
        # seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=-1)

    ips = [vton_img, garm_img, n_steps]
    run_button.click(fn=process_hd, inputs=ips, outputs=[result_gallery])

demo.queue().launch(share=True)