Готовы ли вы вывести свою модель LipSync PyTorch на новый уровень и преобразовать ее в модель CoreML?

Большой! В этом сообщении блога мы расскажем вам, как преобразовать модель и подготовить ее к интеграции с приложением iOS. Итак, давайте погрузимся!

Прежде всего, давайте начнем с некоторой справочной информации. Модель, с которой мы будем работать, называется Wav2Lip, и ее можно использовать для создания видео с синхронизацией губ. Вы можете найти более подробную информацию и демо на официальном сайте по адресу https://bhaasha.iiit.ac.in/lipsync/.

Для начала, вот пошаговое руководство по преобразованию модели:

Шаг 1: Клонируйте репозиторий Начните с клонирования репозитория Wav2Lip с GitHub. Вы можете найти его на https://github.com/Rudrabha/Wav2Lip. После клонирования перейдите в папку Wav2Lip.

Шаг 2: Подготовьте зависимости Чтобы все работало гладко, убедитесь, что у вас установлены все необходимые зависимости. Если вы запускаете это на своем локальном компьютере с процессором, все в порядке. Однако, если вы предпочитаете использовать Google Colab, репозиторий предоставляет блокнот Colab для этой цели.

Шаг 3. Загрузите предварительно обученные модели. Вам необходимо загрузить две предварительно обученные модели. Сначала загрузите модель распознавания лиц с https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth и сохраните ее в папке face_detection/detection/sfd/. Затем загрузите основные предварительно обученные веса модели с https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/Eb3LEzbfuKlJiR600lQWRxgBIY27JZg80f7V9jtMfbNDaQ?e=TBFBVW. Вы также можете найти альтернативные ссылки для скачивания на странице GitHub.

Шаг 4. Запуск модели Теперь пришло время запустить модель и создать демонстрационное видео. Откройте терминал или командную строку и выполните следующую команду:

python inference.py --checkpoint_path <ckpt> --face <video.mp4> --audio <an-audio-source>

Обязательно замените <ckpt> на путь к загруженным весам предварительно обученной модели, <video.mp4> на видео, которое вы хотите синхронизировать по губам, и <an-audio-source> на аудиофайл (желательно в формате WAV) длительностью 20 секунд или меньше.

Когда вы будете удовлетворены сгенерированным демо-видео и убедитесь, что модель работает должным образом, пришло время приступить к преобразованию модели CoreML.

Преобразование модели CoreML:

Шаг 1. Импортируйте необходимые библиотеки. В начало файла inference.py добавьте следующий оператор импорта:

from coremltools import ct

Если у вас не установлен coremltools, вы можете сделать это, запустив pip install coremltools в терминале или в командной строке.

Шаг 2. Найдите входные тензоры в коде В файле inference.py найдите строку #260. Здесь вы увидите входные тензоры, которые ожидает модель, в частности тензоры изображения и звука. Под этими строками вы добавите код для преобразования модели CoreML.

Вот пример того, как должен выглядеть код конверсии:

img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)        
traced_model = torch.jit.trace(model, (mel_batch, img_batch))    
print("This is the traced model:")    
print(traced_model)    
coreml_model = ct.converters.convert(
    traced_model,
    inputs=[
        ct.TensorType(name='mel_batch', shape=mel_batch.shape),
        ct.TensorType(name='img_batch', shape=img_batch.shape)
    ],
    minimum_deployment_target=ct.target.iOS15
)
coreml_model.save("model.mlpackage")

Приведенный выше код создает TorchScript-представление исходной модели PyTorch, отслеживает ее, а затем преобразует в модель CoreML. Мы используем функцию convert из coremltools и указываем входные формы, которые ожидает модель. В этом случае мы предоставляем формы тензоров mel_batch и img_batch.

Мы также установили minimum_deployment_target=ct.target.iOS15 для создания программы ML вместо модели нейронной сети. Программа ML меньше по размеру и не зависит от своего веса, что делает ее предпочтительным типом, если вы ориентируетесь на iOS 15 и выше. Это уменьшает размер модели как минимум вдвое по сравнению с моделью нейронной сети.

После запуска кода преобразования вы должны увидеть файл model.mlpackage, сохраненный в основном каталоге репозитория Wav2Lip.

Теперь давайте перейдем к изучению преобразованной модели через Xcode. Вот как это выглядит:

Давайте загрузим модель CoreML с диска, распечатаем ее и посмотрим, как выглядят входные и выходные тензоры:

import coremltools as ct

model = ct.models.MLModel('/path/to/model/Wav2Lip.mlpackage')
print(model)

Большой! Давайте внимательно посмотрим:

input {
  name: "mel_batch"
  type {
    multiArrayType {
      shape: 102
      shape: 1
      shape: 80
      shape: 16
      dataType: FLOAT32
    }
  }
}
input {
  name: "img_batch"
  type {
    multiArrayType {
      shape: 102
      shape: 6
      shape: 96
      shape: 96
      dataType: FLOAT32
    }
  }
}
output {
  name: "var_1049"
  type {
    multiArrayType {
      dataType: FLOAT32
    }
  }
}
metadata {
  userDefined {
    key: "com.github.apple.coremltools.source"
    value: "torch==2.0.1"
  }
  userDefined {
    key: "com.github.apple.coremltools.version"
    value: "7.0b1"
  }
}

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

1- Обнаружение лиц: модель использует библиотеку обнаружения лиц для обнаружения лиц в данном видео. Он сохраняет эти изображения лиц в виде массива 96x96 пикселей. Изначально перед распознаванием лиц размер видео изменяется до 256x256 пикселей. Этот шаг обнаружения лиц выполняется другой библиотекой еще до того, как модель получит входные данные.

2- Предварительная обработка аудио: аудиофайл проходит определенные этапы предварительной обработки, которые вы можете изучить, изучив код в файле Python.

3- Входные данные модели: модель принимает массив лиц с изображениями 96x96 пикселей и 6 каналов (не 3 обычных) в качестве одного из своих входных данных. Количество граней соответствует количеству кадров.

4- Аудиовход: модель также ожидает тензор, содержащий предварительно обработанные аудиоданные.

5- Выходные данные модели. Модель генерирует массив повторно выровненных лиц с изображениями 96x96, где области рта перестраиваются так, как если бы они были синхронизированы по губам с соответствующими данными аудиоспектрограммы. На выходе получается массив выровненных лиц с тем же количеством кадров, что и на входе.

6- Постобработка: файл Python выполняет постобработку, реконструируя выходные кадры лица с исходными видеокадрами и добавляя исходный аудиофайл к видео. Эти шаги выполняются вне модели, в результате чего получается финальное видео размером 256x256 пикселей.

Заключительные примечания

Делает ли это модель хорошей, особенно для обработки в реальном времени? Простой ответ: нет, я так не думаю. Отдельные этапы предварительной и последующей обработки в идеале должны быть интегрированы как внутренние слои исходной модели с оптимизацией для повышения производительности.

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

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

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

Понимая связанные с этим сложности и изучая различные подходы, вы сможете найти наилучшее решение, соответствующее вашим конкретным потребностям.

Удачи вам в преобразовании модели LipSync PyTorch в CoreML, и я надеюсь, что этот пост в блоге был полезен в вашем путешествии!

Где скачать? Загрузите полностью преобразованную модель CoreML здесь.

По любым вопросам или запросам обращайтесь к нам по электронной почте [email protected]. Мы просим только серьезные запросы.