AI enthusiasm #4 - Your stable diffusion chatbot🐠
Tags: tutorial, computervision, python, ai
Published: 2024-04-13 22:09:35

## 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](https://astrabert.github.io/awesome-tiny-sd): make sure to check it out and leave a :star: on GitHub!

First of all, we need to install all the necessary packages:

```bash
python3 -m pip install gradio==4.25.0 diffusers==0.27.2 torch==2.1.2 pydantic==2.6.4 accelerate transformers trl peft
```
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:

```python
from diffusers import DiffusionPipeline
import torch
```

- Define the image-generating pipeline (this will automatically download the stable-diffusion model you specified and all its related components):

```python
pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float32)
```

We chose to use [segmind/small-sd](https://huggingface.co/segmind/small-sd) because it's small and CPU-friendly.

### Block 2: Define chatbot essentials in `app.py`

- Import necessary dependencies:

```python
import gradio as gr
import time
from imgen import *
```

- A simple function to print like and dislikes by the users:

```python
def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)
```

- The function the appends new messages and/or uploaded files to the chatbot history:

```python
def add_message(history, message):
    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:

```python
def bot(history):
    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:

```python
with gr.Blocks() as demo:
    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:

```python
demo.queue()
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", share=False)
```

- Run the script:

```bash
python3 app.py
```

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:

```bash
docker pull ghcr.io/astrabert/awesome-tiny-sd:latest
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](https://github.com/sponsors/AstraBert): 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!:heart:

_Cover image by Google DeepMind_