Build a simple image-generating chatbot
Today I’m guiding you through the creation, from scratch, of an image generating chatbot.
We’ll doing it following the script I used to build my already existent chatbot, awesome-tiny-sd: make sure to check it out and leave a ⭐ on GitHub!
First of all, we need to install all the necessary packages:
Once you did that, make sure to set up your folder like this:
|__ app.py
|__ imgen.py
And let’s begin coding!😎
Block 1: import your favorite stable-diffusion model in imgen.py
Import necessary dependencies:
import torch
Define the image-generating pipeline (this will automatically download the stable-diffusion model you specified and all its related components):
We chose to use segmind/small-sd because it’s small and CPU-friendly.
Block 2: Define chatbot essentials in app.py
Import necessary dependencies:
import time
from imgen import *
A simple function to print like and dislikes by the users:
print(x.index, x.value, x.liked)
The function the appends new messages and/or uploaded files to the chatbot history:
if len(message[“files“]) > 0:
history.append((message[“files“], None))
if message[“text“] is not None and message[“text“] != “”:
history.append((message[“text“], None))
return history, gr.MultimodalTextbox(value=None, interactive=False)
The function that, starting from the text-prompt, generates an image:
if type(history[–1][0]) != tuple: ## text prompt
try:
prompt = history[–1][0]
image = pipeline(prompt).images[0] ## call the model
image.save(“generated_image.png“)
response = (“generated_image.png“,)
history[–1][1] = response
yield history ## return the image
except Exception as e:
response = f“Sorry, the error ‘{e}‘ occured while generating the response; check [troubleshooting documentation](https://astrabert.github.io/awesome-tiny-sd/#troubleshooting) for more“
history[–1][1] = “”
for character in response:
history[–1][1] += character
time.sleep(0.05)
yield history
if type(history[–1][0]) == tuple: ## input are files
response = f“Sorry, this version still does not support uploaded files 🙁“ ## We will see how to add this functionality in the future
history[–1][1] = “”
for character in response:
history[–1][1] += character
time.sleep(0.05)
yield history
Block 3: build the actual chatbot
Define the chatbot blocks with Gradio:
chatbot = gr.Chatbot(
[[None, (“Hi, I am awesome-tiny-sd, a little stable diffusion model that lets you generate images:blush:nJust write me a prompt, I‘ll generate what you ask for:heart:“,)]], ## the first argument is the chat history
label=“awesome-tiny-sd“,
elem_id=“chatbot“,
bubble_full_width=False,
) ## this is the base chatbot architecture
chat_input = gr.MultimodalTextbox(interactive=True, file_types=[“png“,“jpg“,“jpeg“], placeholder=“Enter your image-generating prompt…“, show_label=False) ## types of supported input
chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input]) ## receive a message
bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name=“bot_response“) ## send a message
bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
chatbot.like(print_like_dislike, None, None)
clear = gr.ClearButton(chatbot) ## show clear button
Launch the chatbot:
if __name__ == “__main__“:
demo.launch(server_name=“0.0.0.0“, share=False)
Run the script:
Now the chatbot, once the stable diffusion pipeline is loaded, should be running on localhost:7860 (or 0.0.0.0:7860 for Linux-like OS).
You can give a try on this Hugging Face space: https://huggingface.co/spaces/as-cle-bert/awesome-tiny-sd
Otherwise, you can download awesome-tiny-sd Docker image and run it through container:
docker run -p 7860:7860 ghcr.io/astrabert/awesome-tiny-sd:latest
Give it a try, you won’t be disappointed!!!
Do not forget to sponsor the project on GitHub: if we get far enough with sponsoring, we will upgrade the HF space to a GPU-powered one in order to make image generation faster.
What will be the first image you are going to generate with awesome-tiny-sd? Let me know in the comments below!❤️
Cover image by Google DeepMind
