Изображение, сгенерированное AI по промпту «photograph of a robot drawing in the wild, nature, jungle» («фотография робота, рисующего в природе, джунглях»)

22 августа 2022 года Stability.AI объявила о публичном релизе Stable Diffusion — мощной диффузионной модели text-to-image. Модель способна генерировать различные варианты изображений на основании текстового или графического ввода.

Стоит заметить, что «модель выпущена под лицензией Creative ML OpenRAIL-M. Лицензия допускает коммерческое и некоммерческое использование. Ответственность за этическое использование модели лежит на разработчиках. Это относится и к производным от неё моделям».

В этой статье я расскажу, как точно настраивать эмбеддинги для создания персонализированных изображений на основании произвольных стилей или объектов. Вместо переучивания модели мы можем представить собственный стиль в виде новых слов в пространстве эмбеддингов модели. В результате этого новое слово будет руководить созданием новых изображений интуитивно понятным образом.

Посмотрите на показанное ниже сравнение реального объекта и изображения, сгенерированного при помощи модели Stable Diffusion и точно настроенного эмбеддинга:


Давайте приступим к подготовке и установке.

Установка

Следует учесть, что этот туториал основан на форкнутой версии Stable Diffusion, в которую интегрирована текстовая инверсия для персонализации генерации изображений. В официальном репозитории рекомендуют использовать для установки Anaconda.

Для установки всех необходимых зависимостей изучите следующее руководство.

Если вы намереваетесь устанавливать все зависимости через pip install, то воспользуйтесь следующим файлом требований:

pip install -r requirements.txt

Если у вас возникнут проблемы с установкой, изучите следующее руководство по устранению проблем.

Убедитесь, что у вас есть все требуемые модели и чекпоинты в рабочей папке или в папке .cache.

Или же можно протестировать это при помощи следующего ноутбука Stable Diffusion AI на Google Colab. Следуйте инструкциям по запуску Stable Diffusion в изолированном окружении.

Подготовка данных

Первым делом в корневом каталоге (stable-diffusion) создайте папку training_data. Мы поместим в неё все наши обучающие изображения. Можно дать им любые имена, но они должны иметь следующие свойства:

• размер изображения: 512 x 512
• объекты на изображениях должны быть правильно повёрнуты
• для наилучших результатов используйте 3-5 изображений (если использовать слишком много изображений, модель может не сойтись)
• изображения должны содержать схожую контекстуальную информацию. Для точной настройки под стиль нужно использовать согласованные цветовые схемы и стиль графики. Для точной настройки под объекты на изображениях должен присутствовать одинаковый объект под разными углами (рекомендуется сохранять одно направление для объектов, передняя и задняя части которых различаются).

Для демонстрации я воспользуюсь своим брелоком в качестве данных обучения для точной настройки под объект.


Показанные выше фото сделаны камерой моего телефона, поэтому мне нужно было уменьшить их до размера 512 x 512. Внутри папки training_data я создал ещё две папки:

• raw — для сырых изображений, требующих изменения размера
• key — подпапка для всех обучающих изображений (512 x 512)

Для изменения размера обучающих изображений вы можете воспользоваться в качестве справки следующим скриптом:

from PIL import Image
import glob

size = (512, 512)

# получаем все файлы в папке, проверяем, что все они являются изображениями
files = glob.glob('./training_data/raw/*')

for fil in files:
    # при необходимости реализуем здесь проверку типа файла
    
    # получаем базовое имя, например "dragon.jpg" -> ("dragon", ".jpg")
    basename = os.path.splitext(os.path.basename(fil))[0]

    with Image.open(fil) as img:
        # изменяем размер до 512 x 512
        img = img.resize(size)
        
        # при необходимости поворачиваем изображение
        # img = img.rotate(90)
        
        # сохраняем изменённое изображение, при необходимости измените способ ресэмплинга и папку для вывода
        img.save(f"./training_data/key/{basename}.png", format="PNG", resample=Image.Resampling.NEAREST)

Конфигурация

Перейдите в папку <root>/configs/stable-diffusion/, в ней вы должны увидеть следующие файлы:

• v1-finetune.yaml
• v1-finetune_style.yaml
• v1-inference.yaml

Файл v1-finetune.yaml предназначен для точной настройки под объекты. Для точной настройки под стиль нужно использовать в качестве файла конфигурации v1-finetune_style.yaml.

Рекомендуется создать резервные копии файлов конфигурации на случай, если вы поломаете конфигурацию.

Стандартная конфигурация требует для обучения не менее 20 ГБ VRAM. Мы можем снизить требования к памяти, уменьшив размер батча и количество воркеров. Измените следующие строки:

data:
  target: main.DataModuleFromConfig
  params:
    batch_size: 2
    num_workers: 16
...
lightning:
  callbacks:
    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 500
        max_images: 8

чтобы они выглядели так:

data:
  target: main.DataModuleFromConfig
  params:
    batch_size: 1
    num_workers: 8
...
lightning:
  callbacks:
    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 500
        max_images: 1

• batch_size с 2 до 1
• num_workers с 16 до 8
• max_images с 8 до 1

Я протестировал новую конфигурацию, она работает на GeForce RTX 2080.

Для уменьшения объёма памяти можно изменить разрешение обучающего изображения в first_state_config. Рекомендую использовать 256 или 512, в зависимости от устройств GPU.

first_stage_config:
      target: ldm.models.autoencoder.AutoencoderKL
      params:
        embed_dim: 4
        monitor: val/rec_loss
        ddconfig:
          double_z: true
          z_channels: 4
          resolution: 256

