app.py

import os
import gradio as gr
import torch
import numpy as np
import cv2
import base64
import time
from PIL import Image
from openai import OpenAI
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler
from peft import PeftModel
from transformers import AutoModel

# OpenAI API
api_key = os.getenv('OPENAI_API_KEY')
client = OpenAI(api_key=api_key)

# text to image 設定
access_token = os.getenv('HF_TOKEN')
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe_t2i = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_safetensors=True, token=access_token).to("cuda")
pipe_t2i.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_t2i.scheduler.config)

# 載入 LoRA 權重
lora_path = "./tbh368-sdxl.safetensors"
pipe_t2i.load_lora_weights(lora_path, adapter_name="milton-glaser")
pipe_t2i.set_adapters(["milton-glaser"], adapter_weights=[1.0])

import jieba
import numpy as np
import pandas as pd

valence_scores = []
arousal_scores = []
conversation_times = 0

# 載入 CVAW Corpus 資料
cvaw_data = pd.read_csv('./CVAW_all_SD.csv', delimiter='\t')
cvaw_dict = dict(zip(cvaw_data['Word'], zip(cvaw_data['Valence_Mean'], cvaw_data['Arousal_Mean'])))

def analyze_sentiment_corpus(text):
  words = jieba.cut(text)
  global conversation_times
  global valence_scores
  global arousal_scores

  conversation_times += 1

  for word in words:
    if word in cvaw_dict:
        valence, arousal = cvaw_dict[word]
        valence_scores.append(valence)
        arousal_scores.append(arousal)
        # 只保留最後三句使用者輸入的內容
        valence_scores = valence_scores[-3:]
        arousal_scores = arousal_scores[-3:]

  if conversation_times < 4: # 當對話次數<4的時候，返回10，代表不會進入放鬆模式
    return 10, 10
  else:
    avg_valence = np.mean(valence_scores)
    avg_arousal = np.mean(arousal_scores)
    return avg_valence, avg_arousal

def call_gpt(input_text, history):
  messages = [{"role":"system", "content":"對話請以繁體中文進行：你是一位熟悉現象學的諮商實習生，擅長引導使用者描述他當下的所知覺到的事物。回答問題的時候必須有同理心，請同理使用者說的內容，再繼續回答，且不要超過20個字。"}]

  for h in history:
    messages.append({"role": "user", "content": h[0]})
    messages.append({"role": "assistant", "content": h[1]})

  messages.append({"role": "user", "content": input_text})

  chat_reply = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    temperature=0.8
  )

  return chat_reply.choices[0].message.content

# Testing meditation function
meditation_flag = True # 判斷是否進入放鬆模式

def chat_with_bot(input_text, history):
  global meditation_flag
  response = ""
  med_confirm_layout = False # 是否顯示放鬆選項
  jump2med_btn = True # 是否允許跳轉到放鬆介面

  # 進行情感分析
  valence, arousal = analyze_sentiment_corpus(input_text)

  # 判斷是否建議放鬆練習
  if 3 <= arousal <= 4 and meditation_flag is True:  # 詢問是否進行放鬆練習
    response = "我知道你的處境了\n我有一個建議，我們來進行一個可以讓自己放鬆的呼吸覺察練習好嗎？"
    history.append((input_text, response))
    med_confirm_layout = True
    return history, med_confirm_layout, jump2med_btn
  elif meditation_flag is False:  # 已經放鬆過，不顯示跳轉按鈕
    response = call_gpt(input_text, history)
    history.append((input_text, response))
    jump2med_btn = False
    return history, med_confirm_layout, jump2med_btn
  else: # 繼續對話
    response = call_gpt(input_text, history)
    history.append((input_text, response))
    return history, med_confirm_layout, jump2med_btn

def translate_to_english(text):
  character = "You are a professional text to image prompt generator, please use the following text to generate prompt in English. It's very important to summarize it in 70 tokens."
  messages = [{"role":"system", "content":character},
              {"role": "user", "content": text}]

  chat_reply = client.chat.completions.create(
    messages=messages,
    model="gpt-4o-mini",
  )

  return chat_reply.choices[0].message.content

generated_images = None
last_genimg_times = 0

