make colab happy

gradio_app.py

@@ -0,0 +1,319 @@
+import os
+import gradio as gr
+import torch
+import numpy as np
+import imageio
+from PIL import Image
+import uuid
+
+from draggan import utils
+from draggan.draggan import drag_gan
+from draggan import draggan as draggan
+
+device = 'cuda'
+
+
+SIZE_TO_CLICK_SIZE = {
+    1024: 8,
+    512: 5,
+    256: 2
+}
+
+CKPT_SIZE = {
+    'stylegan2/stylegan2-ffhq-config-f.pkl': 1024,
+    'stylegan2/stylegan2-cat-config-f.pkl': 256,
+    'stylegan2/stylegan2-church-config-f.pkl': 256,
+    'stylegan2/stylegan2-horse-config-f.pkl': 256,
+    'ada/ffhq.pkl': 1024,
+    'ada/afhqcat.pkl': 512,
+    'ada/afhqdog.pkl': 512,
+    'ada/afhqwild.pkl': 512,
+    'ada/brecahad.pkl': 512,
+    'ada/metfaces.pkl': 512,
+    'human/stylegan_human_v2_512.pkl': 512,
+    'human/stylegan_human_v2_1024.pkl': 1024,
+    'self_distill/bicycles_256_pytorch.pkl': 256,
+    'self_distill/dogs_1024_pytorch.pkl': 1024,
+    'self_distill/elephants_512_pytorch.pkl': 512,
+    'self_distill/giraffes_512_pytorch.pkl': 512,
+    'self_distill/horses_256_pytorch.pkl': 256,
+    'self_distill/lions_512_pytorch.pkl': 512,
+    'self_distill/parrots_512_pytorch.pkl': 512,
+}
+
+DEFAULT_CKPT = 'ada/afhqcat.pkl'
+
+
+def to_image(tensor):
+    tensor = tensor.squeeze(0).permute(1, 2, 0)
+    arr = tensor.detach().cpu().numpy()
+    arr = (arr - arr.min()) / (arr.max() - arr.min())
+    arr = arr * 255
+    return arr.astype('uint8')
+
+
+def add_points_to_image(image, points, size=5):
+    image = utils.draw_handle_target_points(image, points['handle'], points['target'], size)
+    return image
+
+
+def on_click(image, target_point, points, size, evt: gr.SelectData):
+    if target_point:
+        points['target'].append([evt.index[1], evt.index[0]])
+        image = add_points_to_image(image, points, size=SIZE_TO_CLICK_SIZE[size])
+        return image, not target_point
+    points['handle'].append([evt.index[1], evt.index[0]])
+    image = add_points_to_image(image, points, size=SIZE_TO_CLICK_SIZE[size])
+    return image, not target_point
+
+
+def on_drag(model, points, max_iters, state, size, mask, lr_box):
+    if len(points['handle']) == 0:
+        raise gr.Error('You must select at least one handle point and target point.')
+    if len(points['handle']) != len(points['target']):
+        raise gr.Error('You have uncompleted handle points, try to selct a target point or undo the handle point.')
+    max_iters = int(max_iters)
+    W = state['W']
+
+    handle_points = [torch.tensor(p, device=device).float() for p in points['handle']]
+    target_points = [torch.tensor(p, device=device).float() for p in points['target']]
+
+    if mask.get('mask') is not None:
+        mask = Image.fromarray(mask['mask']).convert('L')
+        mask = np.array(mask) == 255
+
+        mask = torch.from_numpy(mask).float().to(device)
+        mask = mask.unsqueeze(0).unsqueeze(0)
+    else:
+        mask = None
+
+    step = 0
+    for image, W, handle_points in drag_gan(W, model['G'],
+                                            handle_points, target_points, mask,
+                                            max_iters=max_iters, lr=lr_box):
+        points['handle'] = [p.cpu().numpy().astype('int') for p in handle_points]
+        image = add_points_to_image(image, points, size=SIZE_TO_CLICK_SIZE[size])
+
+        state['history'].append(image)
+        step += 1
+        yield image, state, step
+
+
+def on_reset(points, image, state):
+    return {'target': [], 'handle': []}, state['img'], False
+
+
+def on_undo(points, image, state, size):
+    image = state['img']
+
+    if len(points['target']) < len(points['handle']):
+        points['handle'] = points['handle'][:-1]
+    else:
+        points['handle'] = points['handle'][:-1]
+        points['target'] = points['target'][:-1]
+
+    image = add_points_to_image(image, points, size=SIZE_TO_CLICK_SIZE[size])
+    return points, image, False
+
+
+def on_change_model(selected, model):
+    size = CKPT_SIZE[selected]
+
+    G = draggan.load_model(utils.get_path(selected), device=device)
+    model = {'G': G}
+    W = draggan.generate_W(
+        G,
+        seed=int(1),
+        device=device,
+        truncation_psi=0.8,
+        truncation_cutoff=8,
+    )
+    img, _ = draggan.generate_image(W, G, device=device)
+
+    state = {
+        'W': W,
+        'img': img,
+        'history': []
+    }
+
+    return model, state, img, img, size
+
+
+def on_new_image(model, seed):
+    G = model['G']
+    W = draggan.generate_W(
+        G,
+        seed=int(seed),
+        device=device,
+        truncation_psi=0.8,
+        truncation_cutoff=8,
+    )
+    img, _ = draggan.generate_image(W, G, device=device)
+
+    state = {
+        'W': W,
+        'img': img,
+        'history': []
+    }
+
+    points = {'target': [], 'handle': []}
+    target_point = False
+    return img, img, state, points, target_point
+
+
+def on_max_iter_change(max_iters):
+    return gr.update(maximum=max_iters)
+
+
+def on_save_files(image, state):
+    os.makedirs('draggan_tmp', exist_ok=True)
+    image_name = f'draggan_tmp/image_{uuid.uuid4()}.png'
+    video_name = f'draggan_tmp/video_{uuid.uuid4()}.mp4'
+    imageio.imsave(image_name, image)
+    imageio.mimsave(video_name, state['history'])
+    return [image_name, video_name]
+
+
+def on_show_save():
+    return gr.update(visible=True)
+
+
+def on_image_change(model, image_size, image):
+    image = Image.fromarray(image)
+    result = inverse_image(
+        model.g_ema,
+        image,
+        image_size=image_size
+    )
+    result['history'] = []
+    image = to_image(result['sample'])
+    points = {'target': [], 'handle': []}
+    target_point = False
+    return image, image, result, points, target_point
+
+
+def on_mask_change(mask):
+    return mask['image']
+
+
+def on_select_mask_tab(state):
+    img = to_image(state['sample'])
+    return img
+
+
+def main():
+    torch.cuda.manual_seed(25)
+
+    with gr.Blocks() as demo:
+        gr.Markdown(
+            """
+            # DragGAN
+            
+            Unofficial implementation of [Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold](https://vcai.mpi-inf.mpg.de/projects/DragGAN/)
+            
+            [Our Implementation](https://github.com/Zeqiang-Lai/DragGAN) | [Official Implementation](https://github.com/XingangPan/DragGAN) (Not released yet)
+
+            ## Tutorial
+            
+            1. (Opklional) Draw a mask indicate the movable region.
+            2. Setup a least one pair of handle point and target point.
+            3. Click "Drag it". 
+            
+            ## Hints
+            
+            - Handle points (Blue): the point you want to drag.
+            - Target points (Red): the destination you want to drag towards to.
+            
+            ## Primary Support of Custom Image.
+            
+            - We now support dragging user uploaded image by GAN inversion.
+            - **Please upload your image at `Setup Handle Points` pannel.** Upload it from `Draw a Mask` would cause errors for now.
+            - Due to the limitation of GAN inversion, 
+                - You might wait roughly 1 minute to see the GAN version of the uploaded image.
+                - The shown image might be slightly difference from the uploaded one.
+                - It could also fail to invert the uploaded image and generate very poor results.
+                - Idealy, you should choose the closest model of the uploaded image. For example, choose `stylegan2-ffhq-config-f.pkl` for human face. `stylegan2-cat-config-f.pkl` for cat.
+                
+            > Please fire an issue if you have encounted any problem. Also don't forgot to give a star to the [Official Repo](https://github.com/XingangPan/DragGAN), [our project](https://github.com/Zeqiang-Lai/DragGAN) could not exist without it.
+            """,
+        )
+        G = draggan.load_model(utils.get_path(DEFAULT_CKPT), device=device)
+        model = gr.State({'G': G})
+        W = draggan.generate_W(
+            G,
+            seed=int(1),
+            device=device,
+            truncation_psi=0.8,
+            truncation_cutoff=8,
+        )
+        img, F0 = draggan.generate_image(W, G, device=device)
+
+        state = gr.State({
+            'W': W,
+            'img': img,
+            'history': []
+        })
+        points = gr.State({'target': [], 'handle': []})
+        size = gr.State(CKPT_SIZE[DEFAULT_CKPT])
+        target_point = gr.State(False)
+
+        with gr.Row():
+            with gr.Column(scale=0.3):
+                with gr.Accordion("Model"):
+                    model_dropdown = gr.Dropdown(choices=list(CKPT_SIZE.keys()), value=DEFAULT_CKPT,
+                                                 label='StyleGAN2 model')
+                    seed = gr.Number(value=1, label='Seed', precision=0)
+                    new_btn = gr.Button('New Image')
+                with gr.Accordion('Drag'):
+                    with gr.Row():
+                        lr_box = gr.Number(value=2e-3, label='Learning Rate')
+                        max_iters = gr.Slider(1, 500, 20, step=1, label='Max Iterations')
+
+                    with gr.Row():
+                        with gr.Column(min_width=100):
+                            reset_btn = gr.Button('Reset All')
+                        with gr.Column(min_width=100):
+                            undo_btn = gr.Button('Undo Last')
+                    with gr.Row():
+                        btn = gr.Button('Drag it', variant='primary')
+
+                with gr.Accordion('Save', visible=False) as save_panel:
+                    files = gr.Files(value=[])
+
+                progress = gr.Slider(value=0, maximum=20, label='Progress', interactive=False)
+
+            with gr.Column():
+                with gr.Tabs():
+                    with gr.Tab('Setup Handle Points', id='input'):
+                        image = gr.Image(img).style(height=512, width=512)
+                    with gr.Tab('Draw a Mask', id='mask') as masktab:
+                        mask = gr.ImageMask(img, label='Mask').style(height=512, width=512)
+
+        image.select(on_click, [image, target_point, points, size], [image, target_point])
+        image.upload(on_image_change, [model, size, image], [image, mask, state, points, target_point])
+        mask.upload(on_mask_change, [mask], [image])
+        btn.click(on_drag, inputs=[model, points, max_iters, state, size, mask, lr_box], outputs=[image, state, progress]).then(
+            on_show_save, outputs=save_panel).then(
+            on_save_files, inputs=[image, state], outputs=[files]
+        )
+        reset_btn.click(on_reset, inputs=[points, image, state], outputs=[points, image, target_point])
+        undo_btn.click(on_undo, inputs=[points, image, state, size], outputs=[points, image, target_point])
+        model_dropdown.change(on_change_model, inputs=[model_dropdown, model], outputs=[model, state, image, mask, size])
+        new_btn.click(on_new_image, inputs=[model, seed], outputs=[image, mask, state, points, target_point])
+        max_iters.change(on_max_iter_change, inputs=max_iters, outputs=progress)
+        masktab.select(lambda: gr.update(value=None), outputs=[mask]).then(on_select_mask_tab, inputs=[state], outputs=[mask])
+    return demo
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', default='cuda')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('-p', '--port', default=None)
+    parser.add_argument('--ip', default=None)
+    args = parser.parse_args()
+    device = args.device
+    demo = main()
+    print('Successfully loaded, starting gradio demo')
+    demo.queue(concurrency_count=1, max_size=20).launch(share=args.share, server_name=args.ip, server_port=args.port)