Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import legacy | |
| import dnnlib | |
| import numpy as np | |
| import torch | |
| from find_direction import find_direction | |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| with dnnlib.util.open_url("./pretrained/ffhq.pkl") as f: | |
| G = legacy.load_network_pkl(f)['G_ema'].to(device) | |
| DESCRIPTION = '''# <a href="https://github.com/catlab-team/stylemc"> StyleMC:</a> Multi-Channel Based Fast Text-Guided Image Generation and Manipulation | |
| ''' | |
| FOOTER = 'This space is built by <a href = "https://github.com/catlab-team">Catlab Team</a>.' | |
| def main(): | |
| with gr.Blocks(css='style.css') as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Box(): | |
| gr.Markdown('''## Step 1 (Finding a global manipulation direction) | |
| - Please enter the target **text prompt** and **identity loss weight** to find global manipulation direction: | |
| - Hit the **Find Direction** button. | |
| ''') | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| text = gr.Textbox( | |
| label="Enter your prompt", | |
| show_label=False, | |
| max_lines=1, | |
| placeholder="Enter your prompt", | |
| ).style( | |
| container=False, | |
| ) | |
| identity_loss_weight = gr.Slider(0.1, | |
| 10, | |
| value=0.5, | |
| step=0.1, | |
| label='Identity Loss Weight', | |
| interactive=True) | |
| btn = gr.Button("Find Direction").style(full_width=False) | |
| with gr.Box(): | |
| gr.Markdown('''## Step 2 (Manipulation) | |
| - Please upload an image for manipulation: | |
| - You can also select the **previous directions** and determine the **manipulation strength**. | |
| - Hit the **Generate** button. | |
| ''') | |
| with gr.Row(): | |
| identity_loss_weight = gr.Slider(0.1, | |
| 100, | |
| value=50, | |
| step=0.1, | |
| label='Manipulation Strength', | |
| interactive=True) | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| input_image = gr.Image(label='Input Image', | |
| type='filepath') | |
| with gr.Row(): | |
| generate_button = gr.Button('Generate') | |
| with gr.Column(): | |
| with gr.Row(): | |
| generated_image = gr.Image(label='Generated Image', | |
| type='numpy', | |
| interactive=False) | |
| # with gr.Box(): | |
| # gr.Markdown('''## Step 2 (Select Style Image) | |
| # - Select **Style Type**. | |
| # - Select **Style Image Index** from the image table below. | |
| # ''') | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # style_type = gr.Radio(model.style_types, | |
| # label='Style Type') | |
| # text = get_style_image_markdown_text('cartoon') | |
| # style_image = gr.Markdown(value=text) | |
| # style_index = gr.Slider(0, | |
| # 316, | |
| # value=26, | |
| # step=1, | |
| # label='Style Image Index') | |
| # with gr.Row(): | |
| # example_styles = gr.Dataset( | |
| # components=[style_type, style_index], | |
| # samples=[ | |
| # ['cartoon', 26], | |
| # ['caricature', 65], | |
| # ['arcane', 63], | |
| # ['pixar', 80], | |
| # ]) | |
| # with gr.Box(): | |
| # gr.Markdown('''## Step 3 (Generate Style Transferred Image) | |
| # - Adjust **Structure Weight** and **Color Weight**. | |
| # - These are weights for the style image, so the larger the value, the closer the resulting image will be to the style image. | |
| # - Hit the **Generate** button. | |
| # ''') | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # with gr.Row(): | |
| # structure_weight = gr.Slider(0, | |
| # 1, | |
| # value=0.6, | |
| # step=0.1, | |
| # label='Structure Weight') | |
| # with gr.Row(): | |
| # color_weight = gr.Slider(0, | |
| # 1, | |
| # value=1, | |
| # step=0.1, | |
| # label='Color Weight') | |
| # with gr.Row(): | |
| # structure_only = gr.Checkbox(label='Structure Only') | |
| # with gr.Row(): | |
| # generate_button = gr.Button('Generate') | |
| # with gr.Column(): | |
| # result = gr.Image(label='Result') | |
| # with gr.Row(): | |
| # example_weights = gr.Dataset( | |
| # components=[structure_weight, color_weight], | |
| # samples=[ | |
| # [0.6, 1.0], | |
| # [0.3, 1.0], | |
| # [0.0, 1.0], | |
| # [1.0, 0.0], | |
| # ]) | |
| gr.Markdown(FOOTER) | |
| # preprocess_button.click(fn=model.detect_and_align_face, | |
| # inputs=input_image, | |
| # outputs=aligned_face) | |
| # aligned_face.change(fn=model.reconstruct_face, | |
| # inputs=aligned_face, | |
| # outputs=[ | |
| # reconstructed_face, | |
| # instyle, | |
| # ]) | |
| # style_type.change(fn=update_slider, | |
| # inputs=style_type, | |
| # outputs=style_index) | |
| # style_type.change(fn=update_style_image, | |
| # inputs=style_type, | |
| # outputs=style_image) | |
| # generate_button.click(fn=model.generate, | |
| # inputs=[ | |
| # style_type, | |
| # style_index, | |
| # structure_weight, | |
| # color_weight, | |
| # structure_only, | |
| # instyle, | |
| # ], | |
| # outputs=result) | |
| # example_images.click(fn=set_example_image, | |
| # inputs=example_images, | |
| # outputs=example_images.components) | |
| # example_styles.click(fn=set_example_styles, | |
| # inputs=example_styles, | |
| # outputs=example_styles.components) | |
| # example_weights.click(fn=set_example_weights, | |
| # inputs=example_weights, | |
| # outputs=example_weights.components) | |
| demo.launch( | |
| # enable_queue=args.enable_queue, | |
| # server_port=args.port, | |
| # share=args.share, | |
| ) | |
| if __name__ == '__main__': | |
| main() |