DragGAN: An AI-Based Image Editing Tool

DragGAN or Drag your GAN is a newly released open-source AI tool that allows user to manipulate or edit images by dragging certain point on the input image.

We have seen lots of Text to image ai tool, but DragGAN is totally different. This is one of the first ai tool to do image editing. Once this tool will mature, It can be a killer for Photoshop.

DragGAN Demo

Researchers released some demos in DragGAN paper. Let’s see some of those interesting demo videos to understand how this new deep learning model works.

Isn’t it so interesting? You just need to select some point and it will do AI magic for you. Most important part is that, it will do everything in seconds.

You can do almost anything with it. You can see in above demo that it is changing the perspective or angle of the image in a few clicks. It is converting a neutral face into a laughing face. Not only that it can work for natural landscapes also.

How DragGAN works

DragGAN is an AI editing tool developed by researchers from Google, the Max Planck Institute of Informatics, and MIT CSAIL. It utilizes a technique called Generative Adversarial Network (GAN) to enable users to manipulate images with ease.

DragGAN leverages a pre-trained GAN, which is a type of neural network consisting of two components: a generator and a discriminator.

1. Generator: It generates new images based on the user’s input
2. Discriminator: Evaluates the realism of those generated images

During the training phase, the generator tries to generate images. Then the discriminator decides whether they are real or not. This kind of adversarial training process helps the DragGAN ai tool to produce more realistic and convincing images.

DragGAN ai tool website

If you are interested to know more about this ai tool, visit DragGAN’s official website for demos and the DragGAN paper to explore their innovative work.

Download DragGAN

You can visit DragGAN GitHub site to download code of this ai tool. The code is not yet released. It will be released in June 2023.

I will keep my eyes on it. As soon as the code is released, I will try this and share the result about how you can access and use DragGAN to edit your own image.

Official DragGAN

Official DraGAN code finally released. You can visit DragGAN GitHub site to download code of this ai tool. There are two ways you can access official DragGAN AI tool: without code and with code.

Official DragGAN without Code

If you want to use DragGAN AI-based image editing tool without any coding, go to this website. This is basically a Huggingface link where the official DragGan application is deployed.

Once you open that link the official DragGAN AI tool will open. Let me give you a quick demo about how you can properly use that application.

In this DragGAN application, you need to follow some steps, then only it will work as per your desire. First, you need to select a proper algorithm. They call it pre-trained model.

In this hugging face DragGAN tool, only two models are available: stylegan_human_v2_512 and stylegan2_dogs_1024_pytorch. One for editing humans and another for dogs.

Next, you need to select the source and destination point. In the above image, red dot is the source point and blue dot is the destination point. After doing this you can simply click on the Start button to see the Official DragGAN tool in action.

As you can see in the video above, the official dragGAN AI tool in Huggingface is displaying low-quality input images. However, in this particular case, I discovered that the unofficial DragGAN is providing high-quality input images. Editing quality is I would say Okay if not good.

Official DragGAN with Code

As a coder, we always try to debug anything with code. This way we can explore what is happening in the background. So let’s implement official DragGAN in Google Colab.

Why google colab? To implement DragGAN, you need at least 14 GB GPU memory which I don’t have. So thought of utilizing free 14 GB GPU of Google Colab.

To avail free GPU, open a new Colab notebook. Go to Runtime -> Change runtime type then select Hardware accelerator as GPU and Save.

Let me break the entire code into some steps. Run each block of code in Colab Notebook.

Step1: Download files and Install Libraries

# Install Required libraries
!pip install gradio
!pip install Ninja

# Clone DragGAN official ripository
!git clone https://github.com/XingangPan/DragGAN.git

# Go inside that directory
%cd /content/DragGAN/

# Download pre trained models
!python scripts/download_model.py

In the above code we are installing required libraries to run DragGAN AI tool in google colab. Next downloading DragGAN official repository and finally downloading pre-trained models. All downloaded models will be saved inside ./checkpoints folder.

