why10w/kg.py

# encoding: utf-8
import erniebot
import json
from answerQuestions import generateStory

erniebot.api_type = "aistudio"
erniebot.access_token = "ff1531c8c0f429f92adbc2eaed2e23bfb5349e0f"


def res2json(res: str):
    res = res[8:-4]
    res_json = json.loads(res)
    return res_json


def extract_KG(description: str):
    res = ""
    prompt = """你是专门进行实体抽取的专家。请从input中抽取出实体和实体间的关系，最多给出2对最主要的实体关系。实体只能是具体事物：如动物、植物、人、物品等。
请按照JSON字符串的格式回答: [{"entity": "", "relation": "", "object": ""}]"""
    response = erniebot.ChatCompletion.create(
        model="ernie-4.0",
        messages=[{"role": "user", "content": prompt + "\ninput:\n" + description}],
    )
    result = response["result"]
    res = res2json(result)
    print(f"知识图谱抽取结果：{res}")
    return res


def split_img_describe(story: str):
    print("\n开始对故事进行分段")
    prompt = """你是专门进行绘本描绘的专家。请从input中抽取成一到二张图片的描述，每张图片只包含一个场景和一个动作，主体越少越好。请按照JSON字符串的格式回答: [{"caption":"","description":""}]```"""
    response = erniebot.ChatCompletion.create(
        model="ernie-4.0",
        messages=[{"role": "user", "content": prompt + "\ninput:\n" + story}],
    )
    result = response["result"]

    print(f"故事片段切分结果：{result}")
    return res2json(result)


text2img_prompt = """
文生图prompt要求:

```md
每张图片prompt描述是英文

以“反推法”，头脑中先生成画面，再落地成prompt，最后输入大模型，验证模型的生成效果。
首先我们要在大脑中生成一个画面，每个画面的构成有三到要素：人、景、物，每个要素都可以单独生成画面，也可以随意的排列组合形成画面。

1.人·动作定格
    首先我们以人物的动作定格，这里面的人物包括拟人化的动物。
    从形象的角度出发，人物照片又分为全身照、半身照(不再示例)和细节照。
* 以全身照为例，生成的画面着重在人物整体，轻表情细节。人物整体描述(prompt)先下三大定义：数量+性别+年龄，再作化整为零描述着装（外形）的形状和颜色：头+上半身+下半身+脚，最后进行大动作描述：跑步、跳舞、站立.....
    示例：一个中年男士，黑色短发，面容冷峻，黑色西装，白色衬衫，深褐色领带，黑色皮带和西装裤，深褐色袜子，黑色皮鞋，向我跑过来。
    翻译：A middle-aged man with short black hair and a cold face ran towards me in a black suit, white shirt, dark brown tie, black belt and suit pants, dark brown socks, and black leather shoes.
* 以细节照为例，生成的画面着重在动态定格，轻相对静态。面部细节描述(prompt)先下两大定义：性别+年龄+外形，再做化整为零描述动作（五官）的动作和形态：眉毛+眼睛+鼻子+嘴巴+耳朵，最后是整个表情的定义：开心、失落、愤怒.......
    示例：一个10岁左右的小女孩，眉毛皱在一起，眼睛瞪的圆圆的看着我，鼻孔撑大，下嘴唇轻咬上嘴唇，看起来很伤心的样子。
    翻译：A little girl around 10 years old, her eyebrows furrowed together, her eyes wide and round looking at me, her nostrils widened, her lower lip gently biting her upper lip, looking very sad.


2.景·以静画动
    再次我们以风景的元素作拆解，这里面的风景可以是狭义的景色，也可以是广义的背景。
风景元素有很多，包括青山、绿水、芦苇、水鸟、汽车、马路等等。
    我们首先进行脑补构图，脑补先确认是否有明确的分界线，例如
* 上下的分界：地平线将画面分为天空和大地、海平面将画面分为天空和大海，诸如此类的海岸线、桌线等等不作赘述。
* 左右的分界：高速公路将地面分为左右两块，装订线将书本分为左右两个页面，诸如此类的顶梁柱、人中线等等同不作赘述。
* 其他分界线如曲线、圆形等等也可以套用此分类。
    然后按照分界线的顺序进行分解描述，如上线分界线，先介绍上面或下面再描述另一区域。每个区域逐个进行元素描述（prompt）。
示例：蔚蓝的大海上，漂泊着一只小白船，天上飞着很多海鸥。
（本示例为生成背景，prompt均轻量化描述）
翻译：On the vast blue sea, a small white boat drifts, and many seagulls fly in the sky。

3.物·光影留帧
    最后我们以物体的形象作描述，这里面的物体可以是画面中重点突出的对象，也可以是混入整体背景的参照物。
    物体的分类有很多，电脑、插座、床、快递箱等等。
    对于参照物的描述（prompt），可以一笔带过；对于重点描述(突出C位)的物体，要有直观感受的细节化描述（prompt）：形状、大小(通过参照物)、颜色等三方面内容，进而再进行细节延伸的描述。
示例：桌子上有一台笔记本电脑，电脑旁边放着一个透明玻璃杯，这是一个圆柱形的杯子，半杯水还冒着热气，杯身折射了些许阳光。
翻译：There is a laptop on the table, and next to it is a transparent glass. This is a cylindrical cup, half of which is still steaming hot, and the cup body reflects some sunlight。"""


