Skip to content

Commit

Permalink
update README.md
Browse files Browse the repository at this point in the history
  • Loading branch information
kennethnym committed Jan 16, 2025
1 parent d2b6e70 commit 1119f00
Showing 1 changed file with 19 additions and 32 deletions.
51 changes: 19 additions & 32 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,14 +1,21 @@
# Multilayer Authenticity Identifier (MAI)

MAI is a research project that attempts to train a CNN model to identify synthetic AI images.
MAI is a research project that attempts to train a machine learning model that identifies AI-generated images.

## Why?

i am bored.

## Architecture

nothing is set in stone, but at the moment, MAI is a simple CNN model that looks like this:
i have experimented with the following model architectures with varying degrees of success.
they reside in their own directory.

### convolutional neural network (CNN)

the code for the CNN model is located at `./resnet`.

the CNN model looks like this:

1. 16-channel, 3x3 convolution layer -> 2x2 max pooling -> relu activation
2. 32-channel, 3x3 convolution layer -> 2x2 max pooling -> relu activation
Expand All @@ -19,40 +26,20 @@ the model expects a 200x200 image as an input and outputs a score, with 1 being

[BCEWithLogitLoss](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html) is used as the loss fn, and [RMSprop](https://pytorch.org/docs/stable/generated/torch.optim.RMSprop.html) as the optimizer.

## Datasets

MAI has been trained on the following datatsets:

- [poloclub/diffusiondb](https://huggingface.co/datasets/poloclub/diffusiondb)
- [nlphuji/flickr30k](https://huggingface.co/datasets/nlphuji/flickr30k)
- [keremberke/painting-style-classification](https://huggingface.co/datasets/keremberke/painting-style-classification)
- [animelover/scenery-images](https://huggingface.co/datasets/animelover/scenery-images)
- [nanxstats/movie-poster-5k](https://huggingface.co/datasets/nanxstats/movie-poster-5k)
- [Alphonsce/metal_album_covers](https://huggingface.co/datasets/Alphonsce/metal_album_covers)

## How to train?

make sure to have [poetry](https://python-poetry.org) installed.

clone the project, and run:

```
poetry install
```

open a shell in the venv created by poetry:
### vision language model (VLM)

```
poetry shell
```
the code for the vlm model finetuning is located at `./moondream`.

run `train.py` to train the model. make sure cuda is available as a cuda-enabled gpu is used to accelerate training. for each epoch, if the validation loss is less than the last epoch, the model is saved locally. you can customize the location easily in `train.py`.
following [the Bi-LoRA paper](https://arxiv.org/abs/2404.01959) which suggests finetuning a vlm on real and ai-generated images, i decided to finetune [moondream](https://moondream.ai/), a small vlm.

## How to run inference?
the training data consists of 50% real images and 50% ai-generated images, along with a q&a pair.
the question is always "Is this image AI-generated?", followed by the answer "Yes." or "No.", depending on whether the image is ai-generated or not.

run `inference.py` instead. place your test images in `test_images/` directory, and don't forget to reference the images in `inference.py`.
preliminary experiments show mixed results. to improve the training data, i decided to generate my own dataset using the following steps:

## More on modal.com
1. ask moondream to caption a real image.
2. use that caption to prompt a model (stable diffusion 3.5 large at the moment) to generate the equivalent ai image.
3. add the image to the row alongside the real image.

i am using (https://modal.com) to run the training and inference, but u can get rid of the modal.com glue pretty easily. you should first remove the decorators above the functions, then at where the functions are invoked, remove `.remote()` and instead invoke the function directly. remove `app` and `vol` variables as well.
hopefully, this will help moondream better understand the difference between real and ai-generated images.

0 comments on commit 1119f00

Please sign in to comment.