Step2: Download Human Model

The above code will download all pre-trained models except human model. So if you want to work with human images you need to download weight of stylegan_human_v2_512 model manually and put them under ./checkpoints.

You can download stylegan_human_v2_512 model from this DragGAN google drive link. I found difficulty to upload that model inside ./checkpoints folder. It was taking so much time to upload that 336 MB model.

To solve that, I created a folder inside my google drive named “draggan_drive” (you can name it anything). Then using below code, I just coppied that model to ./checkpoints folder.

# Mount Google Drive
from google.colab import drive
drive.mount('/content/drive')

# Copy Human model from google drive to checkpoints
!cp /content/drive/MyDrive/draggan_drive/stylegan_human_v2_512.pkl /content/DragGAN/checkpoints

Step5: Run GragGAN Visualizer

Finally below below code to open GragGAN AI tool in a gradio visualizer link:

!python visualizer_drag_gradio.py --share

Once you successfully run above code in your Google Colab, you should see a gradio public URL like this: https://c4dc9240665215f995.gradio.live. Open that URL in any web browser. That’s it you should see the official DragGAN ai tool is running.

All the functionality is same as Huggingface DragGAN tool. Only difference is here in with code version, you will get lots of models (while in Huggingface, only human and dog models were available). Another difference is, since we are running it on a Colab GPU, the process is faster compared to Huggingface.

Official DragGAN with Custom Image

In above examples we used which is provided by DragGAN. Now let’s see how we can edit our own custom image using DragGAN AI tool. I will show you the step by step process.

We will implement DragGAN for custom images in Google Colab. So it does not matter whether you are using Windows or Linux system. Everyone can follow the process. We just need two Google Colab notebooks.

Notebook 1

To do that, first, open a Google Colab notebook then go to Runtime and then select GPU.

Bypass Download Rate Limit

Once you set up your notebook runtime run the below code to bypass Google Drive download limit rate. Once you run this code, you need to Allow permission to access your Google Drive. Simply allow it.

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from oauth2client.client import GoogleCredentials

download_with_pydrive = True #@param {type:"boolean"}

class Downloader(object):
    def __init__(self, use_pydrive):
        self.use_pydrive = use_pydrive

        if self.use_pydrive:
            self.authenticate()

    def authenticate(self):
        auth.authenticate_user()
        gauth = GoogleAuth()
        gauth.credentials = GoogleCredentials.get_application_default()
        self.drive = GoogleDrive(gauth)

    def download_file(self, file_id, file_dst):
        if self.use_pydrive:
            downloaded = self.drive.CreateFile({'id':file_id})
            downloaded.FetchMetadata(fetch_all=True)
            downloaded.GetContentFile(file_dst)
        else:
            !gdown --id $file_id -O $file_dst

downloader = Downloader(download_with_pydrive)

Install Required Packages

Let’s now install all required libraries. We will be going to use pip command and for some of the libraries we will directly fetch from Github.

## Other packages are already builtin in the Colab interpreter
!pip install wandb
!pip install lpips

## Used for faster inference of StyleGAN by enabling C++ code compilation

!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
!sudo unzip ninja-linux.zip -d /usr/local/bin/
!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force

