Pitch: Criei uma IA para autocompletar/gerar prompts para Stable Diffusion e Dall-E 2

Oi, olá! Eu sei que esse é meu Terceiro post aqui sobre AIs que criam imagens, mas eu simplesmente viciei nesse assunto, ao ponto de só trabalhar nos ultimos dias em projetos só envolvendo essas AIs e estou gostando bastante!

Para começar, eu estava aqui no TabNews scrollando pelas postagens, quando vi um post da incrível Flávia Carvalho, onde ela fala sobre uma Startup onde as pessoas podem vender Prompts para Stable Diffusion, Dall-E 2 e MidJourney. No final do artigo, ela cita um designer que criou um modelo que consegue gerar também prompts para obter bons resultados. Foi ai que pensei: "E se eu criasse um modelo que faz a mesma coisa, só que open source?"

De ínicio pensei em criar um único modelo, porém vi que cada AI de geração de imagem tem a sua própria forma de interpretar prompts. Exemplo: Se você usar um prompt no Dall-E pode ser que ele não funcione bem no Stable Diffusion e vice-versa. Então resolvi tentar fazer uma série de modelos que nomeei "MagicPrompt".

Daí, começei a tentar criar meus datasets para treinar os modelos. Para o Stable Diffusion, eu fiz um scraping no site: "lexica.art" (que curiosamente também conheci graças ao TabNews), mas isso foi bem difícil de fazer, já que o site ainda não tinha API e era protegido por cloudflare, daí tive que apelar para coisas como o Undetected Chromedriver e no final, eu consegui coletar mais ou menos 80.000 prompts. Você pode dar uma olhada no dataset clicando aqui.

Para coletar prompts do Dall-E foi bem mais difícil, já que o Dall-E 2 é de uma plataforma fechada e tudo o que temos são o que os usuários compartilham publícamente. Mas consegui encontrar uma boa base de prompts no Web Archive, no subreddit do Dall-E e no dalle2.gallery, totalizando mais ou menos 27.000 prompts.

O próximo passo foi treinar os modelos. para isso eu usei Pytorch e como base eu usei o modelo médio do GPT2, lançado pela OpenAI. Aqui estão alguns resultados:

Treinando o modelo de geração para Stable Diffusion:

Eu treinei o modelo para Stable Diffusion com 150.000 steps e com uma perda média de 0.3 e isso levou cerca de algumas horas para finalizar. Para fins de demonstração eu criei um Space no Hugging Face onde você pode testar o modelo clicando aqui. Também criei um infográfico mostrando como o modelo gera prompts, para isso, usei o exemplo: Landscape of...

Treinando o modelo de geração para Dall-E 2:

Eu treinei o modelo do Dall-E com cerca de 40.000 steps e com uma perda média de 0.2, mas acho que ele não é tão preciso quando o anterior, talvez por ter uma quantidade de prompts como base ainda baixa. Eu também levei algumas horas para finalizar o treinamento. Também criei um Space onde você pode testar o modelo clicando aqui, segue um exemplo onde usei o mesmo prompt anterior como base:

Também pretendo mais para o futuro talvez criar um modelo para MidJourney, mas para isso, também preciso ter uma boa base de prompts, o que é um pouco difícil, visto que o MidJourney só está disponível em um servidor do Discord, onde eu teria que ficar capturando as mensagens, mas estou tentando trabalhar nisso.

Conclusões finais:

Eu sinceramente não esperava uma recepção tão grande do projeto, levando em conta que eu criei esse projeto a pouco mais de uma semana só com a intenção inicial de fazer um post aqui detalhando o processo. O modelo do Stable Diffusion atualmente já conta com quase 40 estrelas e mais de 7.000 downloads, além do seu espaço para teste já ter mais de 100 estrelas, estando entre os projetos em alta no Hugging Face, sem eu nem ter quase divulgado o projeto do MagicPrompt. Inclusive, até o site dreamlike.art disse que está usando o meu modelo para tratamento de prompts e me creditou, o que para mim é surreal ver isso acontecendo.

Atualmente, estou tentando trabalhar em um projeto meio ambicioso que chamei de LemonDore, semelhante ao DreamLike e Dall-E, só que com uma pegada de comunidade, semelhante com a do MidJourney, onde os usuários podem criar imagens e cada um ir as compartilhando entre si. E como esse projeto envolve muita coisa, talvez demore um tempo para sair, mas vou tentar anunciar ele aqui assim que lançar.

E é isso, qualquer sugestão ou dúvida, você pode deixar aqui nos comentários em baixo. E muito obrigado por ler até aqui! :)