app_ark_image.py

import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import urllib.request
import gradio as gr
from pipeline.pipeline_chat import chat_story, create_hint
from pipeline.pipeline_gen_image_video import gen_image_old
import json
import urllib
from datetime import datetime

with open("/home/aiteam1/Videos/Phuong/Do/ai_virtual_gf/jsons/new_characters_story_talk_style2.json", 'r', encoding='utf-8') as file:
    characters = json.load(file)

with open('/home/aiteam1/Videos/Phuong/Do/ai_virtual_gf/jsons/bot_levels.json', 'r') as f:
    level_behaviors = json.load(f)

category_list = [cate["cate_name"] for cate in characters]

with open("/home/aiteam1/Videos/Phuong/Do/ai_virtual_gf/jsons/gifts.json", 'r', encoding='utf-8') as f:
    gifts = json.load(f)
    
gift_list = [gifts[f"{i}"] for i in range(1, 9)]

def get_list_character(category, talk_style):
    for cate in characters:
        if cate['cate_name'] == category:
            character_list = cate['characters']
            if talk_style in ['shy', 'mixed', 'comfortable']:
                character_name_list = []
                for character in character_list:
                    if character['talk_style'] == talk_style:
                        character_name_list.append(character['name'])
            else:
                character_name_list = [character['name'] for character in character_list]
    return gr.update(choices=character_name_list, value=character_name_list[0])

def get_current_bot_level(type_level, talk_style):
    return gr.update(choices=level_behaviors[type_level][talk_style].keys(), value=level_behaviors[type_level][talk_style].keys()[0])

def get_preview(cate_name, character_name):
    for cate in characters:
        if cate['cate_name'] == cate_name:
            for character in cate['characters']:
                if character_name == character['name']:
                    preview = character['preview']
                    return preview
    return None

def hint(chat_history):
    hints = create_hint(chat_history)
    parsed_hints = [h.split(":", 1)[1].strip() if ":" in h else h for h in hints]
    return gr.update(choices=parsed_hints, value=parsed_hints[0] if parsed_hints else None, interactive=True)

def get_parameter(cate_name, character_name):
    for cate in characters:
        if cate['cate_name'] == cate_name:
            for character in cate['characters']:
                if character_name == character['name']:
                    character_id = character['id']
                    introduction = character['introduction']
                    persona_tag = character['persona_tag']
                    personality = character['personality']
                    language_characteristics = character['language_characteristics']
                    character_relationship = character['character_relationship']
                    gender = character['gender']
                    talk_style = character['talk_style']
                    reference_image = character['image_link']
                    return character_id, introduction, persona_tag, personality, language_characteristics, character_relationship, gender, talk_style, reference_image
    return (None,) * 9

# lưu lịch sử chat cho từng nhân vật
chat_histories = {}

def get_images_path(cate, name): 
    now = datetime.now() 
    year, month, day = str(now.year), str(now.month), str(now.day) 
    # id = uuid.uuid4() 
    images_folder = "images_gradio" 
    if not os.path.exists(images_folder): 
        os.makedirs(images_folder)
    cate_folder = os.path.join(images_folder, cate)
    if not os.path.exists(cate_folder): 
        os.makedirs(cate_folder) 
    character_folder = os.path.join(cate_folder, name)
    if not os.path.exists(character_folder): 
        os.makedirs(character_folder) 
    year_folder = os.path.join(character_folder, year) 
    if not os.path.exists(year_folder): 
        os.makedirs(year_folder) 
    month_folder = os.path.join(year_folder, month) 
    if not os.path.exists(month_folder): 
        os.makedirs(month_folder) 
    day_folder = os.path.join(month_folder, day) 
    if not os.path.exists(day_folder): 
        os.makedirs(day_folder)
    i = len(os.listdir(day_folder))
    print(f"the number of {name}'images: {i}")
    images_path = os.path.join(day_folder, f"{i+1}.png") 
    return images_path
    

def chat_inter(message, chat_history, cate, character, type_level, current_bot_level, user_name=None):
    character_id, introduction, persona_tag, personality, language_characteristics, character_relationship, age, gender, talk_style, reference_image = get_parameter(cate, character)

    # Nếu user yêu cầu gen ảnh
    if message == 'Send me your photo':
        url_generated_image, log_infor = gen_image_old(cate, character, age)
        character = character.replace(" ", "_")
        image_path = get_images_path(cate, character)
        urllib.request.urlretrieve(url_generated_image, image_path)

        chat_history.append({"role": "user", "content": "Send me your photo"})
        chat_history.append({"role": "assistant", "content": url_generated_image})
        chat_histories[character] = chat_history
        return "", chat_history

    # Trường hợp bình thường (user chat)
    character_id, replies, messages, log_infor = chat_story(
        chat_history, user_name, character_id, character, introduction, 
        persona_tag, personality, language_characteristics, character_relationship, 
        gender, talk_style, reference_image, level_behaviors, type_level, current_bot_level, 
        gifts, "", ""
    )
    print(f'replies: {replies}')
    chat_history.append({"role": "user", "content": message})
    for reply in replies:
        if reply:
            chat_history.append({"role": "assistant", "content": reply})
    chat_histories[character] = chat_history
    
    print(f"reply: {replies}")
    for reply in replies:
        if reply and 'http' in reply:
            chat_history.append({"role": "assistant", "content": reply})
            chat_histories[character] = chat_history
    return "", chat_history

