Buckets:

mishig's picture
|
download
raw
22 kB

TCDScheduler

Trajectory Consistency Distillation от Jianbin Zheng, Minghui Hu, Zhongyi Fan, Chaoyue Wang, Changxing Ding, Dacheng Tao и Tat-Jen Cham представила метод Strategic Stochastic Sampling (Algorithm 4), способный генерировать качественные образцы за небольшое количество шагов. Являясь продвинутой итерацией многошагового планировщика (Algorithm 1) из Consistency Models, Strategic Stochastic Sampling специально адаптирована для функции траекторной согласованности.

Аннотация из статьи:

Модель латентной консистентности (LCM) расширяет модель консистентности на латентное пространство и использует технику направленной дистилляции консистентности для достижения впечатляющих результатов в ускорении синтеза изображений по тексту. Однако мы наблюдали, что LCM испытывает трудности с генерацией изображений, сочетающих в себе чёткость и детализированную сложность. Чтобы преодолеть это ограничение, мы первоначально исследуем и разъясняем лежащие в основе причины. Наше исследование выявляет, что основная проблема проистекает из ошибок в трёх различных областях. Вследствие этого мы предлагаем дистилляцию траекторной консистентности (TCD), которая включает функцию траекторной консистентности и стратегическое стохастическое выборку. Функция траекторной консистентности снижает ошибки дистилляции, расширяя область граничного условия самоконсистентности и наделяя TCD способностью точно отслеживать всю траекторию ODE потока вероятности. Кроме того, стратегическое стохастическое выборку специально разработано для обхода накопленных ошибок, присущих многошаговому выборке консистентности, и тщательно подобрано для дополнения модели TCD. Эксперименты демонстрируют, что TCD не только значительно повышает качество изображений при низком NFE, но и даёт более детализированные результаты по сравнению с учителем при высоком NFE.

Исходный код можно найти в репозитории jabir-zheng/TCD.

TCDScheduler[[diffusers.TCDScheduler]]

diffusers.TCDScheduler[[diffusers.TCDScheduler]]

diffusers.TCDScheduler(num_train_timesteps: int = 1000, beta_start: float = 0.00085, beta_end: float = 0.012, beta_schedule: str = 'scaled_linear', trained_betas: numpy.ndarray | list[float] | None = None, original_inference_steps: int = 50, clip_sample: bool = False, clip_sample_range: float = 1.0, set_alpha_to_one: bool = True, steps_offset: int = 0, prediction_type: str = 'epsilon', thresholding: bool = False, dynamic_thresholding_ratio: float = 0.995, sample_max_value: float = 1.0, timestep_spacing: str = 'leading', timestep_scaling: float = 10.0, rescale_betas_zero_snr: bool = False)

Источник

Параметры:

num_train_timesteps (int, по умолчанию 1000) : Количество шагов диффузии для обучения модели.

beta_start (float, по умолчанию 0.00085) : начальное значение beta для вывода.

beta_end (float, по умолчанию 0.012) : Итоговое значение beta.

beta_schedule (str, по умолчанию "scaled_linear") : Расписание бета, отображающее диапазон значений бета в последовательность бета для шагов модели. Выберите одно из значений: linear, scaled_linear или squaredcos_cap_v2.

trained_betas (np.ndarray или list[float], необязательный) : Передайте массив бета напрямую в конструктор, чтобы обойти beta_start и beta_end.

original_inference_steps (int, по умолчанию 50) : количество шагов инференса по умолчанию, используемое для создания линейно распределённого расписания временных шагов, из которого в итоге мы берём num_inference_steps равномерно распределённых временных шагов для формирования итогового расписания.

clip_sample (bool, по умолчанию False) : Обрезает предсказанную выборку для численной устойчивости.

clip_sample_range (float, по умолчанию 1.0) : Максимальная величина для обрезки выборок. Действительно только при clip_sample=True.

set_alpha_to_one (bool, по умолчанию True) : Каждый шаг диффузии использует значение произведения альф на этом шаге и на предыдущем. Для последнего шага предыдущего альфа не существует. Если этот параметр установлен в True, произведение предыдущего альфа фиксируется в 1, в противном случае используется значение альфа на шаге 0.

steps_offset (int, по умолчанию 0) : Смещение, добавляемое к шагам инференса, как требуется некоторыми семействами моделей.

