【LLM】构建一个可以图像生成的聊天机器人

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

./
|__ app.py
|__ imgen.py

  from diffusers import DiffusionPipeline
  import torch

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

  import gradio as gr
  import time
  from imgen import *

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

  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)

  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

  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

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

  python3 app.py

docker pull ghcr.io/astrabert/awesome-tiny-sd:latest
docker run -p 7860:7860 ghcr.io/astrabert/awesome-tiny-sd:latest