Buckets:

mishig's picture
|
download
raw
15.9 kB

TCDScheduler

Trajectory Consistency Distillation, de Jianbin Zheng, Minghui Hu, Zhongyi Fan, Chaoyue Wang, Changxing Ding, Dacheng Tao e Tat-Jen Cham, introduziu o Strategic Stochastic Sampling (Algoritmo 4), capaz de gerar bons exemplos em um pequeno número de etapas. Diferenciando-o como uma iteração avançada do multistep scheduler (Algoritmo 1) dos Consistency Models, o Strategic Stochastic Sampling foi especificamente adaptado para a função de consistência de trajetória.

O resumo do artigo é:

O Latent Consistency Model (LCM) estende o Consistency Model para o espaço latente e utiliza a técnica de destilação de consistência guiada para alcançar um desempenho impressionante na aceleração da síntese de texto para imagem. No entanto, observamos que o LCM tem dificuldade em gerar imagens com clareza e intrincados detalhes simultaneamente. Para abordar essa limitação, inicialmente investigamos e elucidamos as causas subjacentes. Nossa investigação identifica que o problema principal decorre de erros em três áreas distintas. Consequentemente, introduzimos a Trajectory Consistency Distillation (TCD), que abrange a função de consistência de trajetória e a amostragem estocástica estratégica. A função de consistência de trajetória reduz os erros de destilação ao ampliar o escopo da condição de fronteira de autoconsistência e ao conferir à TCD a capacidade de rastrear com precisão toda a trajetória da Probability Flow ODE. Além disso, a amostragem estocástica estratégica foi projetada especificamente para contornar os erros acumulados inerentes à amostragem de consistência em múltiplas etapas, sendo meticulosamente adaptada para complementar o modelo TCD. Os experimentos demonstram que a TCD não apenas melhora significativamente a qualidade da imagem em baixos NFEs, mas também produz resultados mais detalhados em comparação com o modelo professor em altos NFEs.

O código original pode ser encontrado em 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)

Fonte

Parâmetros:

num_train_timesteps (int, padrão 1000) : O número de etapas de difusão para treinar o modelo.

beta_start (float, padrão 0.00085) : O valor inicial de beta para a inferência.

beta_end (float, padrão 0.012) : O valor final de beta.

beta_schedule (str, padrão "scaled_linear") : O agendamento de beta, um mapeamento de um intervalo de beta para uma sequência de betas para o avanço do modelo. Escolha entre linear, scaled_linear ou squaredcos_cap_v2.

trained_betas (np.ndarray ou list[float], opcional) : Passe um array de betas diretamente para o construtor para ignorar beta_start e beta_end.

original_inference_steps (int, padrão 50) : O número padrão de etapas de inferência usado para gerar um cronograma de timesteps linearmente espaçados, do qual, em última análise, extrairemos num_inference_steps timesteps igualmente espaçados para formar o cronograma final de timesteps.

clip_sample (bool, padrão False) : Recorta a amostra prevista para estabilidade numérica.

clip_sample_range (float, padrão 1.0) : A magnitude máxima para o recorte de amostras. Válido apenas quando clip_sample=True.

set_alpha_to_one (bool, padrão True) : Cada etapa de difusão usa o valor do produto de alfas nessa etapa e na etapa anterior. Para a etapa final, não há um alfa anterior. Quando esta opção é True, o produto de alfas anterior é fixado em 1; caso contrário, é usado o valor de alfa da etapa 0.

steps_offset (int, padrão 0) : Um deslocamento adicionado às etapas de inferência, conforme exigido por algumas famílias de modelos.

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, padrão False) : Se deve usar o método de "limiar dinâmico". Isso não é adequado para modelos de difusão em espaço latente, como o Stable Diffusion.

dynamic_thresholding_ratio (float, padrão 0.995) : A proporção para o método de limiar dinâmico. Válido apenas quando thresholding=True.

sample_max_value (float, padrão 1.0) : O valor de limiar para o limiar dinâmico. Válido apenas quando thresholding=True.

timestep_spacing (str, padrão "leading") : A forma como os timesteps devem ser escalados. Consulte a Tabela 2 do Common Diffusion Noise Schedules and Sample Steps are Flawed para mais informações.

timestep_scaling (float, padrão 10.0) : O fator pelo qual os timesteps serão multiplicados ao calcular as condições de fronteira do modelo de consistência c_skip e c_out. Aumentar esse valor reduzirá o erro de aproximação (embora o erro de aproximação no padrão de 10.0 já seja bastante pequeno).

rescale_betas_zero_snr (bool, padrão False) : Se deve reescalar os betas para ter SNR terminal zero. Isso permite que o modelo gere amostras muito claras e escuras, em vez de limitá-lo a amostras com brilho médio. Relacionado de forma solta a --offset_noise.

TCDScheduler incorpora o Strategic Stochastic Sampling introduzido pelo artigo Trajectory Consistency Distillation, estendendo o Multistep Consistency Sampling original para permitir a travessia irrestrita de trajetórias.

