| 205 | |
| 206 | |
| 207 | def main(): |
| 208 | torch.cuda.manual_seed(25) |
| 209 | |
| 210 | with gr.Blocks() as demo: |
| 211 | wrapped_model = ModelWrapper(ckpt=DEFAULT_CKPT, size=CKPT_SIZE[DEFAULT_CKPT]) |
| 212 | model = gr.State(wrapped_model) |
| 213 | sample_z = torch.randn([1, 512], device=device) |
| 214 | latent, noise = wrapped_model.g_ema.prepare([sample_z]) |
| 215 | sample, F = wrapped_model.g_ema.generate(latent, noise) |
| 216 | |
| 217 | gr.Markdown( |
| 218 | """ |
| 219 | # DragGAN |
| 220 | |
| 221 | Unofficial implementation of [Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold](https://vcai.mpi-inf.mpg.de/projects/DragGAN/) |
| 222 | |
| 223 | [Our Implementation](https://github.com/Zeqiang-Lai/DragGAN) | [Official Implementation](https://github.com/XingangPan/DragGAN) |
| 224 | |
| 225 | ## Tutorial |
| 226 | |
| 227 | 1. (Optional) Draw a mask indicate the movable region. |
| 228 | 2. Setup a least one pair of handle point and target point. |
| 229 | 3. Click "Drag it". |
| 230 | |
| 231 | ## Hints |
| 232 | |
| 233 | - Handle points (Blue): the point you want to drag. |
| 234 | - Target points (Red): the destination you want to drag towards to. |
| 235 | |
| 236 | ## Primary Support of Custom Image. |
| 237 | |
| 238 | - We now support dragging user uploaded image by GAN inversion. |
| 239 | - **Please upload your image at `Setup Handle Points` pannel.** Upload it from `Draw a Mask` would cause errors for now. |
| 240 | - Due to the limitation of GAN inversion, |
| 241 | - You might wait roughly 1 minute to see the GAN version of the uploaded image. |
| 242 | - The shown image might be slightly difference from the uploaded one. |
| 243 | - It could also fail to invert the uploaded image and generate very poor results. |
| 244 | - Idealy, you should choose the closest model of the uploaded image. For example, choose `stylegan2-ffhq-config-f.pt` for human face. `stylegan2-cat-config-f.pt` for cat. |
| 245 | |
| 246 | > 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. |
| 247 | """, |
| 248 | ) |
| 249 | state = gr.State({ |
| 250 | 'latent': latent, |
| 251 | 'noise': noise, |
| 252 | 'F': F, |
| 253 | 'sample': sample, |
| 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): |
| 262 | with gr.Accordion("Model"): |
| 263 | model_dropdown = gr.Dropdown(choices=list(CKPT_SIZE.keys()), value=DEFAULT_CKPT, |
| 264 | label='StyleGAN2 model') |