| 202 | |
| 203 | |
| 204 | def main(): |
| 205 | torch.cuda.manual_seed(25) |
| 206 | |
| 207 | with gr.Blocks() as demo: |
| 208 | gr.Markdown( |
| 209 | """ |
| 210 | # DragGAN |
| 211 | |
| 212 | Unofficial implementation of [Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold](https://vcai.mpi-inf.mpg.de/projects/DragGAN/) |
| 213 | |
| 214 | [Our Implementation](https://github.com/Zeqiang-Lai/DragGAN) | [Official Implementation](https://github.com/XingangPan/DragGAN) (Not released yet) |
| 215 | |
| 216 | ## Tutorial |
| 217 | |
| 218 | 1. (Optional) Draw a mask indicate the movable region. |
| 219 | 2. Setup a least one pair of handle point and target point. |
| 220 | 3. Click "Drag it". |
| 221 | |
| 222 | ## Hints |
| 223 | |
| 224 | - Handle points (Blue): the point you want to drag. |
| 225 | - Target points (Red): the destination you want to drag towards to. |
| 226 | |
| 227 | ## Primary Support of Custom Image. |
| 228 | |
| 229 | - We now support dragging user uploaded image by GAN inversion. |
| 230 | - **Please upload your image at `Setup Handle Points` pannel.** Upload it from `Draw a Mask` would cause errors for now. |
| 231 | - Due to the limitation of GAN inversion, |
| 232 | - You might wait roughly 1 minute to see the GAN version of the uploaded image. |
| 233 | - The shown image might be slightly difference from the uploaded one. |
| 234 | - It could also fail to invert the uploaded image and generate very poor results. |
| 235 | - 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. |
| 236 | |
| 237 | > 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. |
| 238 | """, |
| 239 | ) |
| 240 | G = draggan.load_model(utils.get_path(DEFAULT_CKPT), device=device) |
| 241 | model = gr.State({'G': G}) |
| 242 | W = draggan.generate_W( |
| 243 | G, |
| 244 | seed=int(1), |
| 245 | device=device, |
| 246 | truncation_psi=0.8, |
| 247 | truncation_cutoff=8, |
| 248 | ) |
| 249 | img, F0 = draggan.generate_image(W, G, device=device) |
| 250 | |
| 251 | state = gr.State({ |
| 252 | 'W': W, |
| 253 | 'img': img, |
| 254 | 'history': [] |
| 255 | }) |
| 256 | points = gr.State({'target': [], 'handle': []}) |
| 257 | size = gr.State(CKPT_SIZE[DEFAULT_CKPT]) |
| 258 | target_point = gr.State(False) |
| 259 | |
| 260 | with gr.Row(): |
| 261 | with gr.Column(scale=0.3): |