Personalized headshot

01 · CONTEXT

The problem: a model that knows a person

Chapter ten of the book Generative AI System Design Interview is titled Personalized Headshot Generation, and the book gathers ten real interview questions about generative AI systems, with detailed solutions.

A personalized headshot service receives a few photos of a person and returns professional portraits of that same person. In an interview, the challenge of machine learning systems is to combine per-user training with a job queue and consent control.

According to the DreamBooth authors, large text-to-image models lack the ability to mimic the appearance of subjects from a reference set.

02 · PERSONALIZATION

DreamBooth and LoRA: teaching a person to the model

DreamBooth fine-tunes a pretrained text-to-image model with a few images of a subject, so that it learns to bind a unique identifier to that subject. With the subject embedded in the model, the identifier can be used to synthesize realistic images of the subject in different scenes.

In the paper, a set of three to five images of the subject is enough for good results, and training takes about five minutes on an NVIDIA A100 with Stable Diffusion.

The authors point out that fine-tuning text-conditioned layers gives rise to language drift. To prevent it, they propose a class-specific prior preservation loss.

LoRA freezes the pretrained model weights and injects trainable rank decomposition matrices into each layer of the Transformer architecture, greatly reducing the number of trainable parameters for downstream tasks.

The LoRA authors focused on Transformer language models, but they write that the principles are generally applicable to any neural network with dense layers. They also describe creating many custom models that can be swapped in real time on machines that keep the pretrained weights in VRAM.

03 · ARCHITECTURE

The end-to-end pipeline

The diagram below shows a proposal with filas and separate workers for training each person's adapter and for generating the portraits.

1. In the diagram, the browser is connected to the API gateway, and the gateway to the order service.

2. The order service is connected to the photo storage and to the orders and consents database.

3. The training workers are connected to the photos and to the adapter storage. DreamBooth describes fine-tuning a pretrained model with a few images of a subject.

4. The generation workers are connected to the adapter storage. LoRA describes a shared pretrained model used to build many small modules for different tasks.

5. The generated portraits storage is connected to the generation workers and to the order service.

04 · GPU

Scheduling GPUs for training and generation

In the design below, training and generation have separate queues and worker groups, in the style of a job-oriented serverless execution.

1. The order service is connected to two separate queues, the training one and the generation one.

2. Each queue feeds its own group of GPU workers.

3. The DreamBooth paper reports about five minutes of training on an NVIDIA A100 with Stable Diffusion.

4. The two worker groups are connected to the adapter storage, the adapter registry, and monitoring.

05 · CONSENT

Consent, withdrawal, and erasure

The GDPR text says that, when processing is based on consent, the controller must be able to demonstrate that the data subject has consented (art. 7).

The same article establishes that it must be as easy to withdraw consent as to give it.

Art. 17 gives the data subject the right to obtain erasure of personal data without undue delay and lists, among the grounds, the withdrawal of the consent on which the processing is based.

In the design, the orders database holds each order's consent record, and a deletion flow removes the person's photos and adapter. This text is study material, not legal advice, and the requirement for conformidade depends on the market served.