# coding=utf-8
#
# Copyright 2024 Toshihiko Aoki
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from diffusers import AutoPipelineForText2Image, LCMScheduler
from llama_cpp import Llama
import gradio as gr

width = 512
height = 512
num_inference_steps = 4
guidance_scale = 1.0

from openvino_pipe import LatentConsistencyEngine
pipe = LatentConsistencyEngine(
    "sd-1.5-lcm-openvino"
)

llm = Llama(
    model_path="llm-jp-1.3b-v1.0_staircaptions-FT_Q2_K.gguf",
)

def ja2prompt(ja_prompt):
    response = llm(f"### Instruction:\n{ja_prompt}\n### Response:\n", max_tokens=128)
    return response['choices'][0]['text']


def prompt2img(sd_prompt):
    return pipe(
        sd_prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=1.0,
    ).images[0]


with gr.Blocks(title="tiny sd web-ui") as demo:
    gr.Markdown(f"## Japanese translation and hallucinations for Stable Diffusion")
    with gr.Row():
        with gr.Column(scale=3):
            ja = gr.Text(label="日本語")
            translate = gr.Button("変換")
            prompt = gr.Text(label="プロンプト")
        with gr.Column(scale=2):
            result = gr.Image()
            t2i = gr.Button("生成")
    translate.click(ja2prompt, ja, prompt)
    t2i.click(prompt2img, prompt, result)

demo.launch()