stable diffusion keras ft
1.0.0
Este repositório fornece código para difusão estável de ajuste fino em Keras. Está adaptado a partir deste roteiro abraçando o rosto. O modelo pré-treinado usado para ajuste fino vem do KERASCV. Para saber sobre o modelo original, consulte esta documentação.
O código fornecido neste repositório é apenas para fins de pesquisa . Confira esta seção para saber mais sobre os possíveis casos de uso e limitações.
Ao carregar este modelo, você aceita a licença Rail-M do CreativEML em https://raw.githubusercontent.com/compvis/stable-diffusion/main/license.
Se você está apenas procurando os recursos que o acompanham deste repositório, aqui estão os links:
Índice :
Este repositório possui um repositório irmão (Keras-SD que serve) que abrange vários padrões de implantação para difusão estável.
ATUALIZAÇÃO 13 de janeiro de 2023 : Este projeto garantiu o 2º lugar na primeira competição de prêmios da comunidade Keras organizada pelo Google.
Após o script original do Hugging Face, este repositório também usa o conjunto de dados Pokemon. Mas foi regenerado para funcionar melhor com tf.data . A versão regenerada do conjunto de dados está hospedada aqui. Confira esse link para obter mais detalhes.
O código de ajuste fino é fornecido no finetune.py . Antes de executar o treinamento, verifique se você possui as dependências (consulte requirements.txt ) instalados.
Você pode iniciar o treinamento com os argumentos padrão executando python finetune.py . Execute python finetune.py -h para saber sobre os argumentos da linha de comando suportados. Você pode ativar o treinamento de precisão mista passando o sinalizador --mp .
Ao iniciar o treinamento, um ponto de verificação do modelo de difusão será gerado em termos de época apenas se a perda atual for menor que a anterior.
Para evitar o OOM e o treinamento mais rápido, é recomendável usar pelo menos uma GPU V100. Usamos um A100.
Alguns detalhes importantes a serem observados :
Detalhes de treinamento :
Nós ajustamos o modelo em duas resoluções diferentes: 256x256 e 512x512. Vários apenas o tamanho do lote e o número de épocas para ajuste fino com essas duas resoluções diferentes. Como não usamos acumulação de gradiente, usamos esse trecho de código para derivar o número de épocas.
python finetune.py --batch_size 4 --num_epochs 577python finetune.py --img_height 512 --img_width 512 --batch_size 1 --num_epochs 72 --mpPara a resolução de 256x256, reduzimos intencionalmente o número de épocas para economizar tempo de computação.
Pesos afinados :
Você pode encontrar os pesos do modelo de difusão ajustados aqui.
O conjunto de dados Pokemon padrão usado neste repositório vem com a seguinte estrutura:
pokemon_dataset/
data.csv
image_24.png
image_3.png
image_550.png
image_700.png
... data.csv parece assim:
Desde que seu conjunto de dados personalizado siga essa estrutura, você não precisará alterar nada na base de código atual, exceto para o dataset_archive .
Caso seu conjunto de dados tenha várias legendas por imagem, você pode selecionar aleatoriamente um do pool de legendas por imagem durante o treinamento.
Com base no conjunto de dados, você pode precisar ajustar os hiperparâmetros.
import keras_cv
import matplotlib . pyplot as plt
from tensorflow import keras
IMG_HEIGHT = IMG_WIDTH = 512
def plot_images ( images , title ):
plt . figure ( figsize = ( 20 , 20 ))
for i in range ( len ( images )):
ax = plt . subplot ( 1 , len ( images ), i + 1 )
plt . title ( title )
plt . imshow ( images [ i ])
plt . axis ( "off" )
# We just have to load the fine-tuned weights into the diffusion model.
weights_path = keras . utils . get_file (
origin = "https://huggingface.co/sayakpaul/kerascv_sd_pokemon_finetuned/resolve/main/ckpt_epochs_72_res_512_mp_True.h5"
)
pokemon_model = keras_cv . models . StableDiffusion (
img_height = IMG_HEIGHT , img_width = IMG_WIDTH
)
pokemon_model . diffusion_model . load_weights ( weights_path )
# Generate images.
generated_images = pokemon_model . text_to_image ( "Yoda" , batch_size = 3 )
plot_images ( generated_images , "Fine-tuned on the Pokemon dataset" ) Você pode trazer seus weights_path (deve ser compatível com o diffusion_model ) e reutilizar o snippet de código.
Confira este notebook Colab para jogar com o código de inferência.
Inicialmente, ajustamos o modelo em uma resolução de 256x256. Aqui estão alguns resultados, juntamente com as comparações com os resultados do modelo original.
| Imagens | Avisos |
|---|---|
| Yoda | |
| gato robótico com asas | |
| Hello Kitty |
Podemos ver que o modelo de ajuste fino possui saídas mais estáveis que o modelo original. Embora os resultados possam ser esteticamente melhorados muito mais, os efeitos de ajuste fino são visíveis. Além disso, seguimos os mesmos hiperparâmetros de abraçar o script do rosto para a resolução de 256x256 (além do número de épocas e tamanho do lote). Com melhores hiperparâmetros, os resultados provavelmente melhorarão.
Para a resolução de 512x512, observamos algo semelhante. Então, experimentamos o parâmetro unconditional_guidance_scale e notamos que, quando está definido como 40 (enquanto mantém os outros argumentos fixos), os resultados saíram melhor.
| Imagens | Avisos |
|---|---|
| Yoda | |
| gato robótico com asas | |
| Hello Kitty |
NOTA : O ajuste fino no 512x512 ainda está em andamento até o momento. Mas leva muito tempo para concluir uma única época sem a presença de treinamento distribuído e acúmulo de gradiente. Os resultados acima são do ponto de verificação derivado após a 60ª época.
Com uma receita semelhante (mas treinada para mais etapas de otimização), os Lambda Labs demonstram resultados surpreendentes.
