Creating AvatarAI.me using Supabase and Cerebrium

@levelsio created AvatarAI.me in November that creates AI generated avatars based on photos of you. You can create profile pictures, pet portraits, professional LinkedIn photos and much more!

This comes off the back of a machine learning model called Dreambooth which is a fine-tuned text-to-image diffusion model for subject driven generation. This means Dreambooth can generate images based on a subject such as a person, a dog or a fictional character. From initial release, AvatarAI.me quickly caught traction generating in excess of $200k in just 1 month! @Levelsio does have 200k followers which did make distribution slightly easier but it is still impressive nonetheless.

Due to its traction we thought we would recreate a version of AvatarAI using our simple model deployment framework and various features from Supabase that we have been wanting to try.

Setting up the Supabase Application

To start, we need to setup our application in Supabase so that we can train our Dreambooth model based on the images in our storage bucket. From Supabase, we will make use of Auth, Database, Storage and Edge functions in order to build our application.

You will need to create an account on Supabase but luckily they have a very generous free tier. We will start by using their React Quickstart which sets up Auth, Database and storage functionality for us. Follow the instructions up to the end of “Initialize React App” and stop before “Set up Login component”. We will not be setting up Auth in this tutorial since it is out of scope.

In your Supabase dashboard, go to your table editor and click the button “New table”. Give it the name filePaths and enable “Row level security” as well as “Realtime”. Insert the following columns into the new table:

id: int8 (primary key)
created_at: timestampz
filePath: text

We will store the location of our generated images in this table and use the Supabase realtime function to display them on our frontend application.

You would also need to specify the correct RLS (Row Level Security)policies for both images and your database. We defined very simple rules — you can select and insert as long as you are using the ANON key from your dashboard since we use this across our application. Below is an example of us doing it for images.

Setup and train your Dreambooth model

As we mentioned, Dreambooth performs particularly well generating images around a specific subject. However, we first need to train the model on images of our subject. We created a simple Git repository to do most of the heavy lifting when it comes to training a Dreambooth model. You can clone the repo here to a private repository since you will have credentials stored in the files. You may follow the instructions in the README of the repo or simply follow them below!

Create an instance on Lamdba Labs or AWS/GCP etc. We opted for Lambda Labs since it has the cheapest GPUs available. We recommend using the 1x Nvidia A100.
SSH into that machine and create and SSH keypair by running: ssh-keygen -t rsa -b 4096
We now need to add our generated SSH key to our repository so we can run our repo code on our instance. Run cat ~/.ssh/id_rsa.pub and add the key to https://github.com/settings/ssh/new.
Run git clone <YOUR_REPO_URL> on the instance you just setup. By adding your ssh key from your instance to the repo, it should be authenticated.
Make sure you are in the auto-stable-diffusion folder on the instance. cd auto-stable-diffusion
In the folder run sudo chmod +x setup.sh to make the setup script executable.
You will need to create a profile on Hugging Face (its free) and can get your token here. Make sure you create/use a token that has write access. You will need to add your Hugging Face token to the setup.sh file. We will use this token to upload your model to Hugging Face once training has completed.
Run source ./setup.sh
Upload the images you wish to train the Dreambooth model on to a folder called “images/unprocessed” in Supabase. You can do this through the Supabase web platform. Please upload a minimum of 5 images!

You will see in the unprocessed folder is a folder called 12345678. This is on purpose and trains Dreambooth to relate the subject to this number. For example, if I entered the prompt “12345678 in a Christmas hat” it would show results of me in a Christmas hat. If you generate multiple models each with different subjects then you can use a unique folder id to refer to different subjects.

9. You will also need to update the download_images.py and upload_images.py file with your Supabase credentials. You can grab your credentials from settings->Api on your Supabase profile.

10. Run ./run-all.sh

It will now train the model on the images you uploaded and upload the results to Supabase. You should see a result similar to the image below. In order to get it working with Cerebrium, you need to upload it to Huggingface (this is a temporary workflow).

You should see your model available in Hugging Face under your profile.

We know the experience isn’t great to train the Dreambooth model since you have to do a lot of setup of training instances. We are looking at including this functionality in the Cerebrium framework where you will be able to train your model using 1 line of code. If you would like to be notified of when we release this functionality, you can enter your email here.

Deploying on Cerebrium

Cerebrium allows you to deploy ML models to serverless GPUs using one line of code. We automatically handle model versioning, A/B testing, integration with monitoring tools and much much more.