def generate_images(history):
  global generated_images  # 使用全域變數來儲存圖片
  global last_genimg_times

  if generated_images is not None and last_genimg_times == conversation_times:
    return generated_images  # 如果圖片已生成，直接返回

  user_story = " ".join([h[0] for h in history])
  prompt = translate_to_english(user_story)

  neg_prompt = "realistic, camera, 8k, words, sentence, text, Low quality, error, blurry, deformed, extra, nude, duplicate, ugly"

  seeds = np.random.randint(0, 100000, 4)
  generator = [torch.Generator().manual_seed(int(i)) for i in seeds]
  images = []
  last_genimg_times = conversation_times

  for i in range(4):
    img = pipe_t2i("style of Milton Glaser, mental image, "+prompt,
                negative_prompt=neg_prompt,
                height=720, width=512,
                generator=generator[i],
                num_inference_steps=40,
                guidance_scale=7.5,
                ).images[0]
    images.append(img)

  generated_images = images  # 儲存生成的圖片
  return images

def select_image(choice, img1, img2, img3, img4):
  index = int(choice.split()[-1]) - 1
  images = [img1, img2, img3, img4]
  return images[index]

def chat_about_image(input_text, history, selected_image):
  # 讀取並編碼圖像
  _, buffer = cv2.imencode('.png', selected_image)
  img_str = base64.b64encode(buffer).decode()

  messages = [
      {"role": "system", "content": "對話請以繁體中文進行：你是一位熟悉現象學的諮商實習生，請根據使用者對他所選出的圖像描述進行引導，指出這張圖像與先前對話的關聯，幫助使用者探索他們的分享與該圖像間的連結，並繼續對話"},
      {"role": "user", "content": [
          {"type": "text", "text": f"看到這張圖像，讓我想到 {input_text}"},
          {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_str}"}}
      ]}
  ]

  chat_reply = client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    max_tokens=300
  )

  reply = chat_reply.choices[0].message.content
  history.append((input_text, reply))

  return (
    history,
    history,
    gr.update(visible=True) # chatbot_interface
  )

audio_file = "./meditation_v2.mp3"

# UI handle functions
def handle_chat(input_text, history):
  updated_history, meditation, jump2med_btn = chat_with_bot(input_text, history)
  if meditation:
    return (
        updated_history,
        gr.update(value="", placeholder="現在我們開始進行身體放鬆吧", interactive=False), # msg
        gr.update(visible=False), # submit
        gr.update(visible=False), # jump_to_med
        gr.update(visible=True), # meditation_buttons
        gr.update(visible=False) # clear
      )
  elif jump2med_btn is False:
    return (
        updated_history,
        gr.update(value="", placeholder="今天想跟我分享什麼呢？", interactive=True), # msg
        gr.update(visible=True),  # submit
        gr.update(visible=False), # jump_to_med
        gr.update(visible=False), # meditation_buttons
        gr.update(visible=True)  # clear
      )
  else:
    return (
        updated_history,
        gr.update(value="", placeholder="今天想跟我分享什麼呢？", interactive=True), # msg
        gr.update(visible=True),  # submit
        gr.update(visible=True), # jump_to_med
        gr.update(visible=False), # meditation_buttons
        gr.update(visible=True)  # clear
      )

def start_meditation(history):
  global meditation_flag
  meditation_flag = False
  audio = audio_file
  return (
    audio
  )

def continue_chat():
  return (
    gr.update(value="", placeholder="今天想跟我分享什麼呢？", interactive=True), # msg
    gr.update(visible=True),  # submit
    gr.update(visible=True), # jump_to_med
    gr.update(visible=False), # meditation_buttons
    gr.update(visible=True)  # clear
  )

def return_to_chat():
  return (
    None, # set audio_player to None will stop the audio
    gr.update(visible=True), # main_interface
    gr.update(visible=True), # chatbot_interface
    gr.update(visible=True), # selected_image_interface
    gr.update(visible=False), # audio_interface
    gr.update(interactive=True, placeholder="今天想跟我分享什麼呢？"), # msg
    gr.update(visible=True), # submit
    gr.update(visible=True), # clear
    gr.update(visible=False), # meditation_buttons
    gr.update(visible=True), # gen_other_img
    gr.update(visible=False) # jump_to_med
  )

import asyncio

async def show_loading():
  # 顯示載入訊息
  yield (gr.update(visible=True), # loading_message
        "載入時間約需十秒，建議戴上耳機體驗。\n在接下來的畫面，你會看到播放介面，按下播放鈕後就會開始播放指導語",
        gr.update(visible=False), # main_interface
        gr.update(visible=False)) # audio_interface
  # 等待 10 秒
  await asyncio.sleep(12)
  # 清除訊息
  yield (gr.update(visible=False), # loading_message
        "",
        gr.update(visible=False), # main_interface
        gr.update(visible=True)) # audio_interface

# Testing meditation function
theme = gr.themes.Base(
    primary_hue="amber",
    secondary_hue="sky",
    font=[gr.themes.GoogleFont('Noto Sans TC'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
).set(
    checkbox_background_color_selected_dark='*secondary_400',
    button_border_width='*checkbox_border_width',
    button_shadow='*shadow_drop',
    button_primary_background_fill_hover='*primary_400',
    button_primary_background_fill_hover_dark='*primary_900',
    button_secondary_background_fill='*secondary_100',
    button_secondary_background_fill_dark='*secondary_700',
    button_secondary_background_fill_hover='*secondary_300',
    button_secondary_background_fill_hover_dark='*secondary_900',
    button_secondary_border_color='*secondary_100',
    button_secondary_border_color_dark='*secondary_600',
    button_secondary_text_color='*secondary_700'
)

css = """
.graphic_parent {
    display: flex;
    flex-direction: row;
    flex-wrap: wrap;
    gap: 8px;
    max-height: 100vh;
    max-width: 100vw;
  }
.graphic {
    width: 48%;
  }
"""

with gr.Blocks(theme=theme, css=css) as demo:
  loading_message = gr.Textbox(visible=False, show_label=False)
  with gr.Column(visible=False) as audio_interface:
    audio_player = gr.Audio(label="呼吸覺察練習", show_download_button=False, interactive=False)
    back_to_chat = gr.Button("返回聊天")

  with gr.Row() as main_interface:
    with gr.Column() as chatbot_interface:
      chatbot = gr.Chatbot(label="聊天機器人", bubble_full_width=False, scale=6)
      msg = gr.Textbox(show_label=False, placeholder="今天想要跟我分享什麼呢？", autofocus=True, scale=2)
      with gr.Row():
        submit = gr.Button("送出", variant="primary", scale=2)
        jump_to_med = gr.Button("我想要現在進行呼吸覺察練習", scale=2)
        gen_other_img = gr.Button("結合聯想生成更多圖像", scale=2, visible=False)
        clear = gr.Button("清除對話紀錄", scale=1)
      with gr.Row(visible=False) as meditation_buttons:
        relax_yes = gr.Button("好", variant="primary")
        relax_no = gr.Button("我想再多分享一點")
    with gr.Column(visible=False, elem_classes="graphic_parent") as image_selector_interface:
      image_selector = gr.Radio(choices=["圖像 1", "圖像 2", "圖像 3", "圖像 4"], label="選擇一張圖像")
      image_outputs = [gr.Image(label=f"圖像 {i+1}", interactive=False, elem_classes="graphic") for i in range(4)]
    with gr.Column(visible=False) as selected_image_interface:
      selected_image = gr.Image(interactive=False, label="你選擇的圖像")
      image_chat_input = gr.Textbox(label="這張圖像讓你產生了什麼樣的聯想？")
      image_chat_button = gr.Button("與聊天機器人分享", variant="primary")

  # chatbot events handle
  submit.click(handle_chat, [msg, chatbot], [chatbot, msg, submit, jump_to_med, meditation_buttons, clear])
  msg.submit(handle_chat, [msg, chatbot], [chatbot, msg, submit, jump_to_med, meditation_buttons, clear])
  clear.click(lambda: None, None, chatbot, queue=False)

  # going to meditation events handle
  jump_to_med.click(start_meditation, [chatbot], audio_player)
  relax_yes.click(start_meditation, [chatbot], audio_player)
  relax_no.click(continue_chat, None, [msg, submit, jump_to_med, meditation_buttons, clear])

  # meditation events handle
  jump_to_med.click(show_loading, None, [loading_message, loading_message, main_interface, audio_interface])
  relax_yes.click(show_loading, None, [loading_message, loading_message, main_interface, audio_interface])
  audio_player.play(generate_images, [chatbot], image_outputs)
  audio_player.stop(return_to_chat, None, [audio_player, main_interface, chatbot_interface, selected_image_interface, audio_interface, msg, submit, clear, meditation_buttons, gen_other_img, jump_to_med])
  back_to_chat.click(return_to_chat, None, [audio_player, main_interface, chatbot_interface, selected_image_interface, audio_interface, msg, submit, clear, meditation_buttons, gen_other_img, jump_to_med])

  # images select events handle
  image_selector.change(select_image, [image_selector] + image_outputs, selected_image)
  image_chat_input.submit(chat_about_image, [image_chat_input, chatbot, selected_image], [chatbot, chatbot, chatbot_interface]).then(lambda: None, None, image_chat_input , queue=False)
  image_chat_button.click(chat_about_image, [image_chat_input, chatbot, selected_image], [chatbot, chatbot, chatbot_interface]).then(lambda: None, None, image_chat_input , queue=False)

  # generate other images event handle
  gen_other_img.click(generate_images, [chatbot], image_outputs)

demo.queue(max_threads=10)
demo.launch(show_api=False)