DEV Community

Cover image for AI enthusiasm #4 - Your stable diffusion chatbot🐠
Clelia (Astra) Bertelli
Clelia (Astra) Bertelli

Posted on

2

AI enthusiasm #4 - Your stable diffusion chatbot🐠

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:

python3 -m pip install gradio==4.25.0 diffusers==0.27.2 torch==2.1.2 pydantic==2.6.4 accelerate transformers trl peft
Enter fullscreen mode Exit fullscreen mode

Once you did that, make sure to set up your folder like this:

./
|__ app.py
|__ imgen.py
Enter fullscreen mode Exit fullscreen mode

And let's begin coding!😎

Block 1: import your favorite stable-diffusion model in imgen.py

  • Import necessary dependencies:
from diffusers import DiffusionPipeline
import torch
Enter fullscreen mode Exit fullscreen mode
  • Define the image-generating pipeline (this will automatically download the stable-diffusion model you specified and all its related components):
pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float32)
Enter fullscreen mode Exit fullscreen mode

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 gradio as gr
import time
from imgen import *
Enter fullscreen mode Exit fullscreen mode
  • A simple function to print like and dislikes by the users:
def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)
Enter fullscreen mode Exit fullscreen mode
  • The function the appends new messages and/or uploaded files to the chatbot history:
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)
Enter fullscreen mode Exit fullscreen mode
  • The function that, starting from the text-prompt, generates an image:
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
Enter fullscreen mode Exit fullscreen mode

Block 3: build the actual chatbot

  • Define the chatbot blocks with Gradio:
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
Enter fullscreen mode Exit fullscreen mode
  • Launch the chatbot:
demo.queue()
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", share=False)
Enter fullscreen mode Exit fullscreen mode
  • Run the script:
python3 app.py
Enter fullscreen mode Exit fullscreen mode

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 pull ghcr.io/astrabert/awesome-tiny-sd:latest
docker run -p 7860:7860 ghcr.io/astrabert/awesome-tiny-sd:latest
Enter fullscreen mode Exit fullscreen mode

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

Image of Datadog

The Future of AI, LLMs, and Observability on Google Cloud

Datadog sat down with Google’s Director of AI to discuss the current and future states of AI, ML, and LLMs on Google Cloud. Discover 7 key insights for technical leaders, covering everything from upskilling teams to observability best practices

Learn More

Top comments (0)

Billboard image

The Next Generation Developer Platform

Coherence is the first Platform-as-a-Service you can control. Unlike "black-box" platforms that are opinionated about the infra you can deploy, Coherence is powered by CNC, the open-source IaC framework, which offers limitless customization.

Learn more