radames/Text2Human-API

1#!/usr/bin/env python2 3from __future__ import annotations4 5import argparse6import os7import pathlib8import subprocess9 10import gradio as gr11 12if os.getenv('SYSTEM') == 'spaces':13    import mim14 15    mim.uninstall('mmcv-full', confirm_yes=True)16    mim.install('mmcv-full==1.5.2', is_yes=True)17 18    with open('patch') as f:19        subprocess.run('patch -p1'.split(), cwd='Text2Human', stdin=f)20 21from model import Model22 23DESCRIPTION = '''# Text2Human24 25This is an unofficial demo for <a href="https://github.com/yumingj/Text2Human">https://github.com/yumingj/Text2Human</a> made by <a href="https://huggingface.co/spaces/hysts/Text2Human">@hysts</a>.26You can modify sample steps and seeds. By varying seeds, you can sample different human images under the same pose, shape description, and texture description. The larger the sample steps, the better quality of the generated images. (The default value of sample steps is 256 in the original repo.)27 28Label image generation step can be skipped. However, in that case, the input label image must be 512x256 in size and must contain only the specified colors.29'''30FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=hysts.text2human" />'31 32 33def parse_args() -> argparse.Namespace:34    parser = argparse.ArgumentParser()35    parser.add_argument('--device', type=str, default='cpu')36    parser.add_argument('--theme', type=str)37    parser.add_argument('--share', action='store_true')38    parser.add_argument('--port', type=int)39    parser.add_argument('--disable-queue',40                        dest='enable_queue',41                        action='store_false')42    return parser.parse_args()43 44 45def set_example_image(example: list) -> dict:46    return gr.Image.update(value=example[0])47 48 49def set_example_text(example: list) -> dict:50    return gr.Textbox.update(value=example[0])51 52 53def main():54    args = parse_args()55    model = Model(args.device)56 57    with gr.Blocks(theme=args.theme, css='style.css') as demo:58        gr.Markdown(DESCRIPTION)59 60        with gr.Row():61            with gr.Column():62                with gr.Row():63                    input_image = gr.Image(label='Input Pose Image',64                                           type='pil',65                                           elem_id='input-image')66                    pose_data = gr.Variable()67                with gr.Row():68                    paths = sorted(pathlib.Path('pose_images').glob('*.png'))69                    example_images = gr.Dataset(components=[input_image],70                                                samples=[[path.as_posix()]71                                                         for path in paths])72 73                with gr.Row():74                    shape_text = gr.Textbox(75                        label='Shape Description',76                        placeholder=77                        '''<gender>, <sleeve length>, <length of lower clothing>, <outer clothing type>, <other accessories1>, ...78Note: The outer clothing type and accessories can be omitted.''')79                with gr.Row():80                    shape_example_texts = gr.Dataset(81                        components=[shape_text],82                        samples=[['man, sleeveless T-shirt, long pants'],83                                 ['woman, short-sleeve T-shirt, short jeans']])84                with gr.Row():85                    generate_label_button = gr.Button('Generate Label Image')86 87            with gr.Column():88                with gr.Row():89                    label_image = gr.Image(label='Label Image',90                                           type='numpy',91                                           elem_id='label-image')92 93                with gr.Row():94                    texture_text = gr.Textbox(95                        label='Texture Description',96                        placeholder=97                        '''<upper clothing texture>, <lower clothing texture>, <outer clothing texture>98Note: Currently, only 5 types of textures are supported, i.e., pure color, stripe/spline, plaid/lattice, floral, denim.'''99                    )100                with gr.Row():101                    texture_example_texts = gr.Dataset(102                        components=[texture_text],103                        samples=[['pure color, denim'], ['floral, stripe']])104                with gr.Row():105                    sample_steps = gr.Slider(10,106                                             300,107                                             value=10,108                                             step=10,109                                             label='Sample Steps')110                with gr.Row():111                    seed = gr.Slider(0, 1000000, value=0, step=1, label='Seed')112                with gr.Row():113                    generate_human_button = gr.Button('Generate Human')114 115            with gr.Column():116                with gr.Row():117                    result = gr.Image(label='Result',118                                      type='numpy',119                                      elem_id='result-image')120 121        gr.Markdown(FOOTER)122 123        input_image.change(fn=model.process_pose_image,124                           inputs=input_image,125                           outputs=pose_data)126        generate_label_button.click(fn=model.generate_label_image,127                                    inputs=[128                                        pose_data,129                                        shape_text,130                                    ],131                                    outputs=label_image)132        generate_human_button.click(fn=model.generate_human,133                                    inputs=[134                                        label_image,135                                        texture_text,136                                        sample_steps,137                                        seed,138                                    ],139                                    outputs=result)140        example_images.click(fn=set_example_image,141                             inputs=example_images,142                             outputs=example_images.components)143        shape_example_texts.click(fn=set_example_text,144                                  inputs=shape_example_texts,145                                  outputs=shape_example_texts.components)146        texture_example_texts.click(fn=set_example_text,147                                    inputs=texture_example_texts,148                                    outputs=texture_example_texts.components)149 150    demo.launch(151        enable_queue=args.enable_queue,152        server_port=args.port,153        share=args.share,154    )155 156 157if __name__ == '__main__':158    main()159