Builders gonna build 👷‍♂️

For a year or so Nir Weingarten and I have been playing with the idea of founding a startup. This April I quit my job, Nir submitted his thesis, and we hopped on the ideation-validation wagon. We're both great believers in AI, and think there are certain use cases that AI can really excel in, such as visual media.

After a while in the ideation process we felt we needed to get our hands dirty, so we decided to explore the current capabilities of diffusion models by building an app, and have some fun on the way :) We saw some cool consumer apps that generate AI portraits of people, and thought we could do much better in terms of end result and of user experience.

This is the story of Selfyx: https://selfyx.com, our AI web app, how we built it, and what we've learned in the process 💪🏻

We'll go over the design, tech stack, ML architecture, costs, time we've spent, dev tools we used, dilemmas we had and decisions we took. Overall this project has been intense, interesting, fun to work on, time and money consuming, and we've learned a lot doing it.

App flow and UX 📱

We built the app around these three pivots:

1. Consistently deliver results which absolutely look like the user, for most users.
2. Have best-in-class, easy user experience, which does not require image cherry-picking: both for user input and app output.
3. Be faster and cheaper than other apps - under 30 minutes turnaround times, O(1$) cost per user.

Good ML models == good training data == bad user experience when uploading images. To solve this we built a robust preprocessing pipeline. Diffusion models output is chaotic and unpredictable, to solve this we built a robust postprocessing pipeline. Inference is slow and costly, to solve this we utilized serverless gpu calls and parallelization.

We came up with the following user flow:

1. The user submits 15 images, they are uploaded to an object storage, and are later used to train the user's personalized AI model.
2. A pre-processing pipeline is initiated on a serverless GPU (More on that later ⬇️) - To create a good user experience, that's also compatible to mobile, we allow the users to upload uncropped images with other people, and our preprocessing pipeline identifies the user out of all of the people in the uploaded images, removes other faces from the image, and crops around the user in each image. Images are then resized to model training resolution (in our case 512x512) and uploaded to object storage.
3. We fine-tune a Stable Diffusion 1.5 model on user images using the Dreambooth method (More on that later ⬇️).
4. We save 3 checkpoints for each user - as each user may have a different "best" checkpoint.
5. We generate images in parallel on multiple serverless GPUs - since the image quality (similarity to the user, general prettiness) varies greatly, we chose to generate lots (~500) of images per user and then choose the best ones.
6. We extract features from generated images - since we need to choose the best images to display for the user, we extract several features from each image so we can later use these features for ranking:
   • Similarity to the user - we used different metrics for how similar the generated images are to the user, mostly facial embedding distances from the original user images.
   • Prettiness score - We want to give the users the best looking images, so we used a model for scoring generated images according to human preferences (More on that later ⬇️).
7. After we have the image features, we rank them using an in-house developed model.
8. The user receives the best 100 images according to the ranker model.

The whole shebang takes around 30 minutes, the longest parts are training ~12 min, and generating the images ~10 min. These can be optimized to ~15 minutes all together, by compiling and quantizing the models, a task that we decided not to complete at this point.

ML stack 🤖

Pre processing 🔧

The biggest factor on the results of the model is the quality of the training data.

Problem is, that getting good training data usually put a set of constraints on the user, such as uploading cropped photos of himself. The process of finding 15 images can be annoying on its own, without having to crop these images! This problem gets even worse when considering mobile users.

We decided to allow the user to upload any image of himself, without cropping, and do the dirty work ourselves.

We needed an extensible solution that would enable us to tag images using models during our processing pipeline. Such as having GPT4o discover certain features in an image, or running embedding models on faces.

We did not find any fitting solution, so we have developed an OS micro-framework (pixel-brain) for processing the data (with models) and storing results in a mongoDB database.

pixel-brain: A micro-framework for processing image data with deep learning models

github.com

It's far from perfect, and we have made a lot of shortcuts to save time, but it's pretty useful. Did someone say technical debt? 🤓

In order to identify the user out of all the people in his uploaded images we embedded all the faces in the images using RetinaFace then performed clustering using DBScan and chose the largest cluster.

Base model selection 🤗

We chose to use Stable Diffusion 1.5 as the base model after a short research (we also looked at SDXL, StableCascade and SD2.0), our reasons are that it has the most support in the open-source community, it is the easiest and fastest to train, and it was much easier to control than the other models. While the bigger models can probably yield better results, we chose not to spend the time and money testing them at this point.

Model checkpoint 📈

The base SD1.5 isn't great at generating realistic images of people, but the OS community offers a lot of great fine-tuned versions of SD1.5, available at Civit AI, and we chose the excellent RealisticVision V5.1 by SG_161222, which was finetuned on realistic data.

Finetuning 🪄

To train the model on user images, we used the Dreambooth method. Dream booth takes some token with little semantic meaning and fine tunes a diffusion model such that this token will represent the learned subject: in our case the user.

The original Dreambooth is performed with a full fine-tune of the model, but after some research we have found that training a LoRA is not only much faster and cheaper, but also generates better results, as a result of the inherent regularization effects. Wait, what is a LoRA? LoRA stands for Low Rank Adaptation, and is a method to efficiently fine tune a model by inserting and training a small intermediate layer between the existing model's layer.