To start you need to create an account here in order to get an API key that links all deployments to your account.

Next, install our Python framework.


pip install cerebrium

Deploy your custom Dreambooth model

‍


from cerebrium import model_type, deploy

deploy((model_type.PREBUILT, 'dreambooth-custom'), 'my-custom-model', )

‍

In the above, we specify that we want to use the prebuilt configuration for Dreambooth models, give the deployment a name on our dashboard and then enter the API key from our dashboard. Once run, you should receive the endpoint for inference of your deployment.

‍

For Dreambooth, you can define the following parameters:

model_id: This is the Hugging Face id of your model repo. This is a required parameter
hf_token: This is the token from your HuggingFace profile in order to access your model repo. This is a required parameter.
height: The height of the generated image which defaults to 512px. This is an optional parameter
width: The width of the generated image which defaults to 512px. This is an optional parameter
num_inference_steps: The number of steps you would like Dreambooth to take to generate the image. In general, the results are better the more steps you use, however the more steps, the longer the generation takes. It defaults to 50 steps and is an optional parameter.
guidance_scale: A way to increase the adherence to the conditional signal that guides the generation (text, in this case) as well as overall sample quality. In simple terms it forces the generation to better match the prompt potentially at the cost of image quality or diversity. This defaults to 7.5 and is an optional parameter.
num_images_per_prompt: The number of image variations you would like the model to generate. This is an optional parameter which defaults to 1.
negative_prompt: The negative prompt is a parameter that tells Stable Diffusion what you don’t want to see in the generated images. This is an optional parameter that has an empty default value.
webhook_endpoint: A url that our model can notify once the images have finished generating. We recommend using this parameter as depending on the number of images you want to generate, the REST endpoint could time out after 1 minute. This is an optional parameter.


curl --location --request POST '' \
--header 'Content-Type: application/json' \
--header 'Authorization': '' \
--data-raw '{
    "prompt": "A rabbit on the moon surfing",
    "hf_token": "",
    "model_id": "",
    "num_images_per_prompt": 8,
    "height": 512,
    "width": 512
}'

Setting up the Supabase Application

Now that our model is trained and deployed, we need to tie everything together. We initially did some Supabase setup so now we just need to setup some edge functions and create our frontend React application.

We have to setup two edge functions:

A webhook endpoint that Cerebrium will call once the image(s) have finished generating.
A function to make the API request to Cerebrium in order to generate our image(s)

Image Callback Webhook

Below is the code for our image webhook that Cerebrium will notify when the images from our Dreambooth model have finished generating. The reason we use an image webhook is that depending on the number of images we generate and the number of training steps we select, it could take some time which means our requests could time out.

We will receive base64 encoded strings of the generated images from Cerebrium, convert them into PNG’s and store them in our storage bucket. We will then update the database with the url’s of the newly added URL’s. You will understand when we get to the frontend why we are adding the URL’s to our database.


// Follow this setup guide to integrate the Deno language server with your editor:
// https://deno.land/manual/getting_started/setup_your_environment
// This enables autocomplete, go to definition, etc.

import { serve } from "https://deno.land/std@0.131.0/http/server.ts"
import { createClient } from 'https://esm.sh/@supabase/supabase-js@2.0.0'
import { decode }  from "https://esm.sh/base64-arraybuffer@1.0.2?target=deno&deno-std=0.132.0";

serve(async (req: Request) => {

  const supabaseClient = createClient(
    // Supabase API URL - env var exported by default.
    Deno.env.get('SUPABASE_URL') ?? '',
    // Supabase API ANON KEY - env var exported by default.
    Deno.env.get('SUPABASE_ANON_KEY') ?? '',
    // Create client with Auth context of the user that called the function.
    // This way your row-level-security (RLS) policies are applied.
    { 
      realtime: {
        params: {
          eventsPerSecond: 10
        }
      },
      db: {
        schema: 'public',
      }
    },
    
  )

  const { images } = await req.json()
  if (!images) { 
    console.log("No images returned"); 
    return new Response(
      JSON.stringify({"message": "No images returned"}),
      { headers: { "Content-Type": "application/json" } },
    )
  }

  try {
    var filePaths: string[] = []
    for (const image of images) {
      const uuid = Math.floor(Math.random() * 1000000)
      const filePath = `processed/${uuid}.png`
      filePaths.push(filePath)

      let { error: uploadError } = await supabaseClient.storage
        .from('images')
        .upload(filePath, decode(image), {
          contentType: 'image/png'
        })

      if (uploadError) {
        throw uploadError
      }

      const { error: insertError } = await supabaseClient.from('filePaths')
        .insert({ filePath: filePath  }).single()

      if (insertError) {
        throw insertError
      }                       
      
    }
    
  } catch (error) {
    console.log('Error uploading images!')
    console.log(error)
    return new Response(
      JSON.stringify({"message": "Error uploading images"}),
      { status: 500,
        headers: { "Content-Type": "application/json" } },
    )
  }
  return new Response(
    JSON.stringify({"message": "Success"}),
    { headers: { "Content-Type": "application/json" } },
  )

})