Кроме того, вам нужно будет изменить initializer_words как часть ввода для эмбеддинга. Он принимает список строк:

# точная настройка на изображение игрушки
initializer_words: ["toy"]
# точная настройка на стиль графики
initializer_words: ["futuristic", "painting"]

Обучение

Закончив с этим, выполните следующую команду:

python main.py --base ./configs/stable-diffusion/v1-finetune.yaml \
    -t \
    --actual_resume ./models/ldm/stable-diffusion-v1/model.ckpt \
    -n my_key \
    --gpus 0, \
    --data_root ./training_data/key/

Укажите --no-test в командной строке, чтобы игнорировать тестирование при точной настройке. Для изменения initializer_words можно использовать аргумент --init_word. Стоит заметить, что он работает только для одной строки. Если вы намереваетесь использовать несколько строк, то нужно изменить их в файле конфигурации.

На момент написания статьи обучение можно выполнять только на одном GPU и после числа при указании GPU нужно поставить запятую.

Если всё настроено правильно, начнётся обучение и будет создана папка logs. Внутри неё будет находиться папка для каждого проведённого вами обучения, и в ней хранятся следующие папки:

• checkpoints — содержит выходной эмбеддинг и чекпоинты обучения
• configs — файлы конфигурации в формате YAML
• images — пример текста промпта, изображение промпта и сгенерированные изображения
• testtube — файлы, связанные с тестированием

Обучение будет продолжаться бесконечно или в зависимости от числа max_steps, заданного в файле конфигурации:

trainer:
  benchmark: True
  max_steps: 4000

Для справки см. v1-finetune.yaml, так как эта конфигурация не включена в v1-finetune_style.yaml. Пропустите её, чтобы обучение выполнялось бесконечно.
Чтобы понять качество результатов работы, проверьте папку images. Изображения с префиксом sample_scaled_ представляют собой сгенерированные сэмплы и если обучение происходит правильно, они должны иметь схожие черты с данными обучения. Можно остановить обучение, когда оно достигнет примерно 5000-7000 шагов.

Внутри папки checkpoints находится довольно большое количество файлов:

• embeddings_gs-5699.pt — файл эмбеддинга на шаге 5699
• epoch=000010.ckpt — чекпоинт обучения на 10 эпохах
• last.ckpt — чекпоинт обучения последней эпохи
• embeddings.pt — файл эмбеддинга последнего шага

Файлы ckpt используются для продолжения обучения. Файлы pt — это файлы эмбеддингов, которые должны использоваться совместно с моделью Stable Diffusion. Просто скопируйте нужный файл эмбеддинга и поместите его в удобное место для получения логического вывода (inference).

Логический вывод

▍ Интерактивный терминал

Для использования в командной строке просто задайте флаг embedding_path при запуске ./scripts/invoke.py.

python ./scripts/invoke.py --embedding_path /path/to/embeddings.pt

Затем используйте в интерактивном терминале следующий промпт для text2image:

invokeai> "a photo of *"

Кроме того, можно выполнить и image2image, задав флаг init_img:

invokeai> "waterfall in the style of *" --init_img=./path/to/drawing.png --strength=0.7

▍ Упрощённый API

Также можно использовать модуль ldm.generate для создания логического вывода программным образом. Он имеет функцию prompt2image как единую точку входа для задач text2image и image2image.

Стоит заметить, что старая версия содержит следующие функции, которые устарели, но всё равно работают:

• txt2img
• img2img

Для получения подробной информации по входным аргументам изучите скрипт на Python.

Полный код логического вывода приведён в этом gist:

from ldm.generate import Generate

# конфигурация модели
gr = Generate(
    'stable-diffusion-1.4',
    embedding_path='./models/embeddings.pt'  # измените путь к эмбеддингам
)

# загрузка модели
gr.load_model()

# инициализация переменной
text = "a photo of *"

# логический вывод возвращает список кортежей
results = gr.prompt2image(
    prompt   = text,
    outdir   = "./outputs/",
    iterations=1,
    steps=50
)

# сохраняем изображение в папку outputs
for row in results:
    im   = row[0]
    seed = row[1]
    im.save(f'./outputs/image-{seed}.png')

Результаты

Стоит заметить, что чаще всего результаты могут быть некачественными. Идея в том, чтобы сгенерировать большое количество изображений и использовать то, которое подходит для ваших целей. Давайте взглянем на результат, который я получил после точной настройки эмбеддинга.

Изображение, сгенерированное AI по промпту «watercolor painting of * with trees» («акварельный рисунок * с деревьями»)

Изображение, сгенерированное AI по промпту «a bagpack in the style of *» («рюкзак в стиле *»)

Изображение, сгенерированное AI по промпту «a shirt with * logo» («футболка с логотипом *»

Кроме того, я протестировал задачи image2image и получил следующие результаты:

Изображение справа сгенерировано по промпту «three * on a beach» («три * на пляже»)

Заключение

Выпуск Stable Diffusion предоставляет огромный потенциал и возможности разработчикам всего мира. Он совершенно точно сильно повлияет на большинство современных отраслей, изменив понимание того, чего можно достичь при помощи AI.

Всего лишь за одну неделю я увидел, как разработчики со всего мира совместно работают над созданием собственных опенсорсных проектов, от UI для веба до анимации.

Я надеюсь увидеть новые прекрасные инструменты и современные модели, созданные опенсорсными сообществами.

Ссылки

1. Stability.AI — Stable Diffusion public release
2. Github — Stable Diffusion in a dream bot style interface
3. Github — Textual Inversion