def generate_text2img_prompt(input: str):
    print("\n正在生成 文生图prompt...")
    t_prompt = (
        """你是专门使用文生图的专家。请将input转换为一个`文生图的prompt`，所有词语使用英文. 请直接按照字符串的格式给出prompt"""
        + text2img_prompt
        + "\n切记，只生成一条prompt，并且直接以字符串的格式给出"
    )

    response = erniebot.ChatCompletion.create(
        model="ernie-4.0",
        messages=[{"role": "user", "content": t_prompt + "\ninput:\n" + str(input)}],
    )
    result = response["result"]
    print("\n生成的图片prompt：" + result)
    return result


# story = """
# 小朋友，你听说过“克隆”这个词吗？让我来给你讲一个有趣的故事吧！
#
# 有一天，科学家们想出了一个非常神奇的方法，可以让一个生命体“复制”出另一个一模一样的生命体，这个方法就叫做“克隆”。你可以把“克隆”想象成是一台神奇的复印机，但是这台复印机不是复印纸张，而是复印生命哦！
#
# 1997年的时候，科学家们用这种方法成功地“复印”出了一只名叫多利的绵羊。你知道吗，多利羊并没有爸爸，它是由一只母羊的细胞复制而来的。这就像是我们从电脑上复制一个文件一样，只不过这次我们复制的是一个活生生的羊！
#
# 克隆羊多利的诞生，就像是打开了一扇新的大门，让人们看到了生命的奇妙和科学的无限可能。从此以后，人们开始更加关注和研究克隆技术，希望用它来帮助我们解决更多的问题。
#
# 所以，克隆羊就是通过克隆技术创造出来的一只羊，它没有父亲，只有母亲，是科学家们用神奇的“生命复印机”复制出来的哦！
# """


def judge_img(url: str, kg):
    """
    识图。判断图片和知识图谱是否一致。
    """
    prompt = """你是一个识图专家.
首先，你需要去识别图片中有什么？重点是有哪些事物。
最后，请判断图片和实体描述是否一致。

返回格式要求：
- result: 图片中包含描述的内容，返回true。图片与描述不一致，返回false。
- message：提示信息，只有当result为false时才需要填写提示。

请按照JSON字符串的格式回答: {"result": true/false,"message:""}.

图谱描述: """ + str(
        kg
    )
    print("\n正在校验图片...")
    # prompt = "图片中有什么?"
    from img import judge_image

    res = judge_image(url, prompt)

    res_text = res["output"]["choices"][0]["message"]["content"][0]["text"]
    print("\n校验结果：")
    print(res_text)
    res_json = json.loads(res_text)
    return res_json