prediction_type (str, defaults to "epsilon") : Prediction type of the scheduler function; can be epsilon (predicts the noise of the diffusion process), sample (directly predicts the noisy sample) or v_prediction` (see section 2.4 of Imagen Video paper).

thresholding (bool, по умолчанию False) : Использовать ли метод «динамического порога». Этот метод не подходит для моделей диффузии в латентном пространстве, таких как Stable Diffusion.

dynamic_thresholding_ratio (float, по умолчанию 0.995) : Коэффициент для метода динамического порогового значения. Действительно только при thresholding=True.

sample_max_value (float, по умолчанию 1.0) : Пороговое значение для динамического порогового выделения. Действительно только при thresholding=True.

timestep_spacing (str, defaults to "leading") : The way the timesteps should be scaled. Refer to Table 2 of the Common Diffusion Noise Schedules and Sample Steps are Flawed for more information.

timestep_scaling (float, по умолчанию 10.0) : Коэффициент, на который умножаются шаги времени при вычислении граничных условий модели согласованности c_skip и c_out. Увеличение этого значения снижает ошибку аппроксимации (хотя при значении по умолчанию 10.0 ошибка аппроксимации уже достаточно мала).

rescale_betas_zero_snr (bool, по умолчанию False) : Нужно ли перемасштабировать беты для нулевого конечного SNR. Это позволяет модели генерировать очень яркие и тёмные образцы, а не ограничиваться образцами со средней яркостью. Непосредственно связано с --offset_noise.

TCDScheduler реализует Strategic Stochastic Sampling, предложенное в статье Trajectory Consistency Distillation, расширяя исходный Multistep Consistency Sampling для обеспечения неограниченного обхода траекторий.

Этот код основан на официальном репозитории TCD(https://github.com/jabir-zheng/TCD).

Эта модель наследуется от SchedulerMixin и ConfigMixin. ~ConfigMixin отвечает за сохранение всех атрибутов конфигурации, передаваемых в функцию __init__ планировщика, например num_train_timesteps. К ним можно обратиться через scheduler.config.num_train_timesteps. SchedulerMixin предоставляет общие функции загрузки и сохранения с помощью функций SchedulerMixin.save_pretrained() и from_pretrained().

add_noise[[diffusers.TCDScheduler.add_noise]]

add_noise(original_samples: Tensor, noise: Tensor, timesteps: IntTensor)

Источник

Параметры:

original_samples (torch.Tensor) : исходные образцы, к которым будет добавлен шум.

noise (torch.Tensor) : Шум, который нужно добавить к образцам.

timesteps (torch.IntTensor) : Шаги, указывающие уровень шума для каждого примера.

Возвращает: torch.Tensor

Шумные образцы.

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

get_velocity[[diffusers.TCDScheduler.get_velocity]]

get_velocity(sample: Tensor, noise: Tensor, timesteps: IntTensor)

Источник

Параметры:

sample (torch.Tensor) : Входной образец.

noise (torch.Tensor) : Тензор шума.

timesteps (torch.IntTensor) : Шаги для вычисления скорости.

Возвращает: torch.Tensor

Вычисленная скорость.

Вычислите предсказание скорости из выборки и шума согласно формуле скорости.

index_for_timestep[[diffusers.TCDScheduler.index_for_timestep]]

index_for_timestep(timestep: typing.Union[float, torch.Tensor], schedule_timesteps: typing.Optional[torch.Tensor] = None)

Источник

Параметры:

timestep (float или torch.Tensor) : значение шага времени, которое нужно найти в расписании.

schedule_timesteps (torch.Tensor, необязательный) : Расписание временных шагов, в котором выполняется поиск. Если None, используется self.timesteps.

Возвращает: int

Индекс шага времени в расписании. Для самого первого шага возвращает второй индекс, если существует несколько совпадений, чтобы избежать пропуска сигмы при начале с середины расписания (например, для image-to-image).

Найдите индекс заданного шага времени в расписании шагов времени.

previous_timestep[[diffusers.TCDScheduler.previous_timestep]]

previous_timestep(timestep: int)

Источник

Параметры:

timestep (int) : Текущий шаг.

Возвращает: int или torch.Tensor

Предыдущий шаг по времени.

Вычислите предыдущий шаг временной шкалы в цепочке диффузии.

scale_model_input[[diffusers.TCDScheduler.scale_model_input]]

scale_model_input(sample: Tensor, timestep: typing.Union[int, torch.Tensor, NoneType] = None)

Источник

Параметры:

sample (torch.Tensor) : Входной образец.

timestep (int или torch.Tensor, необязательный) : Текущий шаг в цепочке диффузии.

Возвращает: torch.Tensor

Масштабированный входной образец.

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

set_begin_index[[diffusers.TCDScheduler.set_begin_index]]

set_begin_index(begin_index: int = 0)

Источник

Параметры:

begin_index (int, по умолчанию 0) : Начальный индекс для планировщика.

Устанавливает начальный индекс для планировщика. Эта функция должна выполняться из конвейера перед инференсом.

set_timesteps[[diffusers.TCDScheduler.set_timesteps]]

set_timesteps(num_inference_steps: int | None = None, device: typing.Union[str, torch.device, NoneType] = None, original_inference_steps: int | None = None, timesteps: list[int] | None = None, strength: float = 1.0)

Источник

Параметры:

num_inference_steps (int, optional) : Количество шагов диффузии, используемое при генерации образцов с помощью предобученной модели. Если используется, timesteps должно быть None.

device (str или torch.device, необязательный) : Устройство, на которое следует переместить шаги. Если None, шаги не перемещаются.

original_inference_steps (int, optional) : исходное количество шагов инференса, которое будет использовано для генерации линейно распределённого расписания временных шагов (отличного от стандартной реализации diffusers). Затем мы возьмём num_inference_steps временных шагов из этого расписания, равномерно распределённых по индексам, и используем их в качестве итогового расписания временных шагов. Если не задано, по умолчанию используется атрибут original_inference_steps.

timesteps (list[int], optional) : Пользовательские шаги (timesteps) для поддержки произвольного интервала между ними. Если None, используется стандартная стратегия равномерного распределения шагов по расписанию обучения/дистилляции. Если передан timesteps, num_inference_steps должен быть None.

strength (float, по умолчанию 1.0) : Используется для определения количества шагов, применяемых при инференсе в режимах img2img, inpaint и т. д.

Устанавливает дискретные временные шаги, используемые для цепочки диффузии (выполняется перед инференсом).

шаг[[diffusers.TCDScheduler.step]]

step(model_output: Tensor, timestep: typing.Union[int, torch.Tensor], sample: Tensor, eta: float = 0.3, generator: typing.Optional[torch.Generator] = None, return_dict: bool = True)

Источник

Параметры:

model_output (torch.Tensor) : Прямой вывод из обученной диффузионной модели.

timestep (int или torch.Tensor) : Текущий дискретный шаг времени в цепочке диффузии.

sample (torch.Tensor) : Текущая инстанция сэмпла, созданная процессом диффузии.

eta (float, по умолчанию 0.3) : Стохастический параметр (обозначенный как gamma в статье), используемый для управления степенью стохастичности на каждом шаге. При eta = 0 это соответствует детерминированному сэмплированию, а при eta = 1 — полностью стохастическому сэмплированию.

generator (torch.Generator, необязательный) : генератор случайных чисел.

return_dict (bool, по умолчанию True) : Возвращать ли TCDSchedulerOutput или tuple.

Возвращает: TCDSchedulerOutput или tuple[torch.Tensor, torch.Tensor]

If return_dict is True, TCDSchedulerOutput is returned, otherwise a tuple is returned containing the previous sample and the predicted noised sample.

Предскажите образец из предыдущего шага времени, обращая SDE. Эта функция распространяет процесс диффузии от выходов обученной модели (чаще всего — предсказанного шума).

TCDSchedulerOutput[[diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput]]

diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput[[diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput]]

diffusers.schedulers.scheduling_tcd.TCDSchedulerOutput(prev_sample: Tensor, pred_noised_sample: typing.Optional[torch.Tensor] = None)

Источник

Параметры:

prev_sample (torch.Tensor формы (batch_size, num_channels, height, width) для изображений) : Вычисленная выборка (x_{t-1}) для предыдущего шага. prev_sample следует использовать как следующий вход модели в цикле шумоподавления.

pred_noised_sample (torch.Tensor формы (batch_size, num_channels, height, width) для изображений, необязательный) : Предсказанное зашумленное значение (x_{s}) на основе вывода модели на текущем шаге.

Класс результата для вывода функции step планировщика.

Xet Storage Details

Size:
22 kB
·
Xet hash:
2ae3f07cfda1254684dade768db251da8a56de35ee72c90fef0f50154cbf9ef0

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.