Collecting wandb
  Downloading wandb-0.15.8-py3-none-any.whl (2.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.1/2.1 MB 10.3 MB/s eta 0:00:00
Requirement already satisfied: Click!=8.0.0,>=7.1 in /usr/local/lib/python3.10/dist-packages (from wandb) (8.1.6)
Collecting GitPython!=3.1.29,>=1.0.0 (from wandb)
  Downloading GitPython-3.1.32-py3-none-any.whl (188 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 188.5/188.5 kB 15.2 MB/s eta 0:00:00
Requirement already satisfied: requests<3,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (2.27.1)
Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb) (5.9.5)
Collecting sentry-sdk>=1.0.0 (from wandb)
  Downloading sentry_sdk-1.29.2-py2.py3-none-any.whl (215 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 215.6/215.6 kB 12.1 MB/s eta 0:00:00
.....
.....
.....

Download Pre-Trained Models

In this step we will create and setup a runtime directory and download all required files inside that directory. The created folder name is PTI.

# Create and set up directory
import os
os.chdir('/content')
CODE_DIR = 'PTI'

!git clone https://github.com/danielroich/PTI.git $CODE_DIR

os.chdir(f'./{CODE_DIR}')

# Import required libraries
import os
import sys
import pickle
import numpy as np
from PIL import Image
import torch
from configs import paths_config, hyperparameters, global_config
from utils.align_data import pre_process_images
from scripts.run_pti import run_PTI
from IPython.display import display
import matplotlib.pyplot as plt
from scripts.latent_editor_wrapper import LatentEditorWrapper

current_directory = os.getcwd()
save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR, "pretrained_models")
os.makedirs(save_path, exist_ok=True)

## Download pretrained StyleGAN on FFHQ 1024x1024
downloader.download_file("125OG7SMkXI-Kf2aqiwLLHyCvSW-gZk3M", os.path.join(save_path, 'ffhq.pkl'))

## Download Dlib tool for alingment, used for preprocessing images before PTI
downloader.download_file("1xPmn19T6Bdd-_RfCVlgNBbfYoh1muYxR", os.path.join(save_path, 'align.dat'))

Cloning into 'PTI'...
remote: Enumerating objects: 336, done.
remote: Counting objects: 100% (95/95), done.
remote: Compressing objects: 100% (46/46), done.
remote: Total 336 (delta 61), reused 56 (delta 49), pack-reused 241
Receiving objects: 100% (336/336), 13.70 MiB | 15.18 MiB/s, done.
Resolving deltas: 100% (125/125), done.

Set Up Configuration Files

In this step, we will define some variables, which are required through this code. Let’s define those in the below code block. One thing to note that, you need to change image_name as per the name of your image.

In my case, I am going to use my image as the input of DragGAN. My image name is Anindya-Naskar.jpg

image_dir_name = 'image'

## If set to true download desired image from given url. If set to False, assumes you have uploaded personal image to
## 'image_original' dir
use_image_online = True
# Change this according to your image name
image_name = 'Anindya-Naskar'
use_multi_id_training = False
global_config.device = 'cuda'
paths_config.e4e = '/content/PTI/pretrained_models/e4e_ffhq_encode.pt'
paths_config.input_data_id = image_dir_name
paths_config.input_data_path = f'/content/PTI/{image_dir_name}_processed'
paths_config.stylegan2_ada_ffhq = '/content/PTI/pretrained_models/ffhq.pkl'
paths_config.checkpoints_dir = '/content/PTI/'
paths_config.style_clip_pretrained_mappers = '/content/PTI/pretrained_models'
hyperparameters.use_locality_regularization = False

Create Required Directories

In this step, we need to create two additional folders inside our PTI directory. The names of those folders will be:

• image_original: In this folder we will keep our custom input image
• image_processed: After some pre-processing, output image will be stored inside this folder

os.makedirs(f'./{image_dir_name}_original', exist_ok=True)
os.makedirs(f'./{image_dir_name}_processed', exist_ok=True)
os.chdir(f'./{image_dir_name}_original')

Read Custom Image

Let’s now read a custom image. To do this you need to upload your custom image inside image_original folder. Also, you need to change your image name in the configuration part which I told you earlier.

original_image = Image.open(f'/content/PTI/{image_dir_name}_original/{image_name}.jpg')
original_image

Pre-Process Input Image

In this step, we need to convert our input image to 512×512 format. This is required to match the input image to the model dimension. The pre-processed image will be saved inside image_processed folder.

os.chdir('/content/PTI')
pre_process_images(f'/content/PTI/{image_dir_name}_original')
aligned_image = Image.open(f'/content/PTI/{image_dir_name}_processed/{image_name}.jpeg')
aligned_image.resize((512,512))

Invert Image using

Now we need to invert our image using PTI method. Below code is to do that. This will take around 5-10 minutes to complete.

os.chdir('/content/PTI')
model_id = run_PTI(use_wandb=False, use_multi_id_training=use_multi_id_training)

Once the above code execution will be finished, two files will be created:

• 0.pt: Inside /content/PTI/embeddings/image/PTI/Anindya-Naskar/0.pt
• model_WTRITDZUIRCI_Anindya-Naskar.pt: Inside rood directory

Convert Model to Pickle format

As per requirement of the DragGAN, the model must be in pickle format. That is the reason in the below code we are converting .pt model to .pkl format.

def load_generators(model_id, image_name):
  with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
    old_G = pickle.load(f)['G_ema'].cuda()

  with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
    new_G = torch.load(f_new).cuda()

  return old_G, new_G

generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
old_G, new_G = load_generators(model_id, generator_type)

def export_updated_pickle(new_G,model_id):
  print("Exporting large updated pickle based off new generator and ffhq.pkl")
  with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
    d = pickle.load(f)
    old_G = d['G_ema'].cuda()
    old_D = d['D'].eval().requires_grad_(False).cpu()

  tmp = {}
  tmp['G'] = old_G.eval().requires_grad_(False).cpu()
  tmp['G_ema'] = new_G.eval().requires_grad_(False).cpu()
  tmp['D'] = old_D
  tmp['training_set_kwargs'] = None
  tmp['augment_pipe'] = None


  with open(f'{paths_config.checkpoints_dir}/stylegan2_custom_512_pytorch.pkl', 'wb') as f:
      pickle.dump(tmp, f)


export_updated_pickle(new_G,model_id)

Once you run this code, a pickle file named stylegan2_custom_512_pytorch.pkl will be created inside your root directory.

Copy files to Drive

Now we need to copy that pickle file and the 0.pt embedding file to our google drive to load those in another colab notebook. To do that, just run below block of code.

# Mount Google drive
from google.colab import drive
drive.mount('/content/drive')

# To avoid error: NotImplementedError: A UTF-8 locale is required. Got ANSI_X3.4-1968
import locale
locale.getpreferredencoding = lambda: "UTF-8"

# Copy Model
!cp /content/PTI/stylegan2_custom_512_pytorch.pkl /content/drive/MyDrive/DragGAN_Custom_image

# Copy Image embeddings
!cp /content/PTI/embeddings/image/PTI/Anindya-Naskar/0.pt /content/drive/MyDrive/DragGAN_Custom_image

Note: Before I ran this code I created a folder named DragGAN_Custom_image inside my google drive. You can change this name accordingly.

Create Notebook 2

Now work of Notebook 1 is finished. Let’s now create another colab notebook. In the same way you also need to change runtime of this notebook to GPU.

Clone DragGAN Git

In the second notebook run below commands to clone DragGAN official files from github.

!git clone https://github.com/XingangPan/DragGAN.git

%cd /content/DragGAN/

Cloning into 'DragGAN'...
remote: Enumerating objects: 416, done.
remote: Counting objects: 100% (175/175), done.
remote: Compressing objects: 100% (95/95), done.
remote: Total 416 (delta 112), reused 84 (delta 80), pack-reused 241
Receiving objects: 100% (416/416), 34.08 MiB | 17.09 MiB/s, done.
Resolving deltas: 100% (169/169), done.
/content/DragGAN

Move Trained model files

In this step, we will create an empty folder named checkpoints inside our root directory. Then we will copy model files (created in previous notebook) from google drive to this folder.

import os
os.makedirs(f'./checkpoints', exist_ok=True)

from google.colab import drive
drive.mount('/content/drive')

# Copy Model
!cp /content/drive/MyDrive/DragGAN_Custom_image/stylegan2_custom_512_pytorch.pkl /content/DragGAN/checkpoints

# Copy Image embeddings
!cp /content/drive/MyDrive/DragGAN_Custom_image/0.pt /content/DragGAN/checkpoints

Once you run this code those two files will be copied to the checkpoints folder.

Change in Code

Now one thing you need to do is open visualizer_drag_gradio.py code (double click). Then you need to modify two things.

Modification 1

In the first modification, look for init_pkl variable. In that variable you need to change the model name to our custom model name (stylegan2_custom_512_pytorch).

Previously, it might be written like this: 'stylegan2_custom_512_pytorch‘. You need to change it. You may find this line of code at around line number 164.

Modification 2

In the second and final modification, you need to modify init_images function. Inside this function you need to load 0.pt embedding file and change None to that variable name like below image.

Once you are done with all those changes, save that python file. Now we are ready to launch the final gradio application.

Run DragGAN Application

Once you are done with those two modifications (total 3 changes), we are finally ready to launch our DragGAN application for my own custom image. Let’s run it using below block of code.

!pip install gradio
!pip install Ninja

!python visualizer_drag_gradio.py --share

Above code will generate a public URL. Open that public URL to edit your own image.

Below video is just an demo how I was playing with my own custom image in DragGAN AI tool.

Un-Official DragGAN

If you still want to play with DragGan AI image editing tool, along with the official version, there is an unofficial version you can try with. I will guide you how can use that unofficial DragGan.

There are two ways you can play with that unofficial Dragan AI tool: 1. Without coding, 2. With coding.

1. DragGAN without Coding

If you want to use DragGAN AI-based image editing tool without any coding, go to this website. This is basically a Huggingface link where the unofficial DragGan application is deployed.

Once you open that link the unofficial DragGAN AI tool will open. Let me give you a quick demo about how you can properly use that application.

Application Walkthrough

In this DragGAN application, you need to follow some steps, then only it will work as per your desire. Steps are:

Step1: Select Model

First you need to select a proper algorithm. They call it StyleGAN2 model. There are four algorithms you can choose from.

• stylegan2-ffhq-config-f.pt : This algorithm is for human face
• stylegan2-cat-config-f.pt : If you are working for cat image
• stylegan2-church-config-f.pt : To edit images of houses in DragGAN ai tool
• stylegan2-horse-config-f.pt : If you have horse images

Step2: Mention Iteration Number

In this step, you need to mention the number of iterations you want the diffusion based large AI model to edit the image. By default, number of iteration is set to 20. You can keep it like that.

Step3: Choose Input image

In this third step, you need to select your input image for editing. There are so many pre-built images available that you can play with. Just click New Image button and it will keep on showing other images.

Note: In this Hugging Face application I faced error while uploading a new image. We will solve this error the coding section below.

Step4: Select Mask Tab

There are two tabs: Draw a Mak and Setup Handle Points. In this step you need to select Draw a Mak tab. This tab is used to draw or mark or mask the area you need to modify in the input image.

Step5: Brush size

Before drawing the modification area, the application also gives you the flexibility to choose the drawing brush size. You can also play with it.

Step6: Draw Modification Area

Finally, in step 6 we need to draw the mask area or the area we want to modify in the input image using AI based image editing tool GradGAN.

This step is optional. While testing the application I found that without marking the modification area the application can drag points.

Step7: Select Points Tab

So we are done with the modification area or mask area drawing. Now we need to select Setup Handle Points tab to mention the points we want to drag.

Step8: Select Source Point

For this example, I want this girl smiley. For that, I am selecting her bottom leap as my source point (blue dot). You just need to do a simple left mouse click to select the point.

Step9: Select Target Point

Now to make the girl smile, her bottom leap needs to be down so that her teeth can be visible. So I selected the target point (red dot) below the source point.

Step10: Drag Point

Finally, click on Drag it button to edit your image as you desire using DragGAN AI photo editor. In our case, we wanted to make the girl smile.

Now let’s see how DragGAN converts the input image (normal) to our desire (smiley).

Disadvantage

I found one main disadvantage of this huggingface application of unofficial DragGAN while testing. That is you can not upload a new image. I was keep on getting error while trying to upload an image to this application.

But the same application when I run through code, I was able to upload outside image. Did not find any valid reason.

2. DragGAN with Coding

Now let’s see how you can run this unofficial DragGAN image editing ai tool in your local system. To do that, we need to first install DragGAN. Following command is to do the installation.

git clone https://github.com/Zeqiang-Lai/DragGAN.git
cd DragGAN
conda create -n draggan python=3.7
conda activate draggan
pip install -r requirements.txt

If you want to run this version of DragGAN in CPU you can run the below command.

python gradio_app.py --device cpu

Once you run the above command, you can access the GragGAN ai application using this local host link: http://127.0.0.1:7860

Now if you want to run DragGAN using your NVIDIA GPU, run below commands to install torch for GPU:

conda uninstall pytorch
pip uninstall torch
pip uninstall torch
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia
conda install cudatoolkit

Note: You can find your personalized torch installation command from this link. If you want more clarification, read this post: Train YOLOv8 on Custom dataset in Windows GPU

Second and third line of the above command is to avoid error like below:

OMP: Hint This means that multiple copies of the OpenMP runtime have been linked into the program. That is dangerous

Then run this command to open DragGAN using GPU:

python gradio_app.py

Problem with Local System

While running this unofficial DragGAN ai tool in my personal system I found some issues. Let me explain those.

If you run in your CPU, you will not able to upload new image to this DragGAN ai tool. Same error like Huggingface application explained above (without coding).

To upload new images, you must need to run DragGAN using NVIDIA GPU. Now again the problem is your GPU memory needs to be more than 6GB. I had 4GB graphics card and got below error while uploading new image:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 4.00 GiB total capacity; 3.06 GiB already allocated; 0 bytes free; 3.18 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Run in Google Colab

To avoid those issues, I tried implementing DragGAN ai image editing tool in google collab and I got successful.

First, you need to set up your Colab notebook with GPU backend. To do that go to Edit > Notebook settings or Runtime>Change runtime type and select GPU as Hardware accelerator

Now run below codes in separate cell in Google Colab to open the unofficial DragGAN ai application.

# To check whether you got GPU or not
!nvidia-smi -L

# Clone unofficial Github project of DragGAN
!git clone https://github.com/Zeqiang-Lai/DragGAN.git

# No need to explain
!cd DragGAN

# Install requied libraries
!pip install -r /content/DragGAN/requirements.txt

# Set your working directory
import sys
sys.path.append(".")
sys.path.append('./DragGAN')

# Run DragGAN ai tool in Gradio APP
from gradio_app import main

demo = main()
# To create a public link, set `share=True`
demo.queue(concurrency_count=1, max_size=20).launch(share=True)

Once you successfully run above code in your Google Colab, you should see a gradio public URL like this: https://c4dc9240665215f995.gradio.live

Open that URL in any web browser. That’s it you should see the DragGAN ai tool is running.

End Note

By utilizing the power of GANs, DragGAN enables users to make image edits in real time by simply clicking and dragging. It provides a convenient and intuitive editing experience to the user.

It may replace image editing tool like Photoshop or Illustrator in near future. But at this time DragGAN is in beta mode and code is not released.

Similar Read:

• Top 9 AI Tools Better Than Chat GPT – 100% FREE
• 7 Best Text to Image AI Image Generators Free tool
• Convert Text to Speech with Deep Learning in Python

Anindya

Hi there, I’m Anindya Naskar, Data Science Engineer. I created this website to show you what I believe is the best possible way to get your start in the field of Data Science.