As in any ML task, proper hyperparameter tuning is essential. There are a lot of knobs to turn, but in our case these were the essential ones:

• How many steps to train, we trained for ~2500 steps
• Training batch size - we have found out that despite what is common in the OS community, larger batch sizes are better. We use 10.
• Number of images to train on, we use 15.
• Lora embedding dim, We use 128.
• The dreambooth token chosen, We use both 'ohwx' and 'sks'.

We trained the models using sd-scripts, another amazing OS tool by Kohya Tech on Modal serverless GPU service, which is AMAZING, a single Dreambooth training run is ~12 min on a A100 Modal GPU, which costs ~1$ including the machine. More on Modal later ⬇️.

Generating images 📷✨

ML is a statistical practice, and results are probabilistic. This is especially true for high dimensional tasks such as diffusion models. When it comes to identity preservation there's no free lunch, and no one setting that works for all the users. We chose to tackle this problem by fighting fire with fire: Create a big pool of generated images, generated with different parameters, and use an additional model to pick the best images from that pool. From a probabilistic perspective, we create a meta distribution, and dramatically increase the scope of people we can accurately model.

We have found that there are numerous parameters that are very crucial for getting good results:

1. Prompts - SD models, and SD1.5 in particular, are very sensitive to changes in the positive and the negative prompts. A seemingly small change in the prompt can dramatically affect the results. Good prompts kick @$$, and bad prompts don't preserve the user's identity.
   
2. CFG (Classifier-Free Guidance) scale - this is a hugely important parameter which trades off sample fidelity with diversity. An excellent description can be found here:

Guidance: a cheat code for diffusion models

sander.ai/2022/05/26/guidance.html



3. Checkpoint epoch - we finetuned the models for multiple epochs, and there isn't a one-size-fits-all epoch which works best for all users.

Another problem is figuring out which of the parameters work best for the majority of users.

Thus, our approach is to generate images from multiple combinations of the above parameters and then choose the best images across all generated images.

For generating the images we spun a stable-diffusion-web-ui server, another great OS tool by Kohya Tech, and used its HTTP API - Not very sophisticated, but does the job, like eating cornflakes for dinner 😁.

To speed up generation time we use multiple parallel Modal GPUs to generate the images. We use 4 concurrent A10G for each user for ~10 min which costs about ~0.9$ per user.

Post processing 🚜

Picking out the best 100 images from the ~500 generated is quite a hard problem.

There are multiple factors determining how good an AI portrait is:

1. How similar the generated image is to the user?
2. Does the generated image look generated?
3. How beautiful is the generated image?
4. Is the generated image striking, unique or compelling?

We used an in-house developed model to choose the best images according to the categories above.

App architecture ⚙️

It's amazing how much SW development has advanced in the last couple of years. We've erected our entire stack on light weight and efficient cloud technology, and it took us just over 2 months (!).

• Frontend is purely in react, a big shoutout for Cursor for making such a great editor which really gives you super-powers, great product!!.
• Assets are served via Cloudflare-Pages which is very cheap, fast and easy to use. We started up with Vercel that worked great, but then had some technical issues and didn't manage to reach support, so we migrated.
• The backend is written in python using FastAPI, and served via Heroku, including the PG database. This is another amazing, easy to use and affordable product.
• Data is securely stored in AWS S3, where it is automatically deleted after use.
• CDN - We started off with Cloudinary, since they have an excellent dynamic image transformation API which we use extensively. Sadly this didn't scale well as their storage prices were higher than we thought they might be. So we did a painful migration to Cloudflare-Images, which is very affordable, but lacked certain features which we had to implement on our own.
• All GPU heavy lifting is done on Modal serverless GPUs - Dreambooth training on A100 and all the other stuff on parallel A10G instances. Modal is an AMAZING service, you won't get better and more affordable GPUs elsewhere. They also have GREAT customer service, and a generous credits program. Kudos!!

Modal: High-performance cloud for developers

modal.com



• Apart from the application database we also maintain a mongoDB database for storing image metadata, we use Mongo Atlas for that.
• We use Datadog for logging aggregation - it was a bit of a headache to set up with Modal but hacked around and now get all of our logs in one easy-to-debug place.
• We use Mixpanel for user analytics.
• Payment is handled via stripe.

Closing remarks 🎙️

Building this application taught us a lot about the capabilities of current diffusion models, and how much time and money it takes to put them into production.

In terms of diffusion models, I think that one of the biggest challenges dealing with image generation models is automatically evaluating their output, which is otherwise still too chaotic for use without cherry-picking. It's hard to define metrics on such vague features such as aesthetics and even harder to automatically judge images according to them. This makes it very hard to build an application which consistently improves over time, as the application output is hard to evaluate.

Nonetheless, I believe this problem is solvable, as we showed with Selfyx, and considering that this technology is advancing faster than Nvidia stocks, we expect to see more generated visuals in our daily life in the coming years.

In terms of dev velocity, we see that using the right tools can get you a very long way, very fast.

We have great belief that image generation models will change how we view the world, if we just learn to tame them 🦁.

We would love to hear your thoughts about Selfyx, possible features or anything else! Contact us on X: