from datasets import load_dataset
import gradio as gr
from gradio_client import Client
import json
import torch
from diffusers import FluxPipeline, AutoencoderKL
from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
import spaces

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

llm_client = Client("Qwen/Qwen2.5-72B-Instruct")

ds = load_dataset("MohamedRashad/FinePersonas-Lite", split="train")

prompt_template = """다음 페르소나 설명을 가진 캐릭터를 생성하세요: {persona_description}
다음 설명의 세계에서: {world_description}
다음 필드를 포함하는 JSON 형식으로 캐릭터를 작성하세요:
- name: 캐릭터의 이름
- background: 캐릭터의 배경
- appearance: 캐릭터의 외모
- personality: 캐릭터의 성격
- skills_and_abilities: 캐릭터의 기술과 능력
- goals: 캐릭터의 목표
- conflicts: 캐릭터의 갈등
- backstory: 캐릭터의 과거 이야기
- current_situation: 캐릭터의 현재 상황
- spoken_lines: 캐릭터의 대사 (문자열 리스트)
JSON 형식의 캐릭터 설명만 작성하고 다른 내용은 포함하지 마세요. '```'도 포함하지 마세요.
"""

world_description_prompt = "독특하고 무작위한 세계 설명을 생성하세요 (세계 설명만 작성하고 다른 내용은 포함하지 마세요)."

def get_random_world_description():
    result = llm_client.predict(
        query=world_description_prompt,
        history=[],
        system="반드시 한글로 출력하라. 당신은 도움이 되는 어시스턴트입니다.",
        api_name="/model_chat",
    )
    return result[1][0][-1]

def get_random_persona_description():
    return ds.shuffle().select([100])[0]["persona"]

@spaces.GPU(duration=75)
def infer_flux(character_json):
    for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
        prompt=character_json["appearance"],
        guidance_scale=3.5,
        num_inference_steps=28,
        width=1024,
        height=1024,
        generator=torch.Generator("cpu").manual_seed(0),
        output_type="pil",
        good_vae=good_vae,
    ):
        yield image

def generate_character(world_description, persona_description, progress=gr.Progress(track_tqdm=True)):
    result = llm_client.predict(
        query=prompt_template.format(
            persona_description=persona_description, world_description=world_description
        ),
        history=[],
        system="반드시 한글로 출력하라. 당신은 도움이 되는 어시스턴트입니다.",
        api_name="/model_chat",
    )
    output = json.loads(result[1][0][-1])
    return output

with gr.Blocks(title="캐릭터 자동 생성", theme="Nymbo/Nymbo_Theme") as app:
    with gr.Column():
        gr.HTML("<center><h1>캐릭터 생성기</h1></center>")
        with gr.Column():
            with gr.Row():
                world_description = gr.Textbox(lines=10, label="세계 설명", scale=4)
                persona_description = gr.Textbox(lines=10, label="페르소나 설명", value=get_random_persona_description(), scale=1)
            with gr.Row():
                random_world_button = gr.Button(value="무작위 세계 설명 가져오기", variant="secondary", scale=1)
                submit_button = gr.Button(value="흥미로운 캐릭터 생성하기!", variant="primary", scale=5)
                random_persona_button = gr.Button(value="무작위 페르소나 설명 가져오기", variant="secondary", scale=1)
        with gr.Row():
            character_image = gr.Image(label="캐릭터 이미지")
            character_json = gr.JSON(label="캐릭터 설명")

        examples = gr.Examples(
            [
                "마법이 실존하고 용들이 하늘을 날아다니는 세계에서, 모험가들의 일행이 전설적인 용왕의 검을 찾아 나섭니다.",
                "에토리아에 오신 것을 환영합니다. 이곳은 물리 법칙이 고대 마법의 의지에 따라 휘어지는 광활하고 신비로운 영역입니다. 이 세계는 끝없는 하늘에 떠 있는 수많은 부유 섬으로 구성되어 있으며, 각 섬은 생명과 비밀로 가득 찬 독특한 생태계입니다. 에토리아의 섬들은 울창한 정글부터 황량한 수정 사막까지 다양합니다. 어떤 섬은 도시 블록만큼 작고, 어떤 섬은 수백 마일에 걸쳐 있습니다. 이 분리된 땅덩어리들을 연결하는 것은 순수한 에너지로 만들어진 반짝이는 다리이며, beaten path를 벗어날 만큼 용감한 사람들은 먼 거리를 순간적으로 이동할 수 있는 숨겨진 포털을 찾을 수 있습니다. 에토리아의 주민들은 그 풍경만큼이나 다양합니다. 인간들은 빛의 에테르 존재, 바위 피부를 가진 거인, 분류하기 어려운 변형 생물들과 공존합니다. 고대 유적들이 섬들을 점점이 장식하고 있어, 과학과 마법의 경계를 흐리는 잊혀진 문명과 기술을 암시합니다. 세계는 에테르라는 신비한 물질로 움직이며, 이것은 모든 것을 통해 흐릅니다. 에테르의 힘을 다룰 수 있는 사람들은 현실 자체를 조작할 수 있는 강력한 마법사가 됩니다. 그러나 에테르는 한정된 자원이며, 그 희소성으로 인해 통제권을 놓고 다투는 다양한 세력 사이에 갈등이 생겼습니다. 섬들 사이의 하늘에서는 웅장한 비행선들이 마법의 기류를 타고 항해하며 무역과 탐험을 촉진합니다. 해적들과 하늘 습격자들이 구름 깊숙한 곳에 숨어 방심한 먹잇감을 항상 노리고 있습니다. 부유하는 땅 깊숙히 아래에는 언더보이드라는 어둡고 위험한 영역이 있으며, 악몽 같은 생물들과 말로 표현할 수 없는 부가 가득합니다. 오직 가장 용감한 모험가들만이 그 깊이를 탐험하려 하고, 더 적은 수만이 그 이야기를 들려주기 위해 돌아옵니다. 항상 존재하는 위협으로, 혼돈의 폭풍이 알려진 세계의 가장자리에서 맹위를 떨치며 그 길에 있는 모든 것을 삼켜버릴 위협을 가하고 있습니다. 에토리아의 영웅들에게 세계의 비밀을 밝히고 너무 늦기 전에 다가오는 어둠을 물리칠 방법을 찾는 임무가 주어졌습니다. 에토리아에서는 모든 섬이 이야기를 품고 있고, 모든 생물이 비밀을 가지고 있으며, 모든 모험이 이 경이롭고 위험에 처한 세계의 운명을 바꿀 수 있습니다.",
                "제 상상 속 세계에는 '오라키스'라는 도시가 있습니다. 이 도시는 순수한 빛의 기둥 위에 떠 있습니다. 도시의 벽은 수정 유리로 만들어져 있어 항상 새벽과 황혼의 색을 반사하며 영원한 천상의 광채를 발합니다. 건물들은 계절에 따라 숨을 쉬며 형태를 바꿉니다 - 봄에는 성장하고, 여름에는 강해지며, 가을에는 퇴색하기 시작해 겨울이 되면 안개가 됩니다.",
            ],
            world_description,
        )

    submit_button.click(
        generate_character, [world_description, persona_description], outputs=[character_json]
    ).then(fn=infer_flux, inputs=[character_json], outputs=[character_image])
    random_world_button.click(
        get_random_world_description, outputs=[world_description]
    )
    random_persona_button.click(
        get_random_persona_description, outputs=[persona_description]
    )

app.queue().launch(share=False)