# Load hoặc khởi tạo hội thoại mới
def init_on_load(cate_name, character, user_name, type_level, current_bot_level):
    if character in chat_histories:
        return chat_histories[character]

    character_id, introduction, persona_tag, personality, language_characteristics, character_relationship, gender, talk_style, reference_image= get_parameter(cate_name, character)
    character_id, replies, messages, hints, score, log_infor = chat_story(
        [], user_name, character_id, character, introduction, 
        persona_tag, personality, language_characteristics, character_relationship, 
        gender, talk_style, reference_image, level_behaviors, type_level, current_bot_level, 
        gifts, "", ""
    )
    for reply in replies:
        opening_message = [{"role": "assistant", "content": reply}]
        # chat_histories[character].append(opening_message)
    # print(chat_histories[character])
    return opening_message

# Reset hội thoại cho nhân vật
def reset_chat(character, user_name, type_level, current_bot_level):
    for c in characters:
        if c['name'] == character:
            reply = chat_story(
                [], user_name, c['id'], c['name'], c['introduction'], 
                c['persona_tag'], c['personality'], c['language_characteristics'], c['character_relationship'],
                c['gender'], c['talk_style'], c['reference_image'], level_behaviors, type_level, current_bot_level,
                gifts, '', ''
            )
            opening_message = [{"role": "assistant", "content": reply}]
            chat_histories[character] = opening_message
            return opening_message
    return []

def send_gift(gift):
    pass

with gr.Blocks() as demo:
    with gr.Column():
        user_name = gr.Text(label="Your name", value="User")
        cate_name = gr.Dropdown(label='Name of category', choices=category_list)
        talk_style = gr.Dropdown(label='Talk Style', choices=["shy", "mixed", "comfortable"])
        bot_character = gr.Dropdown(choices=[], label="Bot's character")
        type_level = gr.Dropdown(choices=['low', 'medium', 'high'], label='Type level')
        current_bot_level = gr.Dropdown(label="Current bot level", choices=[])
        gift = gr.Dropdown(choices=gift_list, label='gift')
        show_hint = gr.Dropdown(label='hint', choices=[])  
        
        cate_name.change(
            fn=get_list_character,
            inputs=[cate_name, talk_style],
            outputs=bot_character
        )
        
        talk_style.change(
            fn=get_list_character,
            inputs=[cate_name, talk_style],
            outputs=bot_character
        )
        
        talk_style.change(
            fn=get_current_bot_level,
            inputs=[type_level, talk_style],
            outputs=[current_bot_level]
        )
        
        with gr.Row():
            ask_btn = gr.Button("Ask Image")
            reset_btn = gr.Button("Reset Chat")
            hint_btn = gr.Button("Hints")
            
        preview = gr.Text(label="Preview")
        bot_character.change(
            fn=get_preview,
            inputs=[cate_name, bot_character],
            outputs=[preview]
        )
        
    with gr.Column():
        chatbot = gr.Chatbot(label="Chat", type="messages", height=400)
        msg = gr.Textbox(label="Message", placeholder="Type here...")

        msg.submit(
            chat_inter,
            inputs=[msg, chatbot, bot_character, type_level, current_bot_level, user_name],
            outputs=[msg, chatbot],
        )

        show_hint.select(
            chat_inter,
            inputs=[show_hint, chatbot, bot_character, type_level, current_bot_level, user_name],
            outputs=[msg, chatbot]
        )
        
        bot_character.change(
            fn=init_on_load,
            inputs=[cate_name, bot_character, user_name, type_level, current_bot_level],
            outputs=chatbot
        )
        
        def ask_photo(chat_history, cate_name, character, type_level, current_bot_level, user_name):
            return chat_inter("Send me your photo", chat_history, cate_name, character, type_level, current_bot_level, user_name)
        
        ask_btn.click(
            ask_photo,
            inputs=[chatbot, cate_name, bot_character, type_level, current_bot_level, user_name],
            outputs=[msg, chatbot]
        )

        # Nút Reset Chat
        reset_btn.click(
            reset_chat,
            inputs=[bot_character, user_name, type_level, current_bot_level],
            outputs=chatbot
        )
        
        hint_btn.click(
            hint,
            inputs=[chatbot],
            outputs=show_hint
        )

    # Bot nói ngay khi mở web
    def init_all(cate_name, talk_style, type_level, current_bot_level, user_name):
        char_update = get_list_character(cate_name, talk_style)
        default_character = char_update['value']
        chat_init = init_on_load(cate_name, default_character, type_level, current_bot_level, user_name)
        hints = create_hint(chat_init)
        parsed_hints = [h.split(":", 1)[1].strip() if ":" in h else h for h in hints]
    
        return char_update, chat_init, gr.update(choices=parsed_hints, value=parsed_hints[0] if parsed_hints else None, interactive=True)

    demo.load(
        fn=init_all,
        inputs=[cate_name, talk_style, type_level, current_bot_level, user_name],
        outputs=[bot_character, chatbot, show_hint]
    )

demo.launch(server_port=6666, server_name="192.168.5.155", share=True)