initial commit

.gitignore

@@ -0,0 +1,252 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python,macos,linux,windows
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,macos,linux,windows
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### macOS Patch ###
+# iCloud generated files
+*.icloud
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+# End of https://www.toptal.com/developers/gitignore/api/python,macos,linux,windows
+
+test_images/

README.md

@@ -0,0 +1,58 @@
+# Multilayer Authenticity Identifier (MAI)
+
+MAI is a research project that attempts to train a CNN model to identify synthetic AI 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:
+
+1. 16-channel, 3x3 convolution layer -> 2x2 max pooling -> relu activation
+2. 32-channel, 3x3 convolution layer -> 2x2 max pooling -> relu activation
+3. 64-channel, 3x3 convolution layer -> 2x2 max pooling -> relu activation
+4. 40,000-neuron layer -> relu -> 120-neuron layer -> relu -> 30 -> 1
+
+the model expects a 200x200 image as an input and outputs a score, with 1 being that the input image is absolutely synthetic, and 0 being that it is absolutely authentic.
+
+[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:
+
+```
+poetry shell
+```
+
+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`.
+
+## How to run inference?
+
+run `inference.py` instead. place your test images in `test_images/` directory, and don't forget to reference the images in `inference.py`.
+
+## More on modal.com
+
+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.
+

augmentation.py

@@ -0,0 +1,50 @@
+import albumentations as a
+import numpy as np
+from albumentations.pytorch import ToTensorV2
+from hyperparams import CROP_SIZE
+
+
+preprocess_training = a.Compose(
+    [
+        a.augmentations.PadIfNeeded(min_width=CROP_SIZE, min_height=CROP_SIZE),
+        a.RandomCrop(width=CROP_SIZE, height=CROP_SIZE),
+        a.Flip(p=0.5),
+        a.RandomRotate90(p=0.5),
+        a.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ToTensorV2(),
+    ]
+)
+preprocess_validation = a.Compose(
+    [
+        a.augmentations.PadIfNeeded(min_width=CROP_SIZE, min_height=CROP_SIZE),
+        a.CenterCrop(width=CROP_SIZE, height=CROP_SIZE),
+        a.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ToTensorV2(),
+    ]
+)
+
+
+def transform_training(example):
+    transformed = []
+    for pil_image in example["image"]:
+        array = np.array(pil_image.convert("RGB"))
+        # check if image is in (height, width, channel) shape
+        # if not, do a transpose
+        if array.shape[-1] != 3:
+            array = np.transpose(array, (1, 2, 0))
+        img = preprocess_training(image=array)["image"]
+        transformed.append(img)
+    example["pixel_values"] = transformed
+    return example
+
+
+def transform_validation(example):
+    transformed = []
+    for pil_image in example["image"]:
+        array = np.array(pil_image.convert("RGB"))
+        if array.shape[-1] != 3:
+            array = np.transpose(array, (1, 2, 0))
+        img = preprocess_validation(image=array)["image"]
+        transformed.append(img)
+    example["pixel_values"] = transformed
+    return example

hyperparams.py

@@ -0,0 +1,5 @@
+FILTER_COUNT = 32
+KERNEL_SIZE = 2
+CROP_SIZE = 200
+BATCH_SIZE = 4
+EPOCHS = 15

inference.py

@@ -0,0 +1,38 @@
+import modal
+import torch
+import numpy as np
+from PIL import Image
+from model import mai
+from augmentation import preprocess_validation
+
+MODEL_NAME = "mai_20240424_180855_4"
+
+image = modal.Image.debian_slim().pip_install(
+    "datasets==2.19.0",
+    "albumentations==1.4.4",
+    "numpy==1.26.4",
+    "torch==2.2.2",
+)
+app = modal.App("multilayer-authenticity-identifier", image=image)
+volume = modal.Volume.from_name("model-store")
+model_store_path = "/vol/models"
+
+
+@app.function(timeout=5000, gpu="T4", volumes={model_store_path: volume})
+def load_model_and_run_inference(img):
+    print(f"REMOTE: {img.shape}")
+    mai.load_state_dict(torch.load(f"{model_store_path}/{MODEL_NAME}"))
+    mai.eval()
+    img_batch = np.expand_dims(img, axis=0)
+    img_batch = torch.tensor(img_batch)
+    prediction = mai(img_batch)
+    prediction = torch.sigmoid(prediction)
+    print(prediction)
+
+
+@app.local_entrypoint()
+def main():
+    img = Image.open("test_images/dog.jpg")
+    img = preprocess_validation(image=np.array(img))["image"]
+    print(img.shape)
+    load_model_and_run_inference.remote(img)

label.py

@@ -0,0 +1,8 @@
+def label_fake(example):
+    example["is_synthetic"] = 1.0
+    return example
+
+
+def label_real(example):
+    example["is_synthetic"] = 0.0
+    return example

model.py

@@ -0,0 +1,19 @@
+import torch.nn as nn
+
+mai = nn.Sequential(
+    nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
+    nn.ReLU(),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+    nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
+    nn.ReLU(),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+    nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
+    nn.ReLU(),
+    nn.MaxPool2d(kernel_size=2, stride=2),
+    nn.Flatten(),
+    nn.Linear(25 * 25 * 64, 120),
+    nn.ReLU(),
+    nn.Linear(120, 30),
+    nn.ReLU(),
+    nn.Linear(30, 1),
+)

pyproject.toml

@@ -0,0 +1,28 @@
+[tool.poetry]
+name = "mai"
+version = "0.1.0"
+description = "Multilayer Authenticity Identifier"
+authors = ["Kenneth <kennethnym@outlook.com>"]
+license = "MIT"
+readme = "README.md"
+package-mode = false
+
+
+[tool.poetry.dependencies]
+python = "^3.12"
+mlx = "^0.11.0"
+datasets = "^2.19.0"
+albumentations = "^1.4.4"
+numpy = "^1.26.4"
+modal = "^0.62.103"
+torch = "^2.2.2"
+
+
+[tool.pyright]
+venvPath = "."
+venv = ".venv"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

train.py

@@ -0,0 +1,218 @@
+import datasets
+import torch
+import torch.nn
+import torch.optim
+import torch.utils.data
+import modal
+from datetime import datetime
+from datasets import concatenate_datasets, load_dataset
+from label import label_fake, label_real
+from model import mai
+from augmentation import transform_training, transform_validation
+from hyperparams import BATCH_SIZE, EPOCHS
+
+TEST_SIZE = 0.1
+
+datasets.logging.set_verbosity(datasets.logging.INFO)
+
+image = modal.Image.debian_slim().pip_install(
+    "datasets==2.19.0",
+    "albumentations==1.4.4",
+    "numpy==1.26.4",
+    "torch==2.2.2",
+)
+app = modal.App("multilayer-authenticity-identifier", image=image)
+volume = modal.Volume.from_name("model-store")
+model_store_path = "/vol/models"
+
+
+def collate(batch):
+    pixel_values = []
+    is_synthetic = []
+    for row in batch:
+        is_synthetic.append(row["is_synthetic"])
+        pixel_values.append(torch.tensor(row["pixel_values"]))
+
+    pixel_values = torch.stack(pixel_values, dim=0)
+    is_synthetic = torch.tensor(is_synthetic, dtype=torch.float)
+
+    return pixel_values, is_synthetic
+
+
+def load_data():
+    print("loading datasets...")
+
+    diffusion_db_dataset = load_dataset(
+        "poloclub/diffusiondb",
+        "2m_random_50k",
+        trust_remote_code=True,
+        split="train",
+    )
+    flickr_dataset = load_dataset("nlphuji/flickr30k", split="test[:50%]")
+    painting_dataset = load_dataset(
+        "keremberke/painting-style-classification", name="full", split="train"
+    )
+    anime_scene_datatset = load_dataset(
+        "animelover/scenery-images", "0-sfw", split="train"
+    )
+    movie_poaster_dataset = load_dataset("nanxstats/movie-poster-5k", split="train")
+    metal_album_art_dataset = load_dataset(
+        "Alphonsce/metal_album_covers", split="train[:80%]"
+    )
+
+    diffusion_db_dataset = diffusion_db_dataset.select_columns("image")
+    diffusion_db_dataset = diffusion_db_dataset.map(label_fake)
+
+    flickr_dataset = flickr_dataset.select_columns("image")
+    flickr_dataset = flickr_dataset.map(label_real)
+
+    painting_dataset = painting_dataset.select_columns("image")
+    painting_dataset = painting_dataset.map(label_real)
+
+    anime_scene_datatset = anime_scene_datatset.select_columns("image")
+    anime_scene_datatset = anime_scene_datatset.map(label_real)
+
+    movie_poaster_dataset = movie_poaster_dataset.select_columns("image")
+    movie_poaster_dataset = movie_poaster_dataset.map(label_real)
+
+    metal_album_art_dataset = metal_album_art_dataset.select_columns("image")
+    metal_album_art_dataset = metal_album_art_dataset.map(label_real)
+
+    diffusion_split = diffusion_db_dataset.train_test_split(test_size=TEST_SIZE)
+    flickr_split = flickr_dataset.train_test_split(test_size=TEST_SIZE)
+    painting_split = painting_dataset.train_test_split(test_size=TEST_SIZE)
+    anime_scene_split = anime_scene_datatset.train_test_split(test_size=TEST_SIZE)
+    movie_poaster_split = movie_poaster_dataset.train_test_split(test_size=TEST_SIZE)
+    metal_album_art_split = metal_album_art_dataset.train_test_split(
+        test_size=TEST_SIZE
+    )
+
+    training_ds = concatenate_datasets(
+        [
+            diffusion_split["train"],
+            flickr_split["train"],
+            painting_split["train"],
+            anime_scene_split["train"],
+            movie_poaster_split["train"],
+            metal_album_art_split["train"],
+        ]
+    )
+    validation_ds = concatenate_datasets(
+        [
+            diffusion_split["test"],
+            flickr_split["test"],
+            painting_split["test"],
+            anime_scene_split["test"],
+            movie_poaster_split["test"],
+            metal_album_art_split["test"],
+        ]
+    )
+
+    training_ds = training_ds.map(
+        transform_training, remove_columns=["image"], batched=True
+    )
+    validation_ds = validation_ds.map(
+        transform_validation, remove_columns=["image"], batched=True
+    )
+
+    training_loader = torch.utils.data.DataLoader(
+        training_ds,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+        num_workers=4,
+        collate_fn=collate,
+    )
+    validation_loader = torch.utils.data.DataLoader(
+        validation_ds,
+        batch_size=BATCH_SIZE,
+        shuffle=True,
+        num_workers=4,
+        collate_fn=collate,
+    )
+
+    return training_loader, validation_loader, len(training_ds)
+
+
+@app.function(gpu="T4", timeout=86400, volumes={model_store_path: volume})
+def train():
+    training_loader, validation_loader, sample_size = load_data()
+
+    print(f"sample size: {sample_size}")
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    model = mai.cuda()
+    loss_fn = torch.nn.BCEWithLogitsLoss()
+    optimizer = torch.optim.RMSprop(model.parameters(), lr=0.001)
+    best_vloss = 1_000_000.0
+
+    for epoch in range(EPOCHS):
+        print(f"EPOCH {epoch + 1}:")
+
+        model.train(True)
+
+        running_loss = 0.0
+        last_loss = 0.0
+        correct = 0.0
+        total = 0.0
+        accuracy = 0.0
+        i = 0
+
+        for i, data in enumerate(training_loader):
+            inputs, labels = data
+            inputs = inputs.cuda()
+
+            labels = labels.cuda()
+            labels = labels.view(-1, 1)
+
+            optimizer.zero_grad()
+
+            outputs = model(inputs)
+
+            loss = loss_fn(outputs, labels)
+            loss.backward()
+
+            optimizer.step()
+
+            # Calculate accuracy
+            predicted = (outputs > 0.5).float()  # Applying a threshold of 0.5
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+            accuracy = 100 * correct / sample_size
+
+            running_loss += loss.item()
+            if i % 1000 == 999:
+                last_loss = running_loss / 1000  # loss per batch
+                print("  batch {} loss: {}".format(i + 1, last_loss))
+                running_loss = 0.0
+
+        print(f"ACCURACY {accuracy}")
+
+        # validation step
+        running_vloss = 0.0
+        model.eval()
+        with torch.no_grad():
+            for i, validation_data in enumerate(validation_loader):
+                vinputs, vlabels = validation_data
+                vinputs = vinputs.cuda()
+                vlabels = vlabels.cuda()
+                vlabels = vlabels.view(-1, 1)
+
+                voutputs = model(vinputs)
+                vloss = loss_fn(voutputs, vlabels)
+
+                running_vloss += vloss
+
+        avg_validation_loss = running_vloss / (i + 1)
+        print("LOSS train {} valid {}".format(last_loss, avg_validation_loss))
+
+        if avg_validation_loss < best_vloss:
+            best_vloss = avg_validation_loss
+            model_path = f"{model_store_path}/mai_{timestamp}_{epoch}"
+            torch.save(model.state_dict(), model_path)
+            volume.commit()
+
+
+@app.local_entrypoint()
+def main():
+    train.remote()