def generate_caption_image_iterate(description):
    """

    输入一个图片描述 description

    通过以下步骤迭代符合描述的图片：
    1. 通过大模型抽取知识图谱实体及关系
    2. 根据图谱信息生成文生图prompt
    3. 生成图片
    4. 对图片和图谱信息进行校验，如果失败，则将错误信息输入到第一步，重新执行1-4步骤。
    5. 如果成功，则返回图片URL。
    """

    img_url = ""
    error_message = ""
    try_num = 0
    max_try_num = 3
    while try_num < max_try_num:
        # 1. 通过大模型抽取知识图谱实体及关系
        kg_input = description
        if error_message != "":
            kg_input += "\n\n提示：请注意避免如下的问题: " + error_message
        kg = extract_KG(kg_input)

        # 2. 根据图谱信息生成文生图prompt
        prompt = generate_text2img_prompt(str(kg))

        # 3. 生成图片
        from img import text_to_image

        img_url = text_to_image(prompt)
        print(img_url)

        # 4. 对图片和图谱信息进行校验，如果失败，则将错误信息输入到第一步，重新执行1-4步骤。
        judge_result = judge_img(img_url, kg)
        if judge_result["result"] in [True, "true"]:
            # 通过
            break
        else:
            error_message = judge_result["message"]

        try_num += 1
    return img_url


def story_to_imgs(story: str):
    from pathlib import PurePosixPath
    import requests

    from urllib.parse import urlparse, unquote

    """
    由故事生成绘本图片集

    1. 将故事分段
    2. 对每个故事片段迭代生成图片

    返回:
    [
        {
            "caption":"多利和它的“母亲”在一起",
            "url":""
        }
    ]
    """

    # 1. 将故事分段
    caption_list = split_img_describe(story)

    # 2. 对每个故事片段迭代生成图片
    res = []
    for caption in caption_list:
        print(f"对caption: {caption} 进行图片生成...")
        print("\n ------ ")
        url = generate_caption_image_iterate(caption["description"])
        res.append({"caption": caption["caption"], "url": url})

    print("\n绘本生成结果：")
    for item in res:
        print(item["caption"], item["url"])
        url = item["url"]
        # 保存到本地
        file_name = item["caption"] + ".png"
        with open("images/%s" % file_name, "wb+") as f:
            f.write(requests.get(url).content)
    return res


if __name__ == "__main__":
    question = "什么是克隆羊"
    story = generateStory(question)

    res = story_to_imgs(story)
    print(res)
    # url = f"https://dashscope-result-hz.oss-cn-hangzhou.aliyuncs.com/1d/aa/20240707/522176a8/257fdf50-f48e-47c6-8e8d-26b847e36bf3-1.png?Expires=1720443646&OSSAccessKeyId=LTAI5tQZd8AEcZX6KZV4G8qL&Signature=7kGPYX7l9QruXzIfmOk4%2FuPNXr0%3D"

    # kg = [
    #     {"entity": "克隆羊多利", "relation": "依偎", "object": "母羊"},
    #     {"entity": "草地", "relation": "环绕", "object": "绿树和鲜花"},
    # ]
    # judge_img(url, kg)

    # description = "在户外的草地上，克隆羊多利和一只看起来一模一样的母羊正亲昵地依偎在一起。它们的眼神温柔而充满爱意，仿佛在告诉对方彼此是多么重要。周围是绿树和鲜花，构成了一幅和谐而美好的画面。"
    # generate_caption_image_iterate(description)
    # caption_list = split_img_describe(story)
    # res = extract_KG(caption_list[-1]["description"])
    # print(res)

    #     # [{'entity': '克隆羊多利', 'relation': '提供细胞', 'object': '母羊'}, {'entity': '科学家们', 'relation': '观察', 'object': '克隆羊多利和母羊'}]
    #     # 合并成图片的prompt
    #     prompt = ""
    #     for item in res:
    #         prompt += item["entity"] + " " + item["relation"] + " " + item["object"] + "\n"


#     from img import simple_call, draw

#     simple_call(result)