‍

Setting up the Supabase Application

We have to setup two edge functions:

A webhook endpoint that Cerebrium will call once the image(s) have finished generating.
A function to make the API request to Cerebrium in order to generate our image(s)

Image Callback Webhook

‍

In order to run this locally, you need to create a .env.local file that contains the SUPABASE_URL and SUPABASE_ANON_KEY variables. These are the same credentials you used above for training the Dreambooth model.

In order to deploy this webhook, run the command: supabase functions deploy image-webhook — no-verify-jwt .We use the flag no-verify-jwt since otherwise Supabase requires you to have an Authorization header. This is not neccessary for our use case.

Cerebrium Request

We then need to create a function that will make a request to our Cerebrium model endpoint. Create another Supabase function and then copy the code below:

For the above code snippet, you will need to enter the URL of your image-webhook url. You can find the URL on your Supabase dashboard under Edge functions. It should look similar to https://project-id.functions.supabase.co/image-webhook. You deploy this function using the same command we used above supabase functions deploy cerebrium_endpoint .


import { serve } from "https://deno.land/std@0.131.0/http/server.ts"

const corsHeaders = {
  'Access-Control-Allow-Origin': '*',
  'Access-Control-Allow-Headers': 'authorization, x-client-info, apikey'
}

serve(async (req) => {

  if (req.method === 'OPTIONS') {
    return new Response('ok', { headers: corsHeaders })
  }

  const { prompt, num_images } = await req.json()
  const result = await fetch("", {
      method: 'POST',
      headers: { 
        'Content-Type': 'application/json'
        'Authorization': 
      },
      body: JSON.stringify({ 
        prompt: prompt,
        "hf_token": "",
        "model_id": "",
        "webhook_endpoint": "",
        "num_images_per_prompt": num_images})
    })

  return new Response(
    JSON.stringify({"message": "Success"}),
    { headers: { "Content-Type": "application/json", 'Access-Control-Allow-Origin': '*' } },
  )
})

Frontend

You can take a look at our design in our Git repository here. However, the most important element to note in our frontend application is our use of the Supabase realtime function. Our frontend listens to changes specifically in the filePaths table and displays an image every time a file URL gets added to the database. We also make the request to the Cerebrium endpoint when a user clicks the “generate”

To run the frontend application run yarn start .

Frontend Deployment

Supabase has a great integration with Vercel to deploy your application just by connecting your git repo. In some instances, it might have a “deploy” button in the bottom tab of your website.

You can create an account on Vercel here and simply connect your Github repo to deploy your application. In your settings you can define which branch is your production branch. Every time you push changes to the master branch Vercel will automatically build and deploy a new version of you website in a couple of minutes.

Once your application is deployed. You need to go to your Supabase dashboard and navigate to Authentication and then URL configuration. You then need to add the URL of your Vercel application such as “https://my-vercel-website.vercel.app/”.

Final Product

You have now created your own custom Dreambooth model! You can play with our live version here. There are many enhancements you can make to the product and the quality of images by increasing the number of images and steps used for training.

Review

So to recap we had the following elements of our project in order to create our AvatarAI.me replica.

Cerebrium framework: This allowed us to deploy our custom Dreambooth model to Serverless GPU’s using one line of code.

Training repo: The code used to train our custom dreambooth model on a specific subject and deploy it to Hugging Face.

Supabase application:

Storage: This was used to store images to train on and the finished generated images.
Edge functions: This was used as a image webhook as well as to do the rest API endpoint call to the Cerebrium API.
Database: We used the database to keep track of the location of generated images and used the realtime feature to update our frontend in realtime.

‍

arrow_back

Back to blog