import logging

import gradio as gr
import keras_cv
import numpy as np
import tensorflow as tf
from huggingface_hub import from_pretrained_keras

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__file__)
prompt_token = "<token>"
text_encoder_url = "Dimitre/stablediffusion-canarinho_pistola"

logger.info(f'Inversed token used: "{prompt_token}"')
logger.info(f'Loading text encoder from: "{text_encoder_url}"')

stable_diffusion = keras_cv.models.StableDiffusion()
stable_diffusion.tokenizer.add_tokens(prompt_token)
text_encoder = from_pretrained_keras("Dimitre/stablediffusion-canarinho_pistola")
stable_diffusion._text_encoder = text_encoder
stable_diffusion._text_encoder.compile(jit_compile=True)


def generate_fn(input_prompt: str) -> np.ndarray:
    """Generates images from a text prompt

    Args:
        input_prompt (str): Text input prompt

    Returns:
        np.ndarray: Generated image
    """
    generated = stable_diffusion.text_to_image(
        prompt=input_prompt, batch_size=1, num_steps=50
    )
    return generated[0]


iface = gr.Interface(
    fn=generate_fn,
    title="Textual Inversion",
    description=f'Textual Inversion Demo, use "{prompt_token}" as the textual inversion token as shown in the examples',
    article="Note: Keras-cv uses lazy initialization, so the first use will be slower while the model is initialized.",
    inputs=gr.Textbox(
        label="Prompt",
        show_label=False,
        max_lines=2,
        placeholder="Enter your prompt",
        elem_id="input-prompt",
    ),
    outputs=gr.Image(),
    examples=[[f"A {prompt_token} portrait, 4k, highly detailed, highest quality, 8k"]],
)

if __name__ == "__main__":
    app, local_url, share_url = iface.launch()