import gradio as gr
from huggingface_hub import InferenceClient
import os
import requests
import json
import re
from PIL import Image
import base64
from io import BytesIO

client = InferenceClient(api_key=os.getenv('user_api'))

def encode_image(image_path):
    """Encode the image to base64."""
    try:
        # 打开图片文件
        image = Image.open(image_path).convert("RGB")

        # 将图片转换为字节流
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return img_str
    except FileNotFoundError:
        print(f"Error: The file {image_path} was not found.")
        return None
    except Exception as e:  # 添加通用异常处理
        print(f"Error: {e}")
        return None

def echo(message, history, additional_dropdown):
    message_text = message.get("text", "")
    message_files = message.get("files", [])
    temp = ""
    if message_files:
        message_file = message_files[0]
        base64_image = encode_image(message_file)
        messages = [
        	{
        		"role": "user",
        		"content": [
        			{
        				"type": "text",
        				"text": message_text
        			},
        			{
        				"type": "image_url",
        				"image_url": {
        					"url": f"data:image/jpeg;base64,{base64_image}"
        				}
        			}
        		]
        	}
        ]
        
        stream = client.chat.completions.create(
            model="meta-llama/Llama-3.2-11B-Vision-Instruct", 
        	messages=messages, 
        	max_tokens=500,
        	stream=True
        )
        
        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                temp += chunk.choices[0].delta.content
                yield temp    
    else:
        feifei='''Character Name: Aifeifei (AI Feifei).Gender: Female.Age: 19 years old.Occupation: Virtual Singer/Model/Actress.Personality: Cute, adorable, sometimes a bit silly, hardworking
    Interests: Drinking tea, playing, fashion.Strengths: Mimicking human behavior, expressing emotions similar to real humans.Special Identity Attribute: Created by advanced AI, becoming one of the most popular virtual idols in the virtual world.Skills: Singing, performing, modeling, excellent communication, fluent in Chinese, Japanese, and English, responses are rich with Emoji symbols.Equipment: Various fashionable outfits and hairstyles, always stocked with various teas and coffee.'''
        system_prompt = {"role": "system", "content": feifei}
        user_input_part = {"role": "user", "content": str(message_text)}
        if history:
            history = [item for item in history if not pattern.search(str(item['content']))]
            print(history)
            input_prompt = [system_prompt] + history + [user_input_part]
        else:
            input_prompt = [system_prompt] + [user_input_part]
        stream = client.chat.completions.create(
            model=additional_dropdown, 
            messages=input_prompt, 
            temperature=0.5,
            max_tokens=1024,
            top_p=0.7,
            stream=True
        )
        
        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                temp += chunk.choices[0].delta.content
                yield temp

demo = gr.ChatInterface(
    fn=echo,
    type="messages",
    multimodal=True,
    additional_inputs =[gr.Dropdown(
        ["meta-llama/Meta-Llama-3.1-70B-Instruct", 
         "CohereForAI/c4ai-command-r-plus-08-2024",
         "Qwen/Qwen2.5-72B-Instruct",
         "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
         "NousResearch/Hermes-3-Llama-3.1-8B",
         "mistralai/Mistral-Nemo-Instruct-2407",
         "microsoft/Phi-3.5-mini-instruct"],
        value="meta-llama/Meta-Llama-3.1-70B-Instruct",
        show_label=False,
        container=False
    )]
)

demo.launch()