Este código é baseado no repositório oficial do TCD(https://github.com/jabir-zheng/TCD).

Este modelo herda de SchedulerMixin e ConfigMixin. ~ConfigMixin cuida de armazenar todos os atributos de configuração que são passados na função __init__ do scheduler, como num_train_timesteps. Eles podem ser acessados via scheduler.config.num_train_timesteps. SchedulerMixin fornece funcionalidade geral de carregamento e salvamento por meio das funções SchedulerMixin.save_pretrained() e from_pretrained().

add_noise[[diffusers.TCDScheduler.add_noise]]

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

Fonte

Parâmetros:

original_samples (torch.Tensor) : As amostras originais às quais o ruído será adicionado.

noise (torch.Tensor) : O ruído a ser adicionado às amostras.

timesteps (torch.IntTensor) : Os timesteps que indicam o nível de ruído para cada amostra.

Retorna: torch.Tensor

As amostras ruidosas.

Adicione ruído às amostras originais de acordo com a magnitude do ruído em cada timestep (este é o processo de difusão direta).

get_velocity[[diffusers.TCDScheduler.get_velocity]]

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

Fonte

Parâmetros:

sample (torch.Tensor) : A amostra de entrada.

noise (torch.Tensor) : O tensor de ruído.

timesteps (torch.IntTensor) : Os passos de tempo para o cálculo da velocidade.

Retorna: torch.Tensor

A velocidade calculada.

Calcule a previsão de velocidade a partir da amostra e do ruído de acordo com a fórmula de velocidade.

index_for_timestep[[diffusers.TCDScheduler.index_for_timestep]]

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

Fonte

Parâmetros:

timestep (float ou torch.Tensor) : O valor do timestep a ser encontrado no schedule.

schedule_timesteps (torch.Tensor, opcional) : O cronograma de passos de tempo a ser pesquisado. Se None, usa self.timesteps.

Retorna: int

O índice do timestep no schedule. Para o primeiro passo, retorna o segundo índice se existirem múltiplas correspondências, a fim de evitar pular um sigma ao iniciar no meio do schedule (por exemplo, para image-to-image).

Encontre o índice de um timestep específico na agenda de timesteps.

previous_timestep[[diffusers.TCDScheduler.previous_timestep]]

previous_timestep(timestep: int)

Fonte

Parâmetros:

timestep (int) : O timestep atual.

Retorna: int ou torch.Tensor

O passo de tempo anterior.

Calcule o passo de tempo anterior na cadeia de difusão.

scale_model_input[[diffusers.TCDScheduler.scale_model_input]]

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

Fonte

Parâmetros:

sample (torch.Tensor) : A amostra de entrada.

timestep (int ou torch.Tensor, opcional) : O passo de tempo atual na cadeia de difusão.

Retorna: torch.Tensor

Uma amostra de entrada escalada.

Garante a intercambiabilidade com agendadores que precisam escalar a entrada do modelo de desruidamento dependendo do passo de tempo atual.

set_begin_index[[diffusers.TCDScheduler.set_begin_index]]

set_begin_index(begin_index: int = 0)

Fonte

Parâmetros:

begin_index (int, padrão 0) : O índice inicial para o agendador.

Define o índice inicial para o agendador. Esta função deve ser executada a partir do pipeline antes da inferência.

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)

Fonte

Parâmetros:

num_inference_steps (int, opcional) : O número de etapas de difusão usadas ao gerar amostras com um modelo pré-treinado. Se usado, timesteps deve ser None.

device (str ou torch.device, opcional) : O dispositivo para o qual os timesteps devem ser movidos. Se None, os timesteps não são movidos.

original_inference_steps (int, opcional) : O número original de etapas de inferência, que será usado para gerar um cronograma de timesteps linearmente espaçados (que é diferente da implementação padrão do diffusers). Em seguida, tomaremos num_inference_steps timesteps desse cronograma, espaçados uniformemente em termos de índices, e usaremos isso como nosso cronograma final de timesteps. Se não for definido, isso será padrão para o atributo original_inference_steps.

timesteps (list[int], opcional) : Passos de tempo personalizados usados para suportar espaçamento arbitrário entre os passos de tempo. Se None, a estratégia padrão de espaçamento, com espaçamento igual entre os passos de tempo no agendamento de passos de tempo de treinamento/destilação, é usada. Se timesteps for passado, num_inference_steps deve ser None.

strength (float, padrão 1.0) : Usado para determinar o número de timesteps usados para inferência ao usar img2img, inpaint, etc.

Define os passos discretos usados para a cadeia de difusão (a ser executado antes da inferência).

etapa[[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)

Fonte

Parâmetros:

model_output (torch.Tensor) : A saída direta do modelo de difusão aprendido.

timestep (int ou torch.Tensor) : O passo de tempo discreto atual na cadeia de difusão.

sample (torch.Tensor) : Uma instância atual de uma amostra criada pelo processo de difusão.

eta (float, padrão 0.3) : Um parâmetro estocástico (referido como gamma no artigo) usado para controlar a estocasticidade em cada etapa. Quando eta = 0, representa amostragem determinística, enquanto eta = 1 indica amostragem totalmente estocástica.

generator (torch.Generator, opcional) : Um gerador de números aleatórios.

return_dict (bool, padrão True) : Se deve ou não retornar um TCDSchedulerOutput ou tuple.

Retorna: TCDSchedulerOutput ou 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.

Preveja a amostra do passo de tempo anterior revertendo o SDE. Esta função propaga o processo de difusão a partir das saídas do modelo aprendido (na maioria das vezes, o ruído previsto).

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)

Fonte

Parâmetros:

prev_sample (torch.Tensor de forma (batch_size, num_channels, height, width) para imagens) : Amostra calculada (x_{t-1}) do passo anterior. prev_sample deve ser usado como entrada do próximo modelo no loop de remoção de ruído.

pred_noised_sample (torch.Tensor de forma (batch_size, num_channels, height, width) para imagens, opcional) : A amostra com ruído previsto (x_{s}) com base na saída do modelo do timestep atual.

Classe de saída para a saída da função step do agendador.

Xet Storage Details

Size:
15.9 kB
·
Xet hash:
8b7c0f015fe4f5eac5176797ea7a594dce238de46a35a50ff8bb229c0748cf07

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