HiLo-Token: адаптивное к входным данным сжатие токенов высоких и низких частот для эффективного редактирования изображений

Аннотация

Креативные инструменты редактирования изображений, такие как кнопки «Удалить» или «Генеративная заливка» в Photoshop, являются ключевыми для повседневного использования пользователями и составляют значительную долю трафика в Photoshop и Lightroom. Однако современные генеративные модели ИИ сталкиваются с серьезными проблемами задержки, которые становятся еще более выраженными при переходе от сверточных U-Net к диффузионным трансформерам (DiT). В нашей оценке на сотнях репрезентативных образцов редактирования изображений, охватывающих широкий диапазон коэффициентов маскирования, один только модуль DiT в среднем составляет 73% от общей задержки модели, даже после дистилляции с 50 временных шагов до 8. Для решения этой проблемы мы предлагаем HiLo-Token — адаптивную к входным данным структуру сжатия токенов, которая выделяет больше токенов для высокочастотных областей с богатым контекстом и меньше токенов для низкочастотных областей. В частности, для области редактирования, заданной пользовательской маской, мы сохраняем все токены в пределах расширенной маски для обеспечения сильной локальности и контекстуальной релевантности. За пределами области редактирования мы вводим простую, но эффективную стратегию выбора высокочастотных токенов на основе пространственной частоты для захвата важных локальных деталей, одновременно используя токены из изображения, уменьшенного в 16 раз, для представления низкочастотных компонентов и сохранения размытой, но глобальной структуры. Обширные эксперименты на производственных оценочных данных подтверждают эффективность предложенного метода, достигая ускорения DiT в 3,13x, 2,59x и 1,67x на A100-80GB для задач редактирования изображений с малыми, средними и большими коэффициентами маскирования (средние значения 6,38%, 15,92% и 35,36% соответственно) без какого-либо ухудшения качества генерации.

English

Creative image editing tools, such as Photoshop's Remove or Generative Fill buttons, are central to everyday customer use and account for a major share of traffic in Photoshop and Lightroom. However, current generative AI models face significant latency challenges, which become even more pronounced when transitioning from convolution-based U-Nets to Diffusion Transformers (DiTs). In our evaluation on hundreds of representative image editing samples spanning a wide range of mask ratios, the DiT module alone accounts for an average of 73% of the total model latency, even after being distilled from 50 timesteps down to 8 timesteps. To tackle this challenge, we propose HiLo-Token, an input-adaptive token compression framework that allocates more token budget to high-frequency, rich-context regions while assigning fewer tokens to low-frequency areas. Specifically, for the editing region specified by the user mask, we retain all tokens within a dilated mask to preserve strong locality and contextual relevance. Outside the editing region, we introduce a simple yet effective high-frequency token selection strategy based on spatial frequency to capture important local details, while using tokens from a 16x downsampled image to represent low-frequency components and preserve the blurry but global structure. Extensive experiments on production-level evaluation data validate the effectiveness of the proposed method, achieving 3.13x, 2.59x, and 1.67x DiT speedups on A100-80GB for image editing tasks across small, medium, and large mask ratio categories with average ratios of 6.38%, 15.92%, and 35.36%, respectively, without any